【YOLOv8】自制数据集PyTorch训练并用ONNXRuntime在C++上部署

本篇将详细介绍如何用自制的（Custom）数据集来训练YOLOv8模型，并用模型部署到C++的工程中。其中，会涉及本人开发过程中遇到的bug及解决方法。

一. PyTorch训练

1.1 数据集制作

数据集的制作跟YOLOv5是一样的，使用Labelme进行标注，会生成 .json文件。

YOLO接受的标签数据是包含id和归一化后坐标的 .txt文件，所以自己写一个 .json转 .txt的python 脚本。以下程序首先将没有标签的数据删除，将剩余的图片和标签保存到新的一个文件夹内(/selec_data)；然后将.json标签转为.txt标签；最后随机划分为训练集和验证集保存在两个子文件夹内（/train_data/train和/train_data/val）

import os 
import glob
import cv2
import json
import numpy as np
import random
import shutil

## Labelme保存的图片和标签的文件夹：
in_dir = './custom_data/data/'

## 按匹配对整理，去除没有标签的图片
img_pths = glob.glob(in_dir + "/*.bmp")
for img_pth in img_pths:
    json_pth = img_pth.replace(".bmp", ".json")
    print(json_pth)
    if os.path.exists(json_pth):
        shutil.copy(img_pth, img_pth.replace('/data', '/select_data'))
        shutil.copy(json_pth, json_pth.replace('/data', '/select_data'))


## 将labelme_json标注转_txt
def convert(size, box):
    """
    convert [xmin, xmax, ymin, ymax] to [x_centre, y_centre, w, h]
    """
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

class_names = ["pupil", 'iris'] #数组元素替换为自己数据的类别

out_dir = "./custom_data/select_data/"
txt_dir = "./custom_data/select_data/"
os.makedirs(txt_dir, exist_ok=True)

json_pths = glob.glob(out_dir + "/*.json")

for json_pth in json_pths:
    f1 = open(json_pth, "r")
    json_data = json.load(f1)

    img_pth = os.path.join(json_pth.replace("json", "bmp"))
    img = cv2.imread(img_pth)
    h, w = img.shape[:2]

    tag = os.path.basename(json_pth)
    out_file = open(os.path.join(txt_dir, tag.replace("json", "txt")), "w")
    # print(json_data)
    label_infos = json_data["shapes"]
    for label_info in label_infos:
        label = label_info["label"]
        points = label_info["points"]
        print("+++", len(points))
        if len(points) >= 3:
            points = np.array(points)
            print(points.shape)
            xmin, xmax = max(0, min(np.unique(points[:, 0]))), min(w, max(np.unique(points[:, 0])))
            ymin, ymax = max(0, min(np.unique(points[:, 1]))), min(h, max(np.unique(points[:, 1])))
            print("++++", ymin, ymax)
        elif len(points) == 2:
            x1, y1 = points[0]
            x2, y2 = points[1]
            xmin, xmax = min(x1, x2), max(x1, x2)
            ymin, ymax = min(y1, y2), max(y1, y2)
        else:
            continue
        bbox = [xmin, xmax, ymin, ymax]
        bbox_ = convert((w,h), bbox)
        cls_id = class_names.index(label)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bbox_]) + '\n')


def mymovefile(srcfile, dstpath):                       # 移动函数
    if not os.path.isfile(srcfile):
        print ("%s not exist!"%(srcfile))
    else:
        fpath,fname=os.path.split(srcfile)             # 分离文件名和路径
        if not os.path.exists(dstpath):
            os.makedirs(dstpath)                       # 创建路径
        shutil.move(srcfile, dstpath + fname)          # 移动文件
        print ("move %s -> %s"%(srcfile, dstpath + fname))

# 划分训练集和验证集
out_dir10 = "./custom_data/train_data/train/images"
out_dir11= "./custom_data/train_data/train/labels"
out_dir20 = "./custom_data/train_data/val/images"
out_dir21 = "./custom_data/train_data/val/labels"

os.makedirs(out_dir10, exist_ok=True)
os.makedirs(out_dir11, exist_ok=True)
os.makedirs(out_dir20, exist_ok=True)
os.makedirs(out_dir21, exist_ok=True)

