yolov5 目标检测代码

 

import torch
import cv2
import time
import os

box_color = (0, 255, 0)

def get_milsecond():
    t = time.time()
    return (int(round(t * 1000)))

if __name__ == '__main__':
    # 模型
    model = torch.hub.load('F:\\openvino+yolov5\\yolov5-6.1', 'custom', 'yolov5s.pt', source='local') # yolov5s.pt best.pt

    url = 'bus.jpg' #'smoking.mp4'

    cap = cv2.VideoCapture(url)

    while True:
        ret, frame = cap.read()
        if not ret or frame is None:
            break

        starttime = get_milsecond()
        results = model(frame)
        pd = results.pandas().xyxy[0]

        for obj in pd.to_numpy():
            box_l, box_t = int(obj[0]), int(obj[1])
            box_r, box_b = int(obj[2]), int(obj[3])
            confidence = obj[4]
            classid = obj[5]
            classname = obj[6]
            frame = cv2.rectangle(frame, (box_l, box_t), (box_r, box_b), box_color, 2)
            frame = cv2.putText(frame, str(classid) + " " + str(classname) + " " + str(confidence), (box_l, box_t + 15), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 2)

        cv2.imshow('result', frame)
        endtime = get_milsecond();
        print('检测耗时:' + str(endtime - starttime))

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

    os.system('pause')