yolov5 最简单的测试代码

最简单的测试代码

下载好yolov5，安装好环境后，拥有权重文件，修改自己的路径，即可使用。

#
import torch
import cv2

# 加载：'./yolov5'文件夹路径、'custom'自定义的加载模式、权重文件路径、'local'本地加载
model = torch.hub.load('./yolov5', 'custom', path='./weights/ppe_yolo_n.pt', source='local')

cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()

    # 转为RGB
    img_cvt = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 推理
    results = model(img_cvt)
    # 查看内容 转为numpy格式
    pd = results.pandas().xyxy[0].to_numpy()

    # 遍历每个类别
    for box in pd:
        x, y, w, h = box[:4].astype('int')
        label_id = box[5]
        cv2.rectangle(frame, (x, y), (w, h), (0, 255, 0), 5)

        cv2.imshow('demo', frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break