【图像分类】实战——使用GoogLeNet识别动漫
【摘要】
目录
摘要
制作数据集
导入所需要的库
设置全局参数
图像预处理
读取数据
设置模型
设置训练和验证
测试
摘要
给你一张动漫图片，你能告诉我出自哪个动漫吗？今天我们就用GoogLeNet做到这一点。我选择的的动漫有秦岭神树、吞噬星空和秦时明月。图片样例如下：
秦岭神树
...
目录
摘要
给你一张动漫图片，你能告诉我出自哪个动漫吗？今天我们就用GoogLeNet做到这一点。我选择的的动漫有秦岭神树、吞噬星空和秦时明月。图片样例如下：
秦岭神树
秦时明月
吞噬星空
从上面的样例空可以看出，不同的动漫风格差异还是很大的，下面就说说我如何实现动漫分类的。
制作数据集
制作数据集，需要用到从动漫中抽取图片，抽取图片需要用到ffmpeg工具，具体的安装和使用参考：Windows安装ffmpeg，使用ffmpeg从视频中的抽取图像_AI浩-CSDN博客。
然后我们使用python，调用ffmeg实现对动漫视频的批量抽取的逻辑，在D盘中新建两个文件夹，一个文件夹命名qlss，这个文件夹存放动漫视频，一个文件夹是imags,用于存放抽取的图片。Python的代码如下：
-
import os
-
# 创建三个列表用来存储视频文件以及视频地址
-
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中执行ffmpeg命令实现对图片抽取。
新建项目
新建一个图像分类的项目，data里面放数据集，dataset文件夹中自定义数据的读取方法，这次我不采用默认的读取方式，太简单没啥意思。然后再新建train.py和test.py
在项目的根目录新建train.py，然后在里面写训练代码。
导入所需要的库
我这次选用的模型是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
设置全局参数
设置BatchSize、学习率和epochs，判断是否有cuda环境，如果没有设置为cpu。
-
# 设置全局参数
-
modellr = 1e-4
-
BATCH_SIZE = 32
-
EPOCHS = 10
-
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
图像预处理
在做图像与处理时，train数据集的transform和验证集的transform分开做，train的图像处理出了resize和归一化之外，还可以设置图像的增强，比如旋转、随机擦除等一系列的操作，验证集则不需要做图像增强，另外不要盲目的做增强，不合理的增强手段很可能会带来负作用，甚至出现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])
-
])
读取数据
将数据集解压后放到data文件夹下面，如图：
然后我们在dataset文件夹下面新建 __init__.py和dataset.py，在dataset.py文件夹写入下面的代码：
说一下代码的核心逻辑。
第一步 建立字典，定义类别对应的ID,用数字代替类别。
第二步 在__init__里面编写获取图片路径的方法。测试集只有一层路径直接读取，训练集在train文件夹下面是类别文件夹，先获取到类别，再获取到具体的图片路径。然后使用sklearn中切分数据集的方法，按照7：3的比例切分训练集和验证集。
第三步 在__getitem__方法中定义读取单个图片和类别的方法，由于图像中有位深度32位的，所以我在读取图像的时候做了转换。
-
# 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, '秦岭神树': 1, '秦时明月': 2}
-
class SeedlingData (data.Dataset):
-
-
def __init__(self, root, transforms=None, train=True, test=False):
-
"""
-
主要目标： 获取所有图片的地址，并根据训练，验证，测试划分数据
-
"""
-
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):
-
"""
-
一次返回一张图片的数据
-
"""
-
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)
然后我们在train.py调用SeedlingData读取数据 ，记着导入刚才写的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)
-
# 读取数据
-
print(dataset_train.imgs)
-
-
# 导入数据
-
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)
设置模型
使用CrossEntropyLoss作为loss，模型采用inception_v3，选用预训练模型。更改全连接层，将最后一层类别设置为3，然后将模型放到DEVICE。优化器选用Adam。
-
# 实例化模型并且移动到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)
-
# 选择简单暴力的Adam优化器，学习率调低
-
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
设置训练和验证
inception_v3模型较早，model输出有两个参数，在这里要注意。
output,hid = model(data)
-
# 定义训练过程
-
-
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))
-
-
-
# 验证过程
-
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))
-
-
-
# 训练
-
-
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')
完成上面的代码就可以训练了，这三部动漫的差别较大，所以很快就得到了非常好的结果。
测试
我介绍两种常用的测试方式，第一种是通用的，通过自己手动加载数据集然后做预测，具体操作如下：
测试集存放的目录如下图：
第一步 定义类别，这个类别的顺序和训练时的类别顺序对应，一定不要改变顺序！！！！
第二步 定义transforms，transforms和验证集的transforms一样即可，别做数据增强。
第三步 加载model，并将模型放在DEVICE里，
第四步 读取图片并预测图片的类别，在这里注意，读取图片用PIL库的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=('吞噬星空','秦岭神树','秦时明月')
-
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()]))
选择用了一集秦岭神树做测试集，运行结果：
第二种 使用自定义的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=('吞噬星空','秦岭神树','秦时明月')
-
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))
-
# 对应文件夹的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
-
公众号搜索“AI小浩”，关注后回复“GooLeNet实战”，获取源码、模型和数据集。
文章来源: wanghao.blog.csdn.net，作者：AI浩，版权归原作者所有，如需转载，请联系作者。
原文链接：wanghao.blog.csdn.net/article/details/117852220
【版权声明】本文为华为云社区用户转载文章，如果您发现本社区中有涉嫌抄袭的内容，欢迎发送邮件进行举报，并提供相关证据，一经查实，本社区将立刻删除涉嫌侵权内容，举报邮箱：
cloudbbs@huaweicloud.com
- 点赞
- 收藏
- 关注作者
评论（0）