使用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格式，這裡主要介紹驗證匯出模型是否正確

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，不過這個對於分類任務影響並不太大。