基于C++的YOLOV5在TX2-NX上的实现（二）

接上篇，本篇主要结合代码解释及说明具体实现过程。

一、常量参数设置

const vector<Scalar> colors = {Scalar(255,255,0),Scalar(0, 255, 0),Scalar(0, 255, 255),Scalar(255, 0, 0)};
const float INPUT_WIDTH = 640.0;
const float INPUT_HEIGHT = 640.0;
const float SCORE_THRESHOLD = 0.2;
const float NMS_THRESHOLD = 0.4;
const float CONFIDENCE_THRESHOLD = 0.4;

colors是一个常量容器，可以理解成一个颜色集合，用于表示识别到的不同物体。

INPUT_WIDTH和INPUT_HEIGHT是输入视频图像的宽和高。

CONFIDENCE_THRESHOLD是检测框的置信阈值，当数值大于0.4时，检测框进入下一步处理流程，否则丢弃。

SCORE_THRESHOLD是类别置信阈值，当数值大于0.2时，认为类别识别结果可信，否则丢弃

NMS_THRESHOLD是做检测框非极大值抑制的阈值，IOU（交并比）大于0.4的候选检测框会被删除。

上述常量参数都可以根据实际应用情况进行调整。

这里补充一下nms的基本处理流程

1、根据置信度得分进行排序
2、选择置信度最高的比边界框添加到最终输出列表中，将其从边界框列表中删除
3、计算所有边界框的面积
4、计算置信度最高的边界框与其它候选框的IoU
5、删除IoU大于阈值的边界框
6、重复上述过程，直至边界框列表为空

二、模型加载

void load_net(dnn::Net &net, bool is_cuda)
{
    auto result = dnn::readNet("yolov5s.onnx");
    if (is_cuda)
    {
        cout<<"Attempt to use CUDA\n";
        result.setPreferableBackend(dnn::DNN_BACKEND_CUDA);
        result.setPreferableTarget(dnn::DNN_TARGET_CUDA_FP16);
    }
    else
    {
        cout<<"Running on CPU\n";
        result.setPreferableBackend(dnn::DNN_BACKEND_OPENCV);
        result.setPreferableTarget(dnn::DNN_TARGET_CPU);
    }
    net = result;
}

主要是利用cv中dnn模块的readNet方法读取模型结构和参数，这个方法在opencv4.5之后才有，所以安装环境时需要注意下cv的版本。

另外，在服务器上训练好的模型，模型结构和参数需要保存到同一个文件中，然后通过工具转换成onnx格式，如果配置过cv-cuda，可以使用cuda进行模型推理加速。

三、模型输入输出类型定义

Mat format_yolov5(const Mat &source)
{
    int col = source.cols;
    int row = source.rows;
    int _max = MAX(col,row);
    Mat result = Mat::zeros(_max,_max,CV_8UC3);//w*h*c
    source.copyTo(result(Rect(0,0,col,row)));//x,y,w,h
    return result;
}

定义一个Mat类型的对象result，将输入图像copy到该对象result，CV_8UC3表示每个元素存储的方式：8位无符号char类型3通道数据。

struct Detection
{
    int class_id;
    float confidence;
    Rect box;
};

定义一个结构体Detection用于存放输出结果，包括类别、置信度、候选框等。

四、检测逻辑

void detect(Mat &image,dnn::Net &net,vector<Detection> &output,vector<string> &className)
{
    Mat blob;
    auto input_image = format_yolov5(image);
    dnn::blobFromImage(input_image,blob,1./255,Size(INPUT_WIDTH,INPUT_HEIGHT),Scalar(),true,false);
    net.setInput(blob);
    vector<Mat> outputs;//定义mat类型的网络输出结果
    net.forward(outputs,net.getUnconnectedOutLayersNames());//网络前向推理
    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;//80个类+xywh+score
    const int rows = 25200;//总输出的候选框个数
    vector<int> class_ids;
    vector<float> confidences;
    vector<Rect> boxes;

    for (int i=0;i<rows;++i)
    {
        float confidence = data[4];
        if (confidence >= CONFIDENCE_THRESHOLD)
        {
            float * classes_scores = data + 5;
            Mat scores(1,className.size(),CV_32FC1,classes_scores);//1*80*1*80
            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 x = data[0];
                float y = data[1];
                float w = data[2];
                float h = data[3];
                int left = int((x-0.5*w)*x_factor);
                int top = int((y-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<int> nms_result;
    dnn::NMSBoxes(boxes,confidences,SCORE_THRESHOLD, NMS_THRESHOLD, nms_result);
    for (int i =0;i<nms_result.size();i++)
    {
        int idx = nms_result[i];
        Detection result;
        result.class_id = class_ids[idx];
        result.box = boxes[idx];
        result.confidence = confidences[idx];
        output.push_back(result);
    }

}

需要注意的是输入图像需要进行归一化，x_factor和y_factor是把检测框对应到原图尺寸。

五、主逻辑

int main(int argc,char **argv)
{
    vector<string> class_list = load_class_list();
    Mat frame;
    //VideoCapture capture("sample.mp4");
    VideoCapture capture(0);
    if (!capture.isOpened())
    {
        cerr<<"Error opening video file\n";
        return -1;
    }
    bool is_cuda = argc > 1 && strcmp(argv[1],"cuda") == 0;
    dnn::Net net;
    load_net(net,is_cuda);
    auto start = chrono::high_resolution_clock::now();
    int frame_count = 0;
    float fps = -1;
    int total_frames = 0;
    while (true)
    {
        capture.read(frame);

        /*
        if (frame.empty())
        {
            cout<<"End of stream\n";
            break;
        }
        */
        if (waitKey(50) == 27)//esc
        {
            
            break;
        }
        vector<Detection> output;
        detect(frame,net,output,class_list);
        frame_count++;
        total_frames++;
        int detections = output.size();
        for (int i=0;i<detections;++i)
        {
            auto detection = output[i]; 
            auto box = detection.box;
            auto class_Id = detection.class_id;
            const auto color = colors[class_Id % colors.size()];
            rectangle(frame,box,color,3);
            rectangle(frame,Point(box.x,box.y-20),Point(box.x+box.width,box.y),color,FILLED);
            putText(frame,class_list[class_Id].c_str(),Point(box.x,box.y-5),FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 0, 0));

        }
        if (frame_count>=30)
        {
            auto end = chrono::high_resolution_clock::now();
            fps = frame_count * 1000.0 / chrono::duration_cast<chrono::milliseconds>(end - start).count();

            frame_count = 0;
            start = chrono::high_resolution_clock::now();
        }
        if (fps > 0)
        {
            ostringstream fps_label;
            fps_label << fixed << setprecision(2);
            fps_label << "FPS: " << fps;
            string fps_label_str = fps_label.str();

            //putText(frame, fps_label_str.c_str(), Point(10, 25), FONT_HERSHEY_SIMPLEX, 1, Scalar(0, 0, 255), 2);
            putText(frame, fps_label_str, Point(10, 25), FONT_HERSHEY_SIMPLEX, 1, Scalar(0, 0, 255), 2);
        }
        imshow("output",frame);

    }
    cout<<"Total frames:"<<total_frames<<endl;
    return 0;
}

六、编译执行

g++ --std=c++11 -o yolotest yolo.cpp `pkg-config --cflags --libs opencv4`

之后，通过./yolotest执行，可以看到实时的检测结果。

相关资源下载链接模型和可执行文件