如何在YOLOv5中对mss获取的游戏截图进行自定义字符检测？

1. 安装依赖 #

确保已安装所需库，执行以下命令：

pip install yolov5 mss numpy opencv-python torch

2. 完整实现代码 #

import mss
import numpy as np
import cv2
import torch

# 加载自定义训练的YOLOv5模型，替换为你的模型路径
model = torch.hub.load('ultralytics/yolov5', 'custom', path='./best.pt')  # 示例路径，根据实际修改

# 设置截图区域，可自定义调整坐标和尺寸
# 示例为全屏截图，若需指定区域，修改top/left/width/height即可
monitor = {"top": 0, "left": 0, "width": 1920, "height": 1080}

# 初始化mss截图对象
with mss.mss() as sct:
    while True:
        # 获取屏幕截图
        screenshot = sct.grab(monitor)
        # 将mss截图转为numpy数组，格式为BGRA
        img = np.array(screenshot)
        # 转为cv2支持的BGR格式
        img_bgr = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
        
        # 用YOLOv5进行检测
        results = model(img_bgr)
        
        # 将检测结果渲染到图像上（自动绘制边界框和标签）
        results.render()
        
        # 转换渲染后的图像格式（YOLOv5返回的是RGB，需转BGR给cv2）
        rendered_img = cv2.cvtColor(np.squeeze(results.imgs), cv2.COLOR_RGB2BGR)
        
        # 展示检测结果
        cv2.imshow("Game Character Detection", rendered_img)
        
        # 按下q键退出循环
        if cv2.waitKey(1) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            break

关键细节说明 #

• 模型路径：务必将path='./best.pt'替换为你自己训练好的模型文件路径（通常是best.pt或last.pt）。
• 截图区域：根据游戏窗口位置调整monitor的参数，比如游戏窗口左上角坐标是(100, 100)，宽高为800x600，就设置monitor = {"top": 100, "left": 100, "width": 800, "height": 600}。
• 格式转换：mss截图默认是BGRA格式，必须转为BGR才能被cv2正常处理；YOLOv5渲染后的图像是RGB格式，也需要转回BGR再用cv2.imshow展示。
• 性能优化：如果检测帧率过低，可以缩小截图区域尺寸，或者使用YOLOv5的轻量化模型（如yolov5s）进行训练。

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