AI学习笔记三：编写检测的yolov5测试代码

若该文为原创文章，转载请注明原文出处。

通过detect.py代码测试通过后，阅读detect.py代码发现，有些难以看懂，看得有点蒙蒙的，

所以编写了一个简单的测试程序。

代码如下：

import cv2
import numpy as np
import torch
import time
import pandas as pd

class My_detector:

    def __init__(self):
        # 加载model
        self.model = torch.hub.load('E:/desktop/yolov5-5.0/', 'custom',
                                    'E:/desktop/yolov5-5.0/weights/yolov5s.pt', source='local') 
                                    
        self.model.conf = 0.4

        # 打开摄像头或打开图片
        self.cap = cv2.VideoCapture('./data/images/zidane.jpg')
        #self.cap = cv2.VideoCapture(0)

    def detect(self):
        while True:
            # 读取图片
            ret, frame = self.cap.read()

            if frame is None:
                break
            
            # 翻转
            frame = cv2.flip(frame, 1)
            # 颜色转换
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            # 推理
            results = self.model(frame_rgb)

            
            results_list = results.pandas().xyxy[0]
            print('pd:')
            print(results_list)
            
            # 只要name为person的数据
            person_list = results_list[results_list['name']=='person'].to_numpy()
            print(person_list)

            for box in person_list:
                l, t, r, b = box[:4].astype('int')
                conf = box[4]
                conf_txt =str(round(conf*100,1) ) + '%'
                name = box[6]
                
                cv2.rectangle(frame, (l, t), (r, b), (0, 255, 0), 5)
                cv2.putText(frame, conf_txt, (l,t-35), cv2.FONT_ITALIC,1,(0,255,0),2)           
                cv2.putText(frame, name, (l,t-70), cv2.FONT_ITALIC,1,(0,255,0),2)   
                

            #cv2.imshow('DEMO', frame)
            cv2.imwrite('./result.jpg', frame)

            if cv2.waitKey(10) & 0xFF == ord('q'):
                break
                
        self.cap.release()
        cv2.destroyAllWindows()


detector = My_detector()
detector.detect()

运行后，会把结果保存成results.jpg图片

打印的信息参数有：

xmin  ymin xmax  ymax：  对应了坐标

class： 类别

name： 名字
 

 执行时，如果提示一些module没有安装，直接使用pip install安装。

如有侵权，或需要完整代码，请及时联系博主。