results = model.predict(source='./picShot/screenshot.png')

draw_box.png

完整的实时截屏检测并画框代码如下：

import pyautogui
import cv2
from ultralytics import YOLO

# Load a model
model = YOLO("yolov8n.pt")  # load a pretrained model (recommended for training)


def show_img():
   # 屏幕截图保存
   pyautogui.screenshot().save('./picShot/screenshot.png')

   results = model.predict(source='./picShot/screenshot.png')

   image = cv2.imread('./picShot/screenshot.png')

   image = draw_box(image, results)

   image = cv2.resize(image, (1920, 1080), fx=100, fy=100)
   # 展示屏幕截图图片
   cv2.imshow("show", image)

   cv2.waitKeyEx(1)


# 定义了一个颜色列表 COLORS ，其中包含了一些颜色的BGR值，用于在图像上绘制不同类别的目标框。
COLORS = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (255, 255, 0), (0, 255, 255), (255, 0, 255), (192, 192, 192), (128, 128, 128), (128, 0, 0),(128, 128, 0), (0, 128, 0)]


def draw_box(image, results):
   # 通过一个循环遍历每个边界框。
   for result in results:
       box = result.boxes
       # 匹配的类型数量小于1则不处理，表示未匹配到
       if len(box.cls.tolist()) < 1:
           continue
       # 获取置信度，粗略的使用第一个值
       # conf = box.conf.tolist()[0]
       # 获取标题映射result.names key在box.cls里面
       # {0: 'retry', 1: 'role', 2: 'start1', 3: 'start2', 4: 'start3', 5: 'start4', 6: 'start5'}
       for listItem in box.xyxy.tolist():
           # 获取当前边界框的左上角和右下角坐标，并将其转换为整数类型
           x0, y0, x1, y1 = int(listItem[0]), int(listItem[1]), int(listItem[2]), int(listItem[3])

           color = [int(c) for c in COLORS[int(box.cls.tolist()[0])]]
           # 使用 cv2.rectangle 函数在图像上绘制边界框，传入边界框的左上角坐标和右下角坐标，颜色值以及线宽（这里设定为3）。
           cv2.rectangle(image, (x0, y0), (x1, y1), color, 3)
           # 拿到标题
           text = result.names[int(box.cls.tolist()[0])]
           # cv2.putText 函数在图像上绘制标签文本，传入标签文本内容、文本位置、字体、字体大小、颜色值以及文本厚度
           cv2.putText(image, text, (max(0, x0), max(0, y0 - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
   return image


class RuntimeYolo:
   def main(self):
       while True:
           show_img()



if __name__ == "__main__":
   gameSupport = RuntimeYolo()
   gameSupport.main()