【圖像分類】實(shí)戰(zhàn)——使用GoogLeNet識(shí)別動(dòng)漫

目錄

摘要

制作數(shù)據(jù)集

導(dǎo)入所需要的庫(kù)

設(shè)置全局參數(shù)

圖像預(yù)處理

讀取數(shù)據(jù)

設(shè)置模型

設(shè)置訓(xùn)練和驗(yàn)證

測(cè)試

摘要

給你一張動(dòng)漫圖片，你能告訴我出自哪個(gè)動(dòng)漫嗎？今天我們就用GoogLeNet做到這一點(diǎn)。我選擇的的動(dòng)漫有秦嶺神樹(shù)、吞噬星空和秦時(shí)明月。圖片樣例如下：

秦嶺神樹(shù)

秦時(shí)明月

吞噬星空

從上面的樣例空可以看出，不同的動(dòng)漫風(fēng)格差異還是很大的，下面就說(shuō)說(shuō)我如何實(shí)現(xiàn)動(dòng)漫分類的。

制作數(shù)據(jù)集

制作數(shù)據(jù)集，需要用到從動(dòng)漫中抽取圖片，抽取圖片需要用到ffmpeg工具，具體的安裝和使用參考：Windows安裝ffmpeg，使用ffmpeg從視頻中的抽取圖像_AI浩-CSDN博客。

然后我們使用python，調(diào)用ffmeg實(shí)現(xiàn)對(duì)動(dòng)漫視頻的批量抽取的邏輯，在D盤中新建兩個(gè)文件夾，一個(gè)文件夾命名qlss，這個(gè)文件夾存放動(dòng)漫視頻，一個(gè)文件夾是imags,用于存放抽取的圖片。Python的代碼如下：

import os

# 創(chuàng)建三個(gè)列表用來(lái)存儲(chǔ)視頻文件以及視頻地址

file_list = []

file_list_path=[]

filelist = []

# 源文件目錄

dir_path = 'D:\qlss'

#cmd命令存入str字符串

str = 'ffmpeg ' + '-i {} -ss 00:00:30 -f image2 -vf fps=1/5 -qscale:v 2 ../imags/img_{}%05d.jpg'

filenameList=os.listdir(dir_path)

print(filenameList)

for j in range(len(filenameList)):

filepath=os.path.join(dir_path,filenameList[j])

str_cmd = str.format(filepath, filenameList[j].split('.')[0])

print(str_cmd)

os.popen(str_cmd)

上面代碼的思路：遍歷視頻文件，使用Python在cmd中執(zhí)行ffmpeg命令實(shí)現(xiàn)對(duì)圖片抽取。

新建項(xiàng)目

新建一個(gè)圖像分類的項(xiàng)目，data里面放數(shù)據(jù)集，dataset文件夾中自定義數(shù)據(jù)的讀取方法，這次我不采用默認(rèn)的讀取方式，太簡(jiǎn)單沒(méi)啥意思。然后再新建train.py和test.py

在項(xiàng)目的根目錄新建train.py，然后在里面寫訓(xùn)練代碼。

導(dǎo)入所需要的庫(kù)

我這次選用的模型是inception_v3。

import torch.optim as optim

import torch

import torch.nn as nn

import torch.nn.parallel

import torch.utils.data

import torch.utils.data.distributed

import torchvision.transforms as transforms

from dataset.dataset import SeedlingData

from torch.autograd import Variable

from torchvision.models import inception_v3

設(shè)置全局參數(shù)

設(shè)置BatchSize、學(xué)習(xí)率和epochs，判斷是否有cuda環(huán)境，如果沒(méi)有設(shè)置為cpu。

# 設(shè)置全局參數(shù)

modellr = 1e-4

BATCH_SIZE = 32

EPOCHS = 10

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

圖像預(yù)處理

在做圖像與處理時(shí)，train數(shù)據(jù)集的transform和驗(yàn)證集的transform分開(kāi)做，train的圖像處理出了resize和歸一化之外，還可以設(shè)置圖像的增強(qiáng)，比如旋轉(zhuǎn)、隨機(jī)擦除等一系列的操作，驗(yàn)證集則不需要做圖像增強(qiáng)，另外不要盲目的做增強(qiáng)，不合理的增強(qiáng)手段很可能會(huì)帶來(lái)負(fù)作用，甚至出現(xiàn)Loss不收斂的情況。注：inception_v3模型輸入的size是3×299×299

transform = transforms.Compose([

transforms.Resize((299, 299)),

transforms.ToTensor(),

transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

])

transform_test = transforms.Compose([

transforms.Resize((299, 299)),

transforms.ToTensor(),

transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

])

讀取數(shù)據(jù)

將數(shù)據(jù)集解壓后放到data文件夾下面，如圖：

然后我們?cè)赿ataset文件夾下面新建?__init__.py和dataset.py，在dataset.py文件夾寫入下面的代碼：

