前言
这几天在研究如何使用opencv的dnn模块进行快速搭建以及缩减代码量进行完成图像分类任务，这篇文章中将记录我在做这个任务的过程中遇到的问题以及如何实现的逻辑，如有纰漏，还请大家批评指正！感谢！
构思逻辑步骤
在这里我将主要介绍逻辑构建以及后面出现的模块化设计。其基本逻辑如下：
核心逻辑板块
- 通过pytorch训练一个自己数据集的Model(该Model的格式为PTH)
- PTH2ONNX模块进行把权重转化为ONNX格式
- 使用opencv的 dnn 模块进行调用onnx完成 图像识别
训练自己的Model
这里我们准备好自己的数据集，为了方便大家操作进行，我使用的是手写数字数据集作为我自己的数据集进行训练（该数据集的介绍大家可以自行google），使用的训练方法为目录式读取数据。
构建Model网络，这里的in_c计算方式为：
in_c=1200 = 训练图像的维度 x 图像的宽 x 图像的高
out_c=10 = 类别数
import torchvision.datasets
import time
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch import optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
class Net(nn.Module):
def __init__(self, in_c=1200, out_c=10):
super(Net, self).__init__()
# 定义全连接层
self.fc1 = nn.Linear(in_c, 512)
# 定义激活层
self.act1 = nn.ReLU(inplace=True)
self.fc2 = nn.Linear(512, 256)
self.act2 = nn.ReLU(inplace=True)
self.fc3 = nn.Linear(256, 128)
self.act3 = nn.ReLU(inplace=True)
self.fc4 = nn.Linear(128, out_c)
def forward(self, x):
# x = x.view(-1, 1200)
x = self.act1(self.fc1(x))
x = self.act2(self.fc2(x))
x = self.act3(self.fc3(x))
x = self.fc4(x)
return x
有了上述的网络模型后，需要采用目录式调用数据进行分类，分好训练集与测试记数据，设置好
损失函数与优化器，计算Loss以及acc值，大家在这个里面需要注意的事项有如下几点;
情况一：
若save代码为
torch. save (network.cpu().state_ dict(), model name)
则load的代码应为
network.load_ state_ dict(torch. load(model name))
情况二：
若save代码为
torch. save (network, model_ name )
则load的代码应为
network.1oad_ state_ dict(torch.1oad(model_ name) .cpu().state_ _dict())
在本案例中我将采用情况二 进行存储modle,在后面进行调用的时候也可以避免不必要的报错。到此，我们能够顺利得到训练好的model，大家可以挑选最优解进行后面的测试预测。
if __name__ == '__main__':
# 输入训练和测试集的路径
train_root = 'C:/Users/kiven/Desktop/小麦/demo/datas/train/'
test_root = 'C:/Users/kiven/Desktop/小麦/demo/datas/test/'
# 将文件夹的内容载入dataset
train_dataset = torchvision.datasets.ImageFolder(root=train_root, transform=torchvision.transforms.ToTensor())
test_dataset = torchvision.datasets.ImageFolder(root=test_root, transform=torchvision.transforms.ToTensor())
# DataLoader 读取数据
train_data = DataLoader(dataset=train_dataset, # 输入自己要加载的数据set
batch_size=5, # 一个批量的大小
shuffle=True, # 是否打乱顺序
num_workers=4, # 是否使用多进程，0代表不使用
pin_memory=True, # 是否将数据保存在pin_memory区， pin_memory数据转移到Gpu中会快一些
drop_last=True) # 当为Ture时，dataset中的数据个数不是batch_size整数倍时，将多余出不足一个batch的数据丢弃
test_data = DataLoader(dataset=test_dataset, # 输入自己要加载的数据set
batch_size=5, # 一个批量的大小
shuffle=True, # 是否打乱顺序
num_workers=4, # 是否使用多进程，0代表不使用
pin_memory=True, # 是否将数据保存在pin_memory区， pin_memory数据转移到Gpu中会快一些
drop_last=True) # 当为Ture时，dataset中的数据个数不是batch_size整数倍时，将多余出不足一个batch的数据丢弃
t1 = time.time()
# 搭建网络
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = Net()
cudnn.benchmark = True
net = net.to(device)
# 定义损失函数 -- 交叉熵
criterion = torch.nn.CrossEntropyLoss().to(device)
# 定义优化器 -- 随机梯度下降
optimizer = optim.SGD(net.parameters(), lr=0.01, weight_decay=0.00005)
# 开始训练
losses = [] # 记录训练损失
acces = [] # 记录训练精度
eval_losses = [] # 记录测试损失
eval_acces = [] # 记录测试精度
nums_epoch = 30 # 训练次数
for epoch in range(nums_epoch):
train_loss = 0 # 设置训练损失的初始值
train_acc = 0 # 设置训练精度的初始值
net.train()
for batch, (img, label) in enumerate(train_data):
img = img.reshape(img.size(0), -1)
img = Variable(img)
img = img.to(device)
label = Variable(label)
label = label.to(device)
# 向前传播
out = net(img)
loss = criterion(out, label.long())
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 记录误差
train_loss += loss.item()
# 计算分类正确率
_, pred = out.max(1)
num_correct = (pred == label.long()).sum().item()
acc = num_correct / img.shape[0]
train_acc += acc
losses.append(train_acc / len(train_data))
acces.append(train_acc / len(train_data))
eval_loss = 0
eval_acc = 0
# 测试集不训练
for img, label in test_data:
img = img.reshape(img.size(0), -1)
img = Variable(img)
img = img.to(device)
label = Variable(label)
label = label.to(device)
out = net(img)
loss = criterion(out, label.long())
eval_loss += loss.item()
_, pred = out.max(1)
num_correct = (pred == label.long()).sum().item()
acc = num_correct / img.shape[0]
eval_acc += acc
eval_losses.append(eval_loss / len(test_data))
eval_acces.append(eval_acc / len(test_data))
# 打印参数
set_epoch = epoch + 1
set_lossTrain = train_loss / len(train_data)
set_AccTrain = train_acc / len(train_data)
set_lossEval = eval_loss / len(test_data)
set_AccEval = eval_acc / len(test_data)
print('[INFO] Epoch-{}: Train: Loss-{:.4f},Accuracy-{:.4f} |Test:Loss-{:.4f}, Accuracy-{:.4f}'.format(set_epoch,
set_lossTrain,
set_AccTrain,
set_lossEval,
set_AccEval))
torch.save(net, './model/Epoch-%s-TrainLoss-%s-TestLoss-%s.pth' % (set_epoch,
set_lossTrain,
set_lossEval))
PTH2ONNX模块转换
在上面的训练网络的部分我们得到了当前最优的PTH，我们需要进行转换，在进行转换的时候需要注意我们在调用网络结构需要将这个地方注释掉的地方取消注释，这个1200的计算方法同in_c.后面代码部分的input为图像的（通道数， 图像的宽，图像的高）
import torch
import torch.onnx
from mymodel import Net
def pth_to_onnx(input, checkpoint, onnx_path, input_names=['input1'], output_names=['output'], device='cpu'):
if not onnx_path.endswith('.onnx'):
print('Warning!')
return 0
model = Net()
model.load_state_dict(torch.load(checkpoint).cpu().state_dict()) # 初始化权重
model.eval()
# #指定模型的输入，以及onnx的输出路径
torch.onnx.export(model, input, onnx_path, verbose=True, input_names=input_names,
output_names=output_names) # 指定模型的输入，以及onnx的输出路径
if __name__ == '__main__':
checkpoint = './lenet.pth'
onnx_path = './ModelLenet.onnx'
input = torch.randn(3, 20, 20)
pth_to_onnx(input, checkpoint, onnx_path)
测试模块
在进行测试的过程中，我们有如下模块需要搭建：
- 网咯的读取
- 标签数据的读取
- 图像数据的处理
- 结果的输出
在上述的模块中采用cv2.dnn.readNet读取网络，可以避免过多的代码部分且这样我们 具备高效性，快捷性和通用性。这个里面需要注意的地方是：这里的input1和output需要同转换模块（pth2onnx中的input_names和output_names保持一致）
这里有另一个注意的地方是，代码中调用的class.names需要同训练的时候，调用的label一致
import cv2
def PredictImg(Img):
net = cv2.dnn.readNet("ModelLenet.onnx")
with open('class.names', 'rt') as f:
classes = f.read().rstrip('\n').split('\n')
num_classes = len(classes)
blol = cv2.dnn.blobFromImage(Img, scalefactor=1, size=(20, 20))
net.setInput(blol, 'input1')
prob = net.forward('output')
probMat = prob.reshape(1, num_classes)
# 求出匹配结果的最小值，最大值，并得到最大值，最小值的索引
_, maxVal, _, maxLoc = cv2.minMaxLoc(probMat)
ClassName = classes[maxLoc[0]]
print(ClassName)
SrcImg = cv2.imread("1.jpg")
PredictImg(SrcImg)
展望
这里我是采用python进行识别的，为了提速的化，我们可以采用C++进行改写，后续我将为大家带来C++版本的代码。
补充C++测试代码
#include <opencv2/opencv.hpp>
#include <iostream>
#include <fstream>
using namespace cv;
using namespace cv::dnn;
using namespace std;
int main()
{
Mat img = imread("/home/kiven-yang/TestOpencv/include/1.jpg");
if (img.empty())
{
printf("could not load image...\n");
return -1;
}
//读取分类种类名称
String typeListFile = "/home/kiven-yang/TestOpencv/include/2.txt";
vector<String> typeList;
ifstream file(typeListFile);
if (!file.is_open())
{
printf("请确认分类种类名称是否正确");
return -1;
}
std::string type;
while (!file.eof())
{
//读取名称
getline(file, type);
if (type.length())
typeList.push_back(type);
}
file.close();
// 加载网络
String tf_pb_file = "/home/kiven-yang/TestOpencv/include/ModelLenet.onnx";
Net net = readNet(tf_pb_file);
if (net.empty())
{
printf("请确认模型文件是否为空文件");
return -1;
}
//对输入图像数据进行处理
Mat blob = blobFromImage(img, 1.0f, Size(20, 20), Scalar(), true, false);
//进行图像种类预测
Mat prob;
net.setInput(blob, "input1");
prob = net.forward("output");
// 得到最可能分类输出
Mat probMat = prob.reshape(1, 1);
Point classNumber;
double classProb; //最大可能性
minMaxLoc(probMat, NULL, &classProb, NULL, &classNumber);
string typeName = typeList.at(classNumber.x).c_str();
//检测内容
int down_width = 400;
int down_height = 400;
Mat resized_down;
//resize down
resize(img, resized_down, Size(down_width, down_height), INTER_LINEAR);
string str = typeName + " possibility:" + to_string(classProb);
putText(resized_down, str, Point(15, 15), FONT_HERSHEY_SIMPLEX, 1.0, Scalar(0, 0, 255), 2, 2);
imshow("图像判断结果", resized_down);
waitKey(0);
return 0;
}