txt_dir = "./custom_data/select_data/"
img_dir = "./custom_data/select_data/"
txt_names = glob.glob(txt_dir+'*.txt')

l = len(txt_names)
sel_names = random.sample(txt_names, int(l*0.8)) #0.90
# print(sel_names)
for txt_name in txt_names:
    img_name = txt_name.replace('txt', 'bmp')
    if txt_name in sel_names:  
        print("---", txt_name)
        shutil.copy(txt_name, txt_name.replace('/select_data', '/train_data/train/labels'))          # 复制文件
        shutil.copy(img_name, img_name.replace('/select_data', '/train_data/train/images'))
    else:
        shutil.copy(txt_name, txt_name.replace('/select_data', '/train_data/val/labels'))          # 复制文件
        shutil.copy(img_name, img_name.replace('/select_data', '/train_data/val/images'))

1.2 训练

训练环境

Python >=3.8
PyTorch >= 1.8
再安装其他必须的库即可

在ultralytics官方仓库下载YOLOv8源码，下载后打开整个工程，在根目录下新建文件夹，命名为datasets，把刚刚准备好的数据集复制到./datasets文件夹下。在./ultraltics/cfg/datasets/下新建一个custom_data.yaml文件，path处填写./datasets文件夹下的数据路径，分别有train和val两个子文件夹；类别数量和类别名称按自己的填写，这里的名称要按照你标注时的id顺序填写，否则后面验证和预测的时候会出现混乱。

# my custom dataset
path: custom_data/train_data
train: # train images (relative to 'path') 
  - train
val: # val images (relative to 'path')
  - val

# number of classes
nc: 2

# class names
names: ["pupil", "iris"]

数据准备已经结束了，接下来准备训练文件。在文件夹根目录下新建python文件./train.py：

from ultralytics import YOLO

if __name__ =='__main__':
	# 模型大小从小到大有n, s, m, l, x，需要哪一个就改v8后面的字母
    model = YOLO('ultralytics/cfg/models/v8/yolov8n.yaml')

	# 这里的data必须设置为我们刚才创建的.yaml文件
    model.train(data = 'ultralytics/cfg/datasets/custom_data.yaml', 
                epochs = 200, 
                batch = 32, 
                imgsz = 640,
                device = 0)   # 如果是CPU训练就写'cpu'，否则表示GPU的编号
                
    metrics = model.val()
    
    path = model.export(format="onnx")

因为我们下载了ultralytics的源码，所以不用安装ultralytics库就可以直接from ultralytics import YOLO。

运行train.py即可开始训练，训练完成后，最后一行代码已经导出了后续我们部署需要的.ONNX文件。模型文件都保存在./runs/detect/train/weights/，这里还能看到验证结果的图片。注意一下：文件夹./runs/detect/train名称会每运行一次就累加一次，train2、train3…

1.3 测试

训练完成后先在Python的onnxruntime上测试一下导出的onnx文件能否正常工作。
先安装onnxruntime库：
GPU：pip install onnxruntime-gpu
CPU：pip install onnxruntime

测试代码在文件夹./examples/YOLOv8_ONNXRuntime/main.py，把待测试图片复制到文件夹./ultralytics/assets/下，修改下面的输入参数：

parser = argparse.ArgumentParser()
# 这里的模型文件路径注意修改为自己满意的那个train文件夹
parser.add_argument("--model", type=str, default="runs/detect/train/weights/best.onnx", help="Input your ONNX model.")
parser.add_argument("--img", type=str, default=str(ASSETS / "img.bmp"), help="Path to input image.")
parser.add_argument("--conf-thres", type=float, default=0.5, help="Confidence threshold")
parser.add_argument("--iou-thres", type=float, default=0.5, help="NMS IoU threshold")
    args = parser.parse_args()

运行文件main.py，如果没有问题的话会生成检测结果。

检测结果上面会有物体名称，显示car之类的，因为我们没有修改main.py中类别名称，它默认是COCO数据集的，这个影响不大。

二. C++部署

2.1 环境

