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()
