使用OpenCV的DNN模塊調用pytorch訓練的分類模型，這里記錄一下中間的流程，主要分為模型訓練，模型轉換和OpenCV調用三步。
一、訓練二分類模型
准備二分類數據，直接使用torchvision.models中的resnet18網絡，主要編寫的地方是自定義數據類中的__getitem__，和網絡最后一層。
__getitem__
將同類數據放在不同文件夾下，編寫Mydataset類，在__getitem__函數中增加數據增強。
class Mydataset(Dataset):
......
def __getitem__(self, idx):
# idx-[0->len(images)]
img, label = self.images[idx], self.labels[idx]
tf = transforms.Compose([
lambda x: Image.open(x).convert('RGB'),
transforms.Resize((int(self.resize), int(self.resize))),
# transforms.Resize((int(self.resize * 1.25), int(self.resize * 1.25))),
# transforms.RandomRotation(15),
# transforms.CenterCrop(self.resize),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
img = tf(img)
label = torch.tensor(label)
return img, label
......
- 修改網絡最后一層
依據類別，修改最后一層的輸出，主要代碼如下:
model = resnet18(pretrained=True) # 比較好的 model
model = nn.Sequential(*list(model.children())[:-1], # [b, 512, 1, 1] -> 接全連接層
# [b, 512, 1, 1] -> [b, 512]
torch.nn.Flatten(),
nn.Linear(512, 2)).to(device) # 添加全連接層
# x = torch.randn(2, 3, 224, 224)
# print(model(x).shape)
# 定義損失函數
criterion = nn.CrossEntropyLoss()
# 定義迭代參數的算法
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
二、Pytorch模型轉為ONNX模型
直接調用torch.onnx接口可將模型導出為ONNX格式，這里主要介紹驗證導出模型是否正確，參考鏈接：pytorch官方文檔
import torch
from torchvision import transforms
from PIL import Image
from torchvision.models import resnet18
import torch.nn as nn
import torch.onnx
import onnx
import onnxruntime
import numpy as np
torch_model = "./resnet18-2Class.pkl"
onnx_save_path = "./resnet18-2Class.onnx"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
data = torch.randn(1, 3, 224, 224, dtype=torch.float, device=device)
model = resnet18(pretrained=True)
model = nn.Sequential(*list(model.children())[:-1], # [b, 512, 1, 1] -> 接全連接層
nn.Flatten(), # [b, 512, 1, 1] -> [b, 512]
nn.Linear(512, 2)).to(device)
model.load_state_dict(torch.load(torch_model))
model.eval()
print("Start convert model to onnx...")
torch.onnx.export(model,
data,
onnx_save_path,
opset_version=10,
do_constant_folding=True, # 是否執行常量折疊優化
input_names=["input"], # 輸入名
output_names=["output"], # 輸出名
dynamic_axes={"input": {0: "batch_size"}, # 批處理變量
"output": {0: "batch_size"}}
)
print("convert onnx is Done!")
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
def get_test_transform():
tf = transforms.Compose([
lambda x: Image.open(x).convert('RGB'),
transforms.Resize((224, 224)),
# transforms.CenterCrop(self.resize),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
return tf
img_path = "./1.png"
img = get_test_transform()(img_path)
img = img.unsqueeze(0) # --> NCHW
print("input img mean {} and std {}".format(img.mean(), img.std()))
torch_out = model(img.to(device))
print("torch predict: ", torch_out)
# onnx
resnet_session = onnxruntime.InferenceSession(onnx_save_path)
inputs = {resnet_session.get_inputs()[0].name: to_numpy(img)}
onnx_out = resnet_session.run(None, inputs)[0]
print("onnx predict: ", onnx_out)
三、OpenCV調用ONNX模型進行分類
這里主要工作是對數據進行預處理，在第一部分中的__getitem__函數的增強部分，轉為openCV圖像處理如下，其他直接調用dnn模塊下的readNetFromONNX(modelPath)即可。
cv::Mat img = cv::imread(imgPath);
img.convertTo(img, CV_32FC3);
cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
cv::resize(img, img, cv::Size(224, 224));
img = img / 255.0;
std::vector<float> mean_value{ 0.485, 0.456, 0.406 };
std::vector<float> std_value{ 0.229, 0.224, 0.225 };
cv::Mat dst;
std::vector<cv::Mat> rgbChannels(3);
cv::split(img, rgbChannels);
for (auto i = 0; i < rgbChannels.size(); i++)
{
rgbChannels[i] = (rgbChannels[i] - mean_value[i]) / std_value[i];
}
cv::merge(rgbChannels, dst);
其中有一個注意點，就是同一張圖片用
torchvision.transforms中的Resize()和OpenCV的resize()函數處理的結果會有一點差別，這是因為transforms中默認使用的PIL的resize進行處理，除了默認的雙線性插值，還會進行antialiasing，不過這個對於分類任務影響並不太大。
其他
頭文件隱藏dnn.h
//*.h
class ClassifyByAI
{
public:
ClassifyByAI();
~ClassifyByAI();
bool LoadModel(const std::string onnxModelPath);
bool Predict(core::Image& image, int& classId, int imgSize = 64);
private:
void PreProcess(core::Image& image, int imgSize = 64);
struct CPrivate;
CPrivate* const mpD;
};
// *.cpp
struct ClassifyByAI::CPrivate
{
cv::dnn::Net net;
cv::Mat img;
};
ClassifyByAI::ClassifyByAI() : mpD(new CPrivate)
{
}
ClassifyByAI::~ClassifyByAI()
{
if (mpD) delete mpD;
}