因为看到其他博客说YOLOv8要用opencv库推理的话要求opencv >= 4.7.1，考虑到项目其他部分对opencv版本的要求，所以我不打算升级opencv，使用了ONNXRuntime推理，另外我全程是使用GPU，环境如下：

visual studio 2022
CUDA 11.2   // 官方说需要CUDA >=11.4，我这里验证11.2没问题，应该CUDA11.x都可以
cudnn 8
onnxRuntime-gpu 1.14.1   // 版本略高的ONNX是可以兼容略低版本的CUDA 
opencv >= 4.0.0

2.2 部署

在ultralytics官方仓库下载的YOLOv8源码中，在文件夹./example/YOLOv8-ONNXRuntime-CPP中，核心代码有inference.h和inference.cpp。main.cpp中是进行图片的读取和绘制检测框。

首先因为我用的是GPU部署推理，所以要在visual studio中把GPU的环境都配置好。

在包含目录中添加onnxruntime和CUDA的包含路径（include），当然还有opencv：

这里的路径要写你自己PC上的具体路径。

在库目录中添加库路径：

在链接器->输入->附加依赖项 添加所需的库文件

CUDA的库文件其实没有用到，不须要添加

最后，把onnxruntime-gpu文件夹中的动态库onnxruntime.dll、onnxruntime_providers_cuda.dll、onnxruntime_providers_shared.dll文件复制到项目生成的/x64文件夹中，否则可能会运行报错缺少库文件。

因为inference.h和inference.cpp都有源码，这里不过多解释，直接用就可以了。

这里要强调几个地方和我遇到的bug
1.main.cpp文件中用到了c++ 17读取文件库#include <filesystem>，如果你的c++版本低，就自己写文件读取函数就可以了，网上找找代码。

2.在main.cpp中找到params.modelPath，修改为自己训练好并导出的.onnx文件，按需修改阈值。

3.main.cpp中函数DetectTest()调用了ReadCocoYaml(yoloDetector);，这里是为了读取COCO数据集中类别名称，我们训练的自己的数据集，把这一行改为

// 放你自己数据集的类别名称。
yoloDetector->classes = {"pupil", "iris"};

4.想用GPU推理，要在inference.h中定义宏 USE_CUDA才能启用GPU推理的部分：

#define    USE_CUDA

5.在TensorProcess()函数中，官方源码是这样的:

    case YOLO_DETECT_V8:
    case YOLO_DETECT_V8_HALF:
    {
        int strideNum = outputNodeDims[1];//8400
        int signalResultNum = outputNodeDims[2];//84
        std::vector<int> class_ids;
        std::vector<float> confidences;
        std::vector<cv::Rect> boxes;
        cv::Mat rawData;
        if (modelType == YOLO_DETECT_V8)
        {
            // FP32
            rawData = cv::Mat(strideNum, signalResultNum, CV_32F, output);
        }
        else
        {
            // FP16
            rawData = cv::Mat(strideNum, signalResultNum, CV_16F, output);
            rawData.convertTo(rawData, CV_32F);
        }
        //Note:
        //ultralytics add transpose operator to the output of yolov8 model.which make yolov8/v5/v7 has same shape
        //https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt
        //rowData = rowData.t();

        float* data = (float*)rawData.data;

按照官方代码运行得到的是错误结果，跟踪代码变量的时候，发现下面这两行代码的地方

int strideNum = outputNodeDims[1];//8400
int signalResultNum = outputNodeDims[2];//84

strideNum=6，signalResultNum=8400，与注释里并不相同，搜索资料后得知，84是coco数据集类别数80+4，8400是固定的。我有两个类别，所以是6和8400。但是，两个变量是反的！！ 也就说strideNum应该是8400，signalResultNum应该是6。

这一点去跟踪python版本的onnxruntime推理也可以看出，而且代码下面注释的那段Note:提醒了我可能有问题，到官方仓库的问答中也找到了相关回答。

因此改为：

int strideNum = outputNodeDims[2];//8400
int signalResultNum = outputNodeDims[1];//84

这下就正确检测到目标啦。

如果有任何问题可以在评论区提问