python截取电脑屏幕(搭配yolo使用)

无论是视频还是摄像头，YOLO的目标检测实际上都是对图像进行预测的因此对于电脑屏幕的实时检测。我们可以不断地截屏，将图片逐张推理，再把推理结果返回窗口。

过程分为三步：

①截屏。此处可以使用MSS来进行实时的截屏；

②使用YOLO进行推理。这里我们调用run()方法即可，有的YOLO版本没有run()方法，需要自己改。

③展示推理结果。将推理结果展示到窗口上。

参考原文链接：https://blog.csdn.net/shopkeeper_/article/details/124578725

screenDetect.py的代码

import cv2
import numpy as np
from mss import mss

import detect

import os
import sys
from pathlib import Path

# 截图区域
bounding_box = {'top': 0, 'left': 0, 'width': 1280, 'height': 720}
# 项目根路径
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH

while True:
    # 截屏并保存到指定目录
    sct_img = mss().grab(bounding_box)
    scr_img = np.array(sct_img)
    cv2.imwrite(os.path.join(ROOT / 'data/screenShot', 'test.jpg'), scr_img)

    # 推理
    detect.run(weights=ROOT / 'yolov5s.pt',  # model.pt path(s) 权重
               data=ROOT / 'data/VOC.yaml',  # dataset.yaml path 模型
               source=ROOT / 'data/screenShot/test.jpg',  # 截图文件路径
               exist_ok=True,  # existing project/name ok, do not increment # 覆盖推理结果
               )

    # 读取推理结果并展示
    scr_img = cv2.imread(ROOT / 'runs/detect/exp/test.jpg', flags=1)
    scr_img = np.array(scr_img)
    scr_img = cv2.resize(scr_img, (800, 400))
    cv2.imshow("Screen Realtime", scr_img)
    # 按下q则终止程序
    if (cv2.waitKey(1) & 0xFF) == ord('q'):
        cv2.destroyAllWindows()
        break

因为是截一张图推理一次，所以整个推理的过程比较卡顿，还不能像使用摄像头那样进行流畅地检测。

可能的改进方法：像摄像头一样通过视频流来输入图片，一次处理多张图片。