使用yolov5s.pt权重，结合opencv，实现实时检测

使用yolov5s.pt权重，结合opencv，实现实时检测

---

一、加载目标检测模型

• 加载模型函数
• 注意

torch.hub.load('./yolov5', 'custom',path='./yolov5/pretrained/yolov5s.pt', source='local')

二、读取摄像头视频并实时检测

1.open cv2读取摄像头视频

• open cv2读取摄像头信息，并显示

import cv2

# 读取摄像头
cap = cv2.VideoCapture(0)

while True:
	
	rec, frame = cap.read()
	
	frame = cv2.flip(frame, 1)
	
	cv2.imshow('demo', frame)

	if cv2.waitKey(10) & 0xFF == 27:
		break

cap.release()
cv2.destroyAllWindows()

2.读取yolov5s.pt权重信息

• 查看 github中官方读取的说法
  

3.读取yolov5s.pt权重信息person类

# Inference,进行推理, 返回 result, pandas格式结果
            results = self.yolo_detector(img_rgb)

            # 推理结果pd包含: xmin, ymin, xmax, ymax, confidence, class, name
            pd = results.pandas().xyxy[0]

            # 筛选所需要的目标类别
            person_list = pd[pd['name'] == 'person'].to_numpy()

三、类的实现重构代码

import torch
import cv2


class yolo_demo:

    def __init__(self):
        # 加载目标检测模型
        self.yolo_detector = torch.hub.load('./yolov5', 'custom',
                                            path='./yolov5/pretrained/yolov5s.pt',
                                            source='local')

        # 设置一下置信度阈值
        self.yolo_detector.conf = 0.5

    def main(self):

        # 读取视频文件
        cap = cv2.VideoCapture(0)

        while True:

            rec, frame = cap.read()

            # bgr 转为 rgb
            img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            # Inference,进行推理, 返回 result, pandas格式结果
            results = self.yolo_detector(img_rgb)

            # 推理结果pd包含: xmin, ymin, xmax, ymax, confidence, class, name
            pd = results.pandas().xyxy[0]

            # 筛选所需要的目标类别
            person_list = pd[pd['name'] == 'person'].to_numpy()

            # 遍历筛选的类别, 包含坐标前四个, 浮点类型, 需要转换成 int 类型: l:left  t:top
            for person in person_list:

                l, t, r, b = person[:4].astype('int')

                # 进行绘制边界框
                cv2.rectangle(frame, (l, t), (r, b), (0, 255, 0), 4)

                # 对画面进行缩放
                # frame = cv2.resize(frame, (1080, 1080))

                # 显示
                cv2.imshow('demo', frame)

            # 退出条件
            if cv2.waitKey(10) & 0xFF == 27:
                break

        cap.release()
        cv2.destroyAllWindows()


plate_demo = yolo_demo()
plate_demo.main()

四、显示结果