OpenCV4.5.x DNN + YOLOv5 C++推理

这个预测时间190ms,应该是cpu版本

昨天修改了个OpenCV DNN支持部署YOLOv5,6.1版本的Python代码,今天重新转换为C++代码了!貌似帧率比之前涨了点!说明C++的确是比Python快点!

点击这里可以查看之前的推文:

OpenCV4.5.4 直接支持YOLOv5 6.1版本模型推理

OpenC4 C++部署YOLOv5

我把测试代码封装成一个工具类了,可以直接用,方便大家(生手党)直接部署调用!保重一行代码都不用再写了!

01

类的声明:

定义了一个结构体作为返回结果,主要包括类别id、置信度、检测框。两个方法分别是初始化参数与网络,另外一个完成检测功能与返回结果集。

yolov5_dnn.h

#pragma once

#include 
struct DetectResult {
    int classId;
    float score;
    cv::Rect box;
};

class YOLOv5Detector {
public:
    void initConfig(std::string onnxpath, int iw, int ih, float threshold);
    void detect(cv::Mat & frame, std::vector &result);
private:
    int input_w = 640;
    int input_h = 640;
    cv::dnn::Net net;
    int threshold_score = 0.25;
};

02

类实现:

直接读取YOLOv5 onnx格式模型,完成对图象预处理,模型推理,后处理返回等操作!代码实现如下:

#include 

void YOLOv5Detector::initConfig(std::string onnxpath, int iw, int ih, float threshold) {
    this->input_w = iw;
    this->input_h = ih;
    this->threshold_score = threshold;
    this->net = cv::dnn::readNetFromONNX(onnxpath);
}

void YOLOv5Detector::detect(cv::Mat & frame, std::vector &results) {
    // 图象预处理 - 格式化操作
    int w = frame.cols;
    int h = frame.rows;
    int _max = std::max(h, w);
    cv::Mat image = cv::Mat::zeros(cv::Size(_max, _max), CV_8UC3);
    cv::Rect roi(0, 0, w, h);
    frame.copyTo(image(roi));

    float x_factor = image.cols / 640.0f;
    float y_factor = image.rows / 640.0f;

    // 推理
    cv::Mat blob = cv::dnn::blobFromImage(image, 1 / 255.0, cv::Size(this->input_w, this->input_h), cv::Scalar(0, 0, 0), true, false);
    this->net.setInput(blob);
    cv::Mat preds = this->net.forward();

    // 后处理, 1x25200x85
    // std::cout << "rows: "<< preds.size[1]<< " data: " << preds.size[2] << std::endl;
    cv::Mat det_output(preds.size[1], preds.size[2], CV_32F, preds.ptr());
    float confidence_threshold = 0.5;
    std::vector boxes;
    std::vector classIds;
    std::vector confidences;
    for (int i = 0; i < det_output.rows; i++) {
        float confidence = det_output.at(i, 4);
        if (confidence < 0.45) {
            continue;
        }
        cv::Mat classes_scores = det_output.row(i).colRange(5, 85);
        cv::Point classIdPoint;
        double score;
        minMaxLoc(classes_scores, 0, &score, 0, &classIdPoint);

        // 置信度 0~1之间
        if (score > this->threshold_score)
        {
            float cx = det_output.at(i, 0);
            float cy = det_output.at(i, 1);
            float ow = det_output.at(i, 2);
            float oh = det_output.at(i, 3);
            int x = static_cast((cx - 0.5 * ow) * x_factor);
            int y = static_cast((cy - 0.5 * oh) * y_factor);
            int width = static_cast(ow * x_factor);
            int height = static_cast(oh * y_factor);
            cv::Rect box;
            box.x = x;
            box.y = y;
            box.width = width;
            box.height = height;

            boxes.push_back(box);
            classIds.push_back(classIdPoint.x);
            confidences.push_back(score);
        }
    }

    // NMS
    std::vector indexes;
    cv::dnn::NMSBoxes(boxes, confidences, 0.25, 0.45, indexes);
    for (size_t i = 0; i < indexes.size(); i++) {
        DetectResult dr;
        int index = indexes[i];
        int idx = classIds[index];
        dr.box = boxes[index];
        dr.classId = idx;
        dr.score = confidences[index];
        cv::rectangle(frame, boxes[index], cv::Scalar(0, 0, 255), 2, 8);
        cv::rectangle(frame, cv::Point(boxes[index].tl().x, boxes[index].tl().y - 20), 
            cv::Point(boxes[index].br().x, boxes[index].tl().y), cv::Scalar(0, 255, 255), -1);
        results.push_back(dr);
    }

    std::ostringstream ss;
    std::vector layersTimings;
    double freq = cv::getTickFrequency() / 1000.0;
    double time = net.getPerfProfile(layersTimings) / freq;
    ss << "FPS: " << 1000 / time << " ; time : " << time << " ms";
    putText(frame, ss.str(), cv::Point(20, 40), cv::FONT_HERSHEY_PLAIN, 2.0, cv::Scalar(255, 0, 0), 2, 8);
}

03

代码调用演示

直接调用,使用视频测试!先创建检测对象,然后初始化、就可以循环调用检测方法,完成图象或者视频检测。然后检查返回结果!

std::shared_ptr detector(new YOLOv5Detector());
detector->initConfig("D:/python/yolov5-6.1/yolov5s.onnx", 640, 640, 0.25f);
cv::VideoCapture capture("D:/images/video/Boogie_Up.mp4");
cv::Mat frame;
std::vector results;
while (true) {
    bool ret = capture.read(frame);
    detector->detect(frame, results);
    for (DetectResult dr : results) {
        cv::Rect box = dr.box;
        cv::putText(frame, classNames[dr.classId], cv::Point(box.tl().x, box.tl().y - 10), cv::FONT_HERSHEY_SIMPLEX, .5, cv::Scalar(0, 0, 0));
    }
    cv::imshow("YOLOv5-6.1 + OpenCV DNN - by gloomyfish", frame);
    char c = cv::waitKey(1);
    if (c == 27) { // ESC 退出
        break;
    }
    // reset for next frame
    results.clear();
}

04

运行结果:

OpenCV4.5.x DNN + YOLOv5 C++推理_第1张图片

OpenCV4.5.x DNN + YOLOv5 C++推理_第2张图片

你可能感兴趣的:(c++入门宝典,c++,计算机视觉,算法)