yolov5 opencv DNN 推理

opencv4.5.1

Ubuntu 21.10

切换到yolov5工程中，打开终端运行指令

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

onnx转onnx-sim

python -m onnxsim yolov5s.onnx yolov5s_sim.onnx

使用netron打开onnx文件

 最终的输出1*25200*85，将三个输出整合在一起即：

 25200=3*（20*20）+3*（40*40）+3*（80*80）个候选框

opecv DNN API:
1.读取模型
net = cv2.dnn.readNetFromONNX(model)
2.输入（类似Tensor）
blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (640, 640), swapRB=True, crop=False)
        swapRB：BGR->RGB
net.setInput(blob)
3.推理
preds = net.forward()
4.非极大值抑制
cv2.dnn.NMSBoxes(boxes,confidences,0.25,0.45)

import cv2
import numpy as np


classes=["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"]


def infer(img,pred,shape):
    w_ratio=img.shape[1]/shape[0]
    h_ratio=img.shape[0]/shape[1]
    # print(w_ratio,h_ratio)
    confidences=[]
    boxes=[]
    class_ids=[]
    boxes_num=pred.shape[1]
    data=pred[0]
    # print("data_shape:",data.shape)
    for i in range(boxes_num):
        da=data[i]#[box,conf,cls]
        confidence=da[4]
        if confidence>0.6:
            score=da[5:]*confidence
            _,_,_,max_score_index=cv2.minMaxLoc(score)#
            max_cls_id=max_score_index[1]
            if score[max_cls_id]>0.25:
                confidences.append(confidence)
                class_ids.append(max_cls_id)
                x,y,w,h=da[0].item(),da[1].item(),da[2].item(),da[3].item()
                nx=int((x-w/2.0)*w_ratio)
                ny=int((y-h/2.0)*h_ratio)
                nw=int(w*w_ratio)
                nh=int(h*h_ratio)
                boxes.append(np.array([nx,ny,nw,nh]))

    indexes=cv2.dnn.NMSBoxes(boxes,confidences,0.25,0.45)
    res_ids=[]
    res_confs=[]
    res_boxes=[]
    for i in indexes:
        res_ids.append(class_ids[i])
        res_confs.append(confidences[i])
        res_boxes.append(boxes[i])

    # print(res_ids)
    # print(res_confs)
    # print(res_boxes)
    return res_ids,res_confs,res_boxes

def draw_rect(img,ids,confs,boxes):
    for i in range(len(ids)):
        cv2.rectangle(img, boxes[i], (0,0,255), 2)
        cv2.rectangle(img, (boxes[i][0],boxes[i][1]-20),(boxes[i][0]+boxes[i][2],boxes[i][1]), (200, 200, 200), -1)
        cv2.putText(img, classes[ids[i]], (boxes[i][0], boxes[i][1] - 10), cv2.FONT_HERSHEY_SIMPLEX, .5, (255, 0, 0))
        cv2.putText(img, str(confs[i]), (boxes[i][0]+60, boxes[i][1] - 10), cv2.FONT_HERSHEY_SIMPLEX, .5, (255, 0, 0))

    cv2.imwrite("res.jpg",img)
    cv2.imshow('img',img)
    cv2.waitKey()


if __name__=="__main__":
    import time
    st=time.time()
    shape=(640,640)
    src=cv2.imread("../data/bus.jpg")
    img=src.copy()
    net=cv2.dnn.readNet("../model_m1/yolov5s_sim.onnx")
    blob=cv2.dnn.blobFromImage(img,1/255.,shape,swapRB=True,crop=False)
    net.setInput(blob)
    pred=net.forward()
    print(pred.shape)
    ids,confs,boxes=infer(img,pred,shape)
    et=time.time()
    print("run time:{:.2f}s/{:.2f}FPS".format(et-st,1/(et-st)))
    draw_rect(src,ids,confs,boxes)

运行结果：

cdnn.cpp

#include
#include
#include


void infer_res(cv::Mat& img,cv::Mat& preds,std::vector &boxes,std::vector&confidences,std::vector &classIds,std::vector&indexs)
{
    float w_ratio=img.cols/640.0;
    float h_ratio=img.rows/640.0;

    cv::Mat data(preds.size[1],preds.size[2],CV_32F,preds.ptr());
    for(int i=0;i(i,4);
        if(conf<0.45)
        {
            continue;
        }

        cv::Mat clsP=data.row(i).colRange(5,85)*conf;
        cv::Point IndexId;
        double score;
        minMaxLoc(clsP,0,&score,0,&IndexId);
        if(score>0.25)
        {
            float x=data.at(i,0);
            float y=data.at(i,1);
            float w=data.at(i,2);
            float h=data.at(i,3);

            int nx=int((x-w/2.0)*w_ratio);
            int ny=int((y-h/2.0)*h_ratio);
            int nw=int(w*w_ratio);
            int nh=int(h*h_ratio);

            cv::Rect box;
            box.x=nx;
            box.y=ny;
            box.width=nw;
            box.height=nh;
            boxes.push_back(box);
            classIds.push_back(IndexId.x);
            confidences.push_back(score);

        }

    }

//    std::vectorindexs;
    cv::dnn::NMSBoxes(boxes,confidences,0.25,0.45,indexs);

}
//void draw_label(cv::Mat& img,std::vector &boxes,std::vector&confidences,std::vector &classIds,std::vector&indexs,std::string& classes[])
//{
//    for(int i=0;iboxes;
    std::vectorconfidences;
    std::vector classIds;
    std::vectorindexs;
    infer_res(src,preds,boxes,confidences,classIds,indexs);
    clock_t et=clock();
    std::cout<<"run time:"<<(double)(et-st)/CLOCKS_PER_SEC< 
CMakeLists.txt
 
cmake_minimum_required(VERSION 2.8)
project(cdnn)
find_package(OpenCV 4.5.1 REQUIRED)
include_directories(${OpenCV_INCLUDES_DIRS})
add_executable(cdnn cdnn.cpp)
target_link_libraries(cdnn ${OpenCV_LIBS})
 
执行：
 
mkdir build
cd build
cmake ..
make
./cdnn
 
运行结果：
 
 

                
        
        

  
    
      
        
      
      
        
          暑期编程PK赛
          
        
        得CSDN机械键盘等精美礼品！