說(shuō)一下代碼的核心邏輯。

第一步 建立字典，定義類別對(duì)應(yīng)的ID,用數(shù)字代替類別。

第二步 在__init__里面編寫獲取圖片路徑的方法。測(cè)試集只有一層路徑直接讀取，訓(xùn)練集在train文件夾下面是類別文件夾，先獲取到類別，再獲取到具體的圖片路徑。然后使用sklearn中切分?jǐn)?shù)據(jù)集的方法，按照7：3的比例切分訓(xùn)練集和驗(yàn)證集。

第三步 在__getitem__方法中定義讀取單個(gè)圖片和類別的方法，由于圖像中有位深度32位的，所以我在讀取圖像的時(shí)候做了轉(zhuǎn)換。

# coding:utf8

import os

from PIL import Image

from torch.utils import data

from torchvision import transforms as T

from sklearn.model_selection import train_test_split

Labels={'吞噬星空': 0, '秦嶺神樹(shù)': 1, '秦時(shí)明月': 2}

class SeedlingData (data.Dataset):

def __init__(self, root, transforms=None, train=True, test=False):

"""

主要目標(biāo)： 獲取所有圖片的地址，并根據(jù)訓(xùn)練，驗(yàn)證，測(cè)試劃分?jǐn)?shù)據(jù)

"""

self.test = test

self.transforms = transforms

if self.test:

imgs = [os.path.join(root, img) for img in os.listdir(root)]

self.imgs = imgs

else:

imgs_labels = [os.path.join(root, img) for img in os.listdir(root)]

imgs = []

for imglable in imgs_labels:

for imgname in os.listdir(imglable):

imgpath = os.path.join(imglable, imgname)

imgs.append(imgpath)

trainval_files, val_files = train_test_split(imgs, test_size=0.3, random_state=42)

if train:

self.imgs = trainval_files

else:

self.imgs = val_files

def __getitem__(self, index):

"""

一次返回一張圖片的數(shù)據(jù)

"""

img_path = self.imgs[index]

img_path=img_path.replace("\\",'/')

if self.test:

label = -1

else:

labelname = img_path.split('/')[-2]

label = Labels[labelname]

data = Image.open(img_path).convert('RGB')

data = self.transforms(data)

return data, label

def __len__(self):

return len(self.imgs)

然后我們?cè)趖rain.py調(diào)用SeedlingData讀取數(shù)據(jù) ，記著導(dǎo)入剛才寫的dataset.py(from dataset.dataset import SeedlingData)

dataset_train = SeedlingData('data/train', transforms=transform, train=True)

dataset_test = SeedlingData("data/train", transforms=transform_test, train=False)

# 讀取數(shù)據(jù)

print(dataset_train.imgs)

# 導(dǎo)入數(shù)據(jù)

train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)

設(shè)置模型

使用CrossEntropyLoss作為loss，模型采用inception_v3，選用預(yù)訓(xùn)練模型。更改全連接層，將最后一層類別設(shè)置為3，然后將模型放到DEVICE。優(yōu)化器選用Adam。

# 實(shí)例化模型并且移動(dòng)到GPU

criterion = nn.CrossEntropyLoss()

model_ft = inception_v3(pretrained=True)

num_ftrs = model_ft.fc.in_features

model_ft.fc = nn.Linear(num_ftrs, 3)

model_ft.to(DEVICE)

# 選擇簡(jiǎn)單暴力的Adam優(yōu)化器，學(xué)習(xí)率調(diào)低

optimizer = optim.Adam(model_ft.parameters(), lr=modellr)

def adjust_learning_rate(optimizer, epoch):

"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""

modellrnew = modellr * (0.1 ** (epoch // 50))

print("lr:", modellrnew)

for param_group in optimizer.param_groups:

param_group['lr'] = modellrnew

設(shè)置訓(xùn)練和驗(yàn)證

inception_v3模型較早，model輸出有兩個(gè)參數(shù)，在這里要注意。

output,hid = model(data)

# 定義訓(xùn)練過(guò)程

def train(model, device, train_loader, optimizer, epoch):

model.train()

sum_loss = 0

total_num = len(train_loader.dataset)

print(total_num, len(train_loader))

for batch_idx, (data, target) in enumerate(train_loader):

data, target = Variable(data).to(device), Variable(target).to(device)

output,hid = model(data)

loss = criterion(output, target)

optimizer.zero_grad()

loss.backward()

optimizer.step()

print_loss = loss.data.item()

sum_loss += print_loss

if (batch_idx + 1) % 10 == 0:

print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(

epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),

100. * (batch_idx + 1) / len(train_loader), loss.item()))

ave_loss = sum_loss / len(train_loader)

print('epoch:{},loss:{}'.format(epoch, ave_loss))

# 驗(yàn)證過(guò)程

def val(model, device, test_loader):

model.eval()

