使用opencv的dnn模块来进行yolov5的目标检测

环境

• python 3.8
• pytorch 1.7 + cu110
• yolov5 6.1
• opencv 4.5.5

前言

在 yolov5 的 C++ 部署方案中， opencv 应该是最能被想到的一种，从 3.3 版本后， opencv 就加入了 dnn 这个模块，有了这个模块，很多的机器学习项目就可以通过它来实现部署了，下面我们就来看看具体的实现步骤。

yolov5

由于 opencv 无法直接读取 yolov5 中的 pt 模型文件，因此，需要将原来的 pt 文件转换成 opencv 能直接读取的 onnx 模型文件

使用目前最新的版本 v6.1， https://github.com/ultralytics/yolov5/releases/tag/v6.1，解压后，进入源码目录，下载对应的模型文件，这里使用 yolov5s.pt 为例，地址： https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s.pt

接下来使用源码中的转换脚本 export.py，将 pt 文件转成 onnx 格式

  python export.py --weights yolov5s.pt --include onnx

完成后就可以在当前目录下生成了 yolov5s.onnx 文件

cmake

来到官网， https://cmake.org/download/，下载后，傻瓜式安装，然后，将其也加入到系统环境变量中

opencv

接下来，我们配置一下 opencv 的环境，来到官网 https://opencv.org/releases/，下载 windows 对应的版本，现在最新的版本是 4.5.5

下载完成后解压

然后新增一个环境变量 OpenCV_DIR，对应的值就是解压后的文件夹下的 build 目录，如下

修改系统环境变量 Path，将 opencv 包中的 build\x64\vc\15\bin 添加进去

开始部署

我们创建一个 C++ 项目，包含了源码文件 yolov5.cpp、 CMakeLists.txt、目标名称文本文件 class.txt、 yolov5s.onnx 和若干测试图片

CMakeLists.txt 定义了编译的规则

  cmake_minimum_required(VERSION 3.20)
set(CMAKE_CXX_STANDARD 11)

project(yolov5   dnn)

find_package(OpenCV REQUIRED)

add_executable(yolov5   dnn yolov5.cpp)

include_directories(${OpenCV_INCLUDE_DIRS})
target_link_libraries(yolov5dnn ${OpenCV_LIBS} ${CMAKE_THREAD_LIBS_INIT})

class.txt 是待检测的目标，默认模型是 coco 的80个目标

最后来看源码文件 yolov5.cpp

  #include 
#include 

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

// 常量
const float INPUT_WIDTH = 640.0;
const float INPUT_HEIGHT = 640.0;
const float SCORE_THRESHOLD = 0.5;
const float NMS_THRESHOLD = 0.45;
const float CONFIDENCE_THRESHOLD = 0.45;

// 显示相关
const float FONT_SCALE = 0.7;
const int FONT_FACE = FONT_HERSHEY_SIMPLEX;
const int THICKNESS = 1;
Scalar BLACK = Scalar(0,0,0);
Scalar BLUE = Scalar(255, 178, 50);
Scalar YELLOW = Scalar(0, 255, 255);
Scalar RED = Scalar(0,0,255);

// 画框函数
void draw_label(Mat& input_image, string label, int left, int top)
{
    int baseLine;
    Size label_size = getTextSize(label, FONT_FACE, FONT_SCALE, THICKNESS, &baseLine);
    top = max(top, label_size.height);
    Point tlc = Point(left, top);
    Point brc = Point(left + label_size.width, top + label_size.height + baseLine);
    rectangle(input_image, tlc, brc, BLACK, FILLED);
    putText(input_image, label, Point(left, top + label_size.height), FONT_FACE, FONT_SCALE, YELLOW, THICKNESS);
}

// 预处理
vector pre_process(Mat &input_image, Net &net)
{
    Mat blob;
    blobFromImage(input_image, blob, 1./255., Size(INPUT_WIDTH, INPUT_HEIGHT), Scalar(), true, false);

    net.setInput(blob);

    vector outputs;
    net.forward(outputs, net.getUnconnectedOutLayersNames());

    return outputs;
}

// 后处理
Mat post_process(Mat &input_image, vector &outputs, const vector &class_name) 
{
    vector class_ids;
    vector confidences;
    vector boxes; 

    float x_factor = input_image.cols / INPUT_WIDTH;
    float y_factor = input_image.rows / INPUT_HEIGHT;

    float *data = (float *)outputs[0].data;

    const int dimensions = 85;
    const int rows = 25200;
    for (int i = 0; i < rows; ++i) 
    {
        float confidence = data[4];
        if (confidence >= CONFIDENCE_THRESHOLD) 
        {
            float * classes_scores = data + 5;
            Mat scores(1, class_name.size(), CV_32FC1, classes_scores);
            Point class_id;
            double max_class_score;
            minMaxLoc(scores, 0, &max_class_score, 0, &class_id);
            if (max_class_score > SCORE_THRESHOLD) 
            {
                confidences.push_back(confidence);
                class_ids.push_back(class_id.x);

                float cx = data[0];
                float cy = data[1];
                float w = data[2];
                float h = data[3];
                int left = int((cx - 0.5 * w) * x_factor);
                int top = int((cy - 0.5 * h) * y_factor);
                int width = int(w * x_factor);
                int height = int(h * y_factor);
                boxes.push_back(Rect(left, top, width, height));
            }

        }
        data += 85;
    }

    vector indices;
    NMSBoxes(boxes, confidences, SCORE_THRESHOLD, NMS_THRESHOLD, indices);
    for (int i = 0; i < indices.size(); i++) 
    {
        int idx = indices[i];
        Rect box = boxes[idx];

        int left = box.x;
        int top = box.y;
        int width = box.width;
        int height = box.height;
        rectangle(input_image, Point(left, top), Point(left + width, top + height), BLUE, 3*THICKNESS);

        string label = format("%.2f", confidences[idx]);
        label = class_name[class_ids[idx]] + ":" + label;
        draw_label(input_image, label, left, top);
    }
    return input_image;
}

// 主函数
int main(int argc, char **argv)
{
    vector class_list;
    ifstream ifs("class.txt");
    string line;

    while (getline(ifs, line))
    {
        class_list.push_back(line);
    }

    Mat frame;
    frame = imread(argv[1]);

    Net net;
    net = readNet("yolov5s.onnx");

    vector detections;
    detections = pre_process(frame, net);

    Mat img = post_process(frame.clone(), detections, class_list);

    vector layersTimes;
    double freq = getTickFrequency() / 1000;
    double t = net.getPerfProfile(layersTimes) / freq;
    string label = format("Inference time : %.2f ms", t);
    putText(img, label, Point(20, 40), FONT_FACE, FONT_SCALE, RED);

    imshow("Output", img);
    waitKey(0);

    return 0;
}

一切准备就绪，马上开始编译，进入到工程目录，打开 powershell，然后一次执行

  mkdir build
cmake -B build
cmake --build build --config Release

编译成功后，就可以找图片来测试了

  .\build\Release\yolov5dnn.exe bus.jpg

可以看到图片的检测时间花费了210毫秒，显然是没有用到 gpu，如果需要启用 gpu 加速的话，就需要重新编译 opencv，使能 cuda，这部分内容可以去参考 windows 编译 opencv,支持 cuda 加速 和 ubuntu 下编译 opencv，支持 cuda 加速