【图像分类】实战——使用GoogLeNet识别动漫-伙伴云

【图像分类】实战——使用GoogLeNet识别动漫

网友投稿 684 2022-05-29

摘要

制作数据集

导入所需要的库

设置全局参数

图像预处理

读取数据

设置模型

设置训练和验证

测试

摘要

给你一张动漫图片，你能告诉我出自哪个动漫吗？今天我们就用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(),

【图像分类】实战——使用GoogLeNet识别动漫

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实战”，获取源码、模型和数据集。

机器学习

筛选后显示数据重新下拉序号，展开后序号不变（筛选后序号下拉不递增怎么办）

684 2022-05-29

【图像 分类】实战——使用GoogLeNet识别动漫

分类汇总，显示文字怎样操作（怎么用分类汇总）

分类汇总不出现汇总字眼了，该怎么处理》（分类汇总不了是怎么回事）

筛选后显示数据重新下拉序号，展开后序号不变（筛选后序号下拉不递增怎么办）

推荐文章

企业生产管理是什么，企业生产管理软件

进盘点进销存软件排行榜前十名

进销存系统哪个简单好用？进销存系统优点

工厂生产管理（工厂生产管理流程及制度）

生产管理软件，机械制造业生产管理，制造业生产过程管理软件

进销存软件和ERP有什么区别？进销存与erp软件理解

进销存如何进行库存管理

如何利用excel制作销售订单管理系统？

数据库订单管理系统有哪些功能？数据库订单管理系统怎么设计？

什么是数据库管理系统？

最近发表

热评文章

零代码开发是什么？2022低代码平台排行榜

智能进销存库存管理系统（智慧进销存）

在线文档哪家强？8款在线文档编辑软件推荐

WPS2016怎么绘制简单的价格表?

连锁餐饮管理系统的功能有哪些？餐饮服务系统的构成及工

Excel项目进度表模板，简化您的项目进度管理

友情链接