python 采用OpenCV的Yolov5

li9yvcax 于 2022-12-25 发布在 Python

关注(0)|答案(1)|浏览(208)

I trained a model using the Yolov5 Google Colab notebook with custom data and classes. I used the detect.py script to do object detection on a video with no issues.
我下载了. pt文件，并将其加载到我存储在机器上的脚本中，该脚本基于Yolov5的Github截图推理示例，当我在PC上的图像上使用它时，它工作得很好：https://github.com/ultralytics/yolov5/issues/36
现在我知道了PyTorch文件可以毫无问题地加载，并且检测工作正常，我尝试通过截图来执行实时检测。我可以通过截图来记录我的屏幕，但当我尝试对模型这样做时，我得到了一个错误。以下是我的代码：

import time
import torch
import cv2
import mss
import numpy

# Load model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='last_210_epochs_Size1280.pt') 

with mss.mss() as sct:
    # Part of the screen to capture
    monitor = sct.monitors[2]

    while "Screen capturing":
        last_time = time.time()

        # Get raw pixels from the screen, save it to a Numpy array
        img = numpy.array(sct.grab(monitor))

        results = model(img)
        

        # Display the picture
        cv2.imshow('test', results)

        # Display the picture in grayscale
        # cv2.imshow('OpenCV/Numpy grayscale',
        #            cv2.cvtColor(img, cv2.COLOR_BGRA2GRAY))

        print("fps: {}".format(1 / (time.time() - last_time)))

        # Press "q" to quit
        if cv2.waitKey(25) & 0xFF == ord("q"):
            cv2.destroyAllWindows()
            break

这是我得到的错误：

不知道我做错了什么。

python

来源：https://stackoverflow.com/questions/74648336/yolov5-with-opencv

1条答案

按热度按时间

t9aqgxwy1#

希望对你有帮助。

import cv2
import numpy as np
import torch
import imutils

def draw_text(img, text,
              font=cv2.FONT_HERSHEY_COMPLEX_SMALL ,
              pos=(0, 0),
              font_scale=1,
              font_thickness=1,
              text_color=(255, 255, 255),
              text_color_bg=(0, 0, 0)
              ):

    x, y = pos
    text_size, _ = cv2.getTextSize(text, font, font_scale, font_thickness)
    text_w, text_h = text_size
    cv2.rectangle(img, pos, (x + text_w + 5, y + text_h+ 5), text_color_bg, -1)
    cv2.putText(img, text, (x, y + text_h + font_scale - 1), font, font_scale, text_color, font_thickness)


model = torch.hub.load('ultralytics/yolov5', 'custom', path='last_210_epochs_Size1280.pt')

# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
out = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 30.0, (640, 360))

with mss.mss() as sct:
    # Part of the screen to capture
    monitor = sct.monitors[2]

    while "Screen capturing":
        last_time = time.time()

        # Get raw pixels from the screen, save it to a Numpy array
        img = numpy.array(sct.grab(monitor))

    `   frame = imutils.resize(img, width=640)
        results = model([frame], size=640)

        if results.pandas().xyxy[0].empty is not True:
            bbox = results.pandas().xyxy[0]
    
            for box in zip(bbox.xmin, bbox.ymin, bbox.xmax, bbox.ymax, bbox.confidence, bbox.name):
                if(box[4] >= 0.7):
    
                    draw_text(frame, str(box[5]), font_scale=1, pos=(int(box[0]), int(box[1])), text_color_bg=(0, 0, 0))
                    cv2.rectangle(frame, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (79, 255, 0), 2)
    
        # Write the frame into the file 'output.avi'
        out.write(frame)
    
        # Display the resulting frame
        cv2.imshow('frame',frame)
    
        # Press Q on keyboard to stop recording
        if cv2.waitKey(100) & 0xFF == ord('q'):
            break`

赞(0）回复(0）举报 2022-12-25

我来回答

python 采用OpenCV的Yolov5

1条答案

相关问题

热门标签

最新问答