Я уже обучил пользовательскую модель с помощью yolov7-tiny.
Теперь я пытаюсь использовать его для обнаружения объектов на экране.
Скрипт у меня есть:
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()
Я знаю, что в этом сценарии все работает, однако, когда он, наконец, обнаруживает объект, он хочет нарисовать прямоугольник на экране.
Я получаю следующую ошибку:
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)
Я не совсем уверен, что здесь происходит, но я считаю, что когда я делаю снимок экрана, я преобразовываю его в массив и применяю свою модель. Когда он хочет нарисовать прямоугольник, он не может этого сделать, потому что выходной массив img несовместим с матрицей OpenCV. Как я могу это исправить?
Я попытался воспроизвести вашу проблему и объединил ваш код с доступной демонстрацией yolo, но я не смог найти ни одной проблемы, которая возвращала бы сообщение об ошибке, подобное этому в вашем вопросе. Вы можете проверить это в своей среде:
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()
Выход:
Спасибо вам большое за это! Ваш код работает, похоже, я неправильно вводил данные изображения. Я не слишком знаком с cv2.cvtColor
или cv2.COLOR_BGR2RGB
, но я обязательно изучу это.
Ваша пользовательская модель работает с изображениями, которые просто читаются
cv2.imread
?