Home > Net >  Using custom yolov7 trained model on my screen
Using custom yolov7 trained model on my screen

Time:11-28

What I know

I have already trained a custom model using yolov7-tiny

I am now trying to use it for object detection on screen

The script I have:

import mss
import numpy as np
import cv2
import time
import keyboard
import torch
from hubconf import custom

model = custom(path_or_model='yolov7-tiny-custom.pt')

with mss.mss() as sct:
    monitor = {'top': 30, 'left': 0, 'width': 1152, 'height': 864}
while True:
    t = time.time()

    img = np.array(sct.grab(monitor))

    results = model(img)

    cv2.imshow('s', np.squeeze(results.render()))

    print('fps: {}'.format(1 / (time.time() - t)))

    cv2.waitKey(1)

    if keyboard.is_pressed('q'):
        break

cv2.destroyAllWindows()

The problem

I know that everything works in that script, however when it finally detects an object, it wants to draw a rectangle on screen.

I receive the following error:

Traceback (most recent call last):
  File "c:\Users\ahmed\Desktop\PC\Repos\yolov7-custom\yolov7-custom\aimbot.py", line 20, in <module>
    cv2.imshow('s', np.squeeze(results.render()))
  File "c:\Users\ahmed\Desktop\PC\Repos\yolov7-custom\yolov7-custom\models\common.py", line 990, in render
    self.display(render=True)  # render results
  File "c:\Users\ahmed\Desktop\PC\Repos\yolov7-custom\yolov7-custom\models\common.py", line 964, in display
    plot_one_box(box, img, label=label, color=colors[int(cls) % 10])
  File "c:\Users\ahmed\Desktop\PC\Repos\yolov7-custom\yolov7-custom\utils\plots.py", line 62, in plot_one_box
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
cv2.error: OpenCV(4.6.0) :-1: error: (-5:Bad argument) in function 'rectangle'
> Overload resolution failed:
>  - Layout of the output array img is incompatible with cv::Mat
>  - Expected Ptr<cv::UMat> for argument 'img'
>  - argument for rectangle() given by name ('thickness') and position (4)
>  - argument for rectangle() given by name ('thickness') and position (4)

Summary

I'm not too sure whats going on here but I believe that when I take the screenshot, I convert it to an array and apply my model. When it wants to draw a rectangle, its unable to do so because the output array img is incompatible with OpenCV matrix. How can I fix that?

CodePudding user response:

I tried to reproduce your issue and combined your code with an available yolo-demo, but I couldn't find any issue that would return an error message like that in your question. You can check it in your environment:

import mss
import numpy as np
import cv2
import torch
import time


model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

with mss.mss() as sct:
    monitor = {'top': 50, 'left': 50, 'width': 600, 'height': 400}

    while True:
        t = time.time()

        img = np.array(sct.grab(monitor))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        results = model(img)
        results.render()
        out = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        cv2.imshow('s', out)

        print('fps: {}'.format(1 / (time.time() - t)))

        if cv2.waitKey(1) == 27:
            break

cv2.destroyAllWindows()

Output:

enter image description here

  • Related