Use yolov5s.pt weights, combined with opencv, to achieve real-time detection

Use yolov5s.pt weights, combined with opencv, to achieve real-time detection

---

1. Load the target detection model

• load model function
• Notice

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

2. Read the camera video and detect it in real time

1.open cv2 read camera video

• open cv2 reads the camera information and displays it

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. Read yolov5s.pt weight information

• View the official reading in github
  

3. Read yolov5s.pt weight information person class

# 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()

3. Class implementation refactoring code

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()

4. Display the results