如何让YoloV5对实时屏幕抓取的画面执行目标检测

你不需要对detect.py做大幅修改，直接把屏幕抓取代码和YOLOv5推理逻辑整合即可，核心原理是YOLOv5的PyTorch Hub模型原生支持numpy数组格式的图像输入，只要做好抓取帧的格式适配就行。

完整可运行代码 #

import cv2
import numpy as np
from mss import mss
import torch

# 提前加载YOLOv5模型，放在循环外避免重复加载
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# 定义屏幕抓取区域
bounding_box = {'top': 340, 'left': 800, 'width': 350, 'height': 400}
sct = mss()

while True:
    # 抓取指定区域屏幕
    sct_img = sct.grab(bounding_box)
    # 转numpy数组，mss默认返回BGRA格式，去掉alpha通道转BGR适配OpenCV和YOLO
    scr_img = np.array(sct_img)[:, :, :3]

    # YOLO推理
    results = model(scr_img)

    # 渲染检测框到原图，得到带标注的图像
    annotated_img = results.render()[0]

    # 显示检测结果
    cv2.imshow('YOLO屏幕实时检测', annotated_img)

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

关键改动说明 #

• 模型加载放在循环最外层，避免每次推理都重载模型，大幅提升运行速度
• mss抓取的图像自带alpha通道，需要截取前3个通道转成BGR格式，避免识别结果异常、显示颜色错乱
• 调用results.render()直接生成带检测框、类别、置信度的标注图，不需要自己手动画框

针对detect.py的修改方案 #

如果你需要基于官方detect.py做改造，只需要修改数据源加载部分：

1. 去掉原有parse_opt()里的source参数相关的默认加载逻辑
2. 把上述屏幕抓取的循环逻辑替换掉detect.py里的帧遍历逻辑，每次抓取的屏幕帧作为输入喂给后续的推理、渲染、显示流程即可

实现效果 #

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