test_loss = 0

correct = 0

total_num = len(test_loader.dataset)

print(total_num, len(test_loader))

with torch.no_grad():

for data, target in test_loader:

data, target = Variable(data).to(device), Variable(target).to(device)

output = model(data)

loss = criterion(output, target)

_, pred = torch.max(output.data, 1)

correct += torch.sum(pred == target)

print_loss = loss.data.item()

test_loss += print_loss

correct = correct.data.item()

acc = correct / total_num

avgloss = test_loss / len(test_loader)

print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(

avgloss, correct, len(test_loader.dataset), 100 * acc))

# 訓(xùn)練

for epoch in range(1, EPOCHS + 1):

adjust_learning_rate(optimizer, epoch)

train(model_ft, DEVICE, train_loader, optimizer, epoch)

val(model_ft, DEVICE, test_loader)

torch.save(model_ft, 'model.pth')

完成上面的代碼就可以訓(xùn)練了，這三部動(dòng)漫的差別較大，所以很快就得到了非常好的結(jié)果。

測(cè)試

我介紹兩種常用的測(cè)試方式，第一種是通用的，通過(guò)自己手動(dòng)加載數(shù)據(jù)集然后做預(yù)測(cè)，具體操作如下：

測(cè)試集存放的目錄如下圖：

第一步 定義類別，這個(gè)類別的順序和訓(xùn)練時(shí)的類別順序?qū)?yīng)，一定不要改變順序！！！！

第二步 定義transforms，transforms和驗(yàn)證集的transforms一樣即可，別做數(shù)據(jù)增強(qiáng)。

第三步 加載model，并將模型放在DEVICE里，

第四步 讀取圖片并預(yù)測(cè)圖片的類別，在這里注意，讀取圖片用PIL庫(kù)的Image。不要用cv2，transforms不支持。

import torch.utils.data.distributed

import torchvision.transforms as transforms

from PIL import Image

from torch.autograd import Variable

import os

# classes = ('Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed',

# 'Common wheat','Fat Hen', 'Loose Silky-bent',

# 'Maize','Scentless Mayweed','Shepherds Purse','Small-flowered Cranesbill','Sugar beet')

classes=('吞噬星空','秦嶺神樹(shù)','秦時(shí)明月')

transform_test = transforms.Compose([

transforms.Resize((299, 299)),

transforms.ToTensor(),

transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

])

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = torch.load("model.pth")

model.eval()

model.to(DEVICE)

path='data/test/'

testList=os.listdir(path)

for file in testList:

img=Image.open(path+file)

img=transform_test(img)

img.unsqueeze_(0)

img = Variable(img).to(DEVICE)

out=model(img)

# Predict

_, pred = torch.max(out.data, 1)

print('Image Name:{},predict:{}'.format(file,classes[pred.data.item()]))

選擇用了一集秦嶺神樹(shù)做測(cè)試集，運(yùn)行結(jié)果：

第二種 使用自定義的Dataset讀取圖片

import torch.utils.data.distributed

import torchvision.transforms as transforms

from dataset.dataset import SeedlingData

from torch.autograd import Variable

# classes = ('Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed',

# 'Common wheat','Fat Hen', 'Loose Silky-bent',

# 'Maize','Scentless Mayweed','Shepherds Purse','Small-flowered Cranesbill','Sugar beet')

classes=('吞噬星空','秦嶺神樹(shù)','秦時(shí)明月')

transform_test = transforms.Compose([

transforms.Resize((299, 299)),

transforms.ToTensor(),

transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

])

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = torch.load("model.pth")

model.eval()

model.to(DEVICE)

dataset_test =SeedlingData('data/test/', transform_test,test=True)

print(len(dataset_test))

# 對(duì)應(yīng)文件夾的label

for index in range(len(dataset_test)):

item = dataset_test[index]

img, label = item

img.unsqueeze_(0)

data = Variable(img).to(DEVICE)

output = model(data)

_, pred = torch.max(output.data, 1)

print('Image Name:{},predict:{}'.format(dataset_test.imgs[index], classes[pred.data.item()]))

index += 1

公眾號(hào)搜索“AI小浩”，關(guān)注后回復(fù)“GooLeNet實(shí)戰(zhàn)”，獲取源碼、模型和數(shù)據(jù)集。

機(jī)器學(xué)習(xí)

版權(quán)聲明：本文內(nèi)容由網(wǎng)絡(luò)用戶投稿，版權(quán)歸原作者所有，本站不擁有其著作權(quán)，亦不承擔(dān)相應(yīng)法律責(zé)任。如果您發(fā)現(xiàn)本站中有涉嫌抄襲或描述失實(shí)的內(nèi)容，請(qǐng)聯(lián)系我們jiasou666@gmail.com 處理，核實(shí)后本網(wǎng)站將在24小時(shí)內(nèi)刪除侵權(quán)內(nèi)容。