yolov8目标检测,训练、导出onnx、opencv dnn调用

1、train

ultralytics目录里找到yolov8.yaml,拷贝到项目路径里。把其中的nc改为自己的类别数。

新建个yoloData.yaml,用于指定样本目录和具体类别,内容如下

train: /home/aaa/datasets/detect/dst/images/train
val: /home/aaa/datasets/detect/dst/images/val
test: /home/aaa/datasets/detect/dst/images/test

nc: 2
names:
  0: cat  #注意0前有两个空格
  1: dog

 train.py

from ultralytics import YOLO

def train():
    model=YOLO("./yolov8.yaml").load("./yolov8n.pt")    #n\s\m 模型从小到大,越大越准确但更耗时
    results=model.train(
        data="yoloData.yaml",       #样本
        epochs=300,                 #训练300轮
        imgsz=640,
        batch=16,
        workers=8,
        amp=False   #禁用amp检查,否则会下载最新版本的yolo如11
    )

if __name__=='__main__':
    train()

2、导出onnx

我预测固定尺寸图像,宽640、高160

from ultralytics import YOLO
if __name__=='__main__':
    model = YOLO('./runs/detect/train2/weights/best.pt', task='detect')
    model.export(format='onnx',imgsz=[160,640],batch=1,device='0',opset=12,dynamic=False)

细节请查阅两个博客:

YOLO11在训练和导出时的一些参数设置 - 夕西行 - 博客园

yolo导出,动态batch,固定图片尺寸 - 夕西行 - 博客园

3、opencv dnn调用

注意opset不能太高,opencv的版本不能太旧,建议4.8及以上

// YoloDetect.cpp: 定义应用程序的入口点。
//

#include "YoloDetect.h"
#include "opencv2/opencv.hpp"
#include "opencv2/dnn.hpp"
#include <vector>

using namespace std;
using namespace cv;
using namespace dnn;

// 定义12个类别及其显示颜色
const vector<string> CLASS_NAMES = {
    "cat", "dog"
};
const vector<Scalar> COLORS = {
    Scalar(255, 0, 0), Scalar(0, 255, 0)
};

// YOLOv8后处理函数
void postProcess(Mat& frame, const vector<Mat>& outputs, float confThreshold = 0.5, float nmsThreshold = 0.4) {
    Mat output = outputs[0];                    //shape: [1, 84, 8400]
    const int num_boxes = output.size[2];       // 8400,框的数量,不同像素的图像数量会不同
    const int num_features = output.size[1];    // 84 (xywh + 80 classes)

    //output.data数据顺序:8400个x、8400个y、8400个w、8400个h、第1类的置信度、第2类的置信度...第80类的置信度。
    Mat output_reshape = Mat(Size(num_boxes, num_features), CV_32F, (float*)output.data).t();//[1,84,8400]=>[1,8400,84],即8400行,每行xywh+80类的置信度

    // 获取框、分数(即置信度)和类别
    vector<Rect> vBoxes;
    vector<float> vScores;
    vector<int> vClass_ids;
    for (int r = 0; r < num_boxes; ++r)
    {
        Mat row = output_reshape.row(r);
        Mat scores = row.colRange(4, num_features);
        Point classIdPoint;
        double maxClassScore;
        minMaxLoc(scores, 0, &maxClassScore, 0, &classIdPoint);
        //存储满足置信度的结果
        if (maxClassScore > confThreshold) {
            float x = row.at<float>(0);             // 中心x
            float y = row.at<float>(1);             // 中心y
            float w = row.at<float>(2);             // 宽度
            float h = row.at<float>(3);             // 高度            
            int left = static_cast<int>((x - w/2)); // 转换为左上角坐标
            int top = static_cast<int>((y - h/2));
            int width = static_cast<int>(w);
            int height = static_cast<int>(h);
            vBoxes.emplace_back(left, top, width, height);//收集结果
            vScores.push_back(static_cast<float>(maxClassScore));
            vClass_ids.push_back(classIdPoint.x);
        }
    }

    // 非极大值抑制
    vector<int> indices;
    NMSBoxes(vBoxes, vScores, confThreshold, nmsThreshold, indices);

    // 绘制检测结果
    for (int idx : indices) 
    {
        Rect box = vBoxes[idx];
        int classId = vClass_ids[idx];
        float confidence = vScores[idx];

        // 绘制边界框
        Scalar color = COLORS[classId % COLORS.size()];
        rectangle(frame, box, color, 2);
        // 显示标签和置信度
        string label = cv::format("%s: %.2f", CLASS_NAMES[classId].c_str(), confidence);
        int baseline;
        Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseline);
        rectangle(frame, Point(box.x, box.y - labelSize.height - 5), Point(box.x + labelSize.width, box.y), color, FILLED);
        putText(frame, label, Point(box.x, box.y - 5), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 0, 0), 1);
    }
}
int main()
{
    // 加载模型
    Net net = readNetFromONNX("E:/models/best.onnx");
    net.setPreferableBackend(DNN_BACKEND_OPENCV);
    net.setPreferableTarget(DNN_TARGET_CPU);

    // 加载图像 (640x160彩色图像)
    Mat frame = imread("E:/Samples160_640/1.bmp");
    if (frame.empty())
    {
        cerr << "Could not load image!" << endl;
        return -1;
    }

    // 预处理
    Mat blob;
    blobFromImage(frame, blob, 1.0 / 255.0, Size(640, 160), Scalar(), true, false);
    net.setInput(blob);
    // 前向传播
    vector<Mat> outputs;
    net.forward(outputs, net.getUnconnectedOutLayersNames());
    // 后处理
    postProcess(frame, outputs,0.6,0.4);
    // 显示结果
    imshow("YOLOv8 Detection", frame);
    waitKey(0);
    // 保存结果
    imwrite("output.jpg", frame);

    return 0;
}

 

【参考】

yolov8训练自己的数据集(简单最快上手版)-CSDN博客

posted @ 2020-04-28 11:36  夕西行  阅读(308)  评论(0)    收藏  举报