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

过程分为三步:

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

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

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

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

screenDetect.py的代码

12345678910111213141516171819202122232425262728293031323334353637383940import cv2import numpy as npfrom mss import mssimport detectimport osimport sysfrom pathlib import Path# 截图区域bounding_box = {'top': 0, 'left': 0, 'width': 1280, 'height': 720}# 项目根路径FILE = Path(__file__).resolve()ROOT = FILE.parents[0] # YOLOv5 root directoryif str(ROOT) not in sys.path: sys.path.append(str(ROOT)) # add ROOT to PATHwhile 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

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

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