【Yolo5V6学习】（2）代码框架实现

【Yolo5V6学习】（2）代码框架实现

项目文件内容如下：

代码框架如下，代码中我会尽量写上详细的代码注释

// Author:  ciellee
// Date:    2022/03/01
#include <iostream>
#include <opencv2/opencv.hpp>

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

struct Output {
    int id;             // 检测结果类别 id
    float confidence;   // 检测结果置信度
    Rect box;           // 选框位置    
};

class Yolov5 {
public:
    Yolov5() {}
    ~Yolov5() {}
    bool ReadModel(Net& net, string& netPath, bool isCuda);
    bool Detect(Mat& SrcImg, Net& net, vector<Output>& output);
    void DrawPred(Mat& img, vector<Output> result, vector<Scalar> color);

public:
    // 锚框: 参考默认的 anchors.yaml 文件，采用 anchors_p5_640 中的数据,yolov5-6.0\models\hub\anchors.yaml
    // 三个不同分辨率的特征图上的锚框anchor，能够分别对大、中、小目标进行计算
    const float netAnchors[3][6] = { 
        { 10.0, 13.0, 16.0, 30.0, 33.0, 23.0 },     // P3 / 8
        { 30.0, 61.0, 62.0, 45.0, 59.0, 119.0 },    // P4 / 16
        { 116.0, 90.0, 156.0, 198.0, 373.0, 326.0 } // P5 / 32
    };
    const float netStride[3] = { 8.0, 16.0, 32.0 };

    const int netWidth = 640;   // 宽：图像尺度大小,必须为32的整数倍
    const int netHeight = 640;  // 高：图像尺度大小,必须为32的整数倍
    float nmsThreshold = 0.45;  // 交并比阈值,IOU值：预测框大小∩真实框大小 / 预测框大小∪真实框大小,预测框与真实框的交集与并集的取值。
        //越大，则容易将对于同一个物品的不同预测结果 当成 对多个物品的多个预测结果，导致一个物品出现了多个预测结果。
        //越小，则容易将对于多个物品的不同预测结果 当成 对同一个物品的不同预测结果，导致多个物品只出现了一个预测结果。
    float boxThreshold = 0.25;  
    float classThreshold = 0.25;    // 置信度阈值,只显示预测概率超过conf_thres的预测结果

    // 分类名
    std::vector<std::string> className = { 
        "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
        "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
        "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
        "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
        "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
        "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
        "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
        "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
        "hair drier", "toothbrush" 
    };
};


int main()
{
    string img_path = "bus.jpg";
    string model_path = "yolov5s.onnx";

    Yolov5 yolo;    // 创建Yolov5实例
    Net net;        // 创建神经网络实例

    // 读取模型
    if (yolo.ReadModel(net, model_path, false)) {
        cout << "read Net ok" << endl;
    }else{
        return -1;
    }
 
    // 生成80种随机颜色
    // cv::Scalar的构造函数是cv::Scalar(v1, v2, v3, v4)，前面的三个参数是依次设置BGR的，和RGB相反，第四个参数设置图片的透明度。
    // 当使用opencv提供的库函数imread()、imwrite()和imshow()时，cv::Scalar(v1, v2, v3, v4)的这四个参数就依次是BGRA，即蓝、绿、红和透明度。
    vector<Scalar> color;
    srand(time(0));
    for (int i = 0; i < 80; i++)
    {
        int b = rand() % 256;
        int g = rand() % 256;
        int r = rand() % 256;
        color.push_back(Scalar(b, g, r));
    }

    vector<Output> result;
    
    // 读取图片
    Mat img = imread(img_path);
    //imshow("img", img);       // 显示图片
    //waitKey(0);

    clock_t start, end;
    double duration;
    start = clock();

    // 开始预测
    if (yolo.Detect(img, net, result)) {
        // 预测成功，开始绘制选框
        yolo.DrawPred(img, result, color);
        
        end = clock();
        duration = (double)(end - start) / CLOCKS_PER_SEC;
        printf("\n预测耗时：%f s\n", duration);
    }
    else {
        cout << "Detect Failed!" << endl;
    }

    system("pause");
    return 0;
}

// 读取模型
bool Yolov5::ReadModel(Net& net, string& netPath, bool isCuda)
{
    return true;
}

// 预测函数
bool Yolov5::Detect(Mat& SrcImg, Net& net, vector<Output>& output)
{
    return true;
}

// 绘制选框
void Yolov5::DrawPred(Mat& img, vector<Output> result, vector<Scalar> color)
{

}

编译运行结果如下：

明天开始实现​​Yolov5::ReadModel()​​ 函数，加油