自定义YOLOv7-tiny模型屏幕目标检测绘图报错修复咨询

• 已使用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)

推测问题源于截图转数组后，图像数组格式与OpenCV矩阵不兼容，导致无法绘制矩形框，求解决方法。

问题根源 #

mss截图生成的数组是BGRA格式（包含alpha透明通道），而YOLOv7的绘图函数依赖OpenCV的BGR格式（无alpha通道），格式不匹配导致OpenCV绘制矩形时抛出参数错误。同时原代码存在缩进问题，while循环在mss上下文外，可能导致后续截图异常。

修改步骤及代码 #

1. 将while循环移入with mss.mss()代码块内，保证截图上下文持续有效。
2. 转换截图格式：用cv2.cvtColor将BGRA转为BGR，去掉alpha通道。
3. 优化退出逻辑，用OpenCV原生的按键检测替代第三方库。

修改后的完整代码：

import mss
import numpy as np
import cv2
import time
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))
        # 转换BGRA格式到OpenCV兼容的BGR格式
        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)

        results = model(img)

        # render()返回图像列表，直接取第一个元素显示
        cv2.imshow('s', results.render()[0])

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

        # 按q键退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

cv2.destroyAllWindows()

关键修改说明 #

• 格式转换：cv2.COLOR_BGRA2BGR将4通道的BGRA转为3通道的BGR，完全匹配OpenCV和YOLOv7绘图函数的格式要求。
• 缩进修正：确保mss的截图对象sct在循环过程中始终处于有效上下文，避免潜在的资源释放问题。
• 简化显示：results.render()返回的是包含绘制后图像的列表，直接索引第一个元素即可，无需额外的np.squeeze处理。

内容的提问来源于stack exchange，提问作者ahmedkhaleel2004