35 lines
758 B
Python
35 lines
758 B
Python
from ultralytics import YOLO
|
|
import cv2
|
|
import pyautogui
|
|
import numpy as np
|
|
|
|
|
|
|
|
# 加载训练好的模型
|
|
# best.pt: 最佳模型，适用于生产
|
|
# last.pt: 最后一轮训练的模型，适用于继续训练
|
|
yolo = YOLO('runs/detect/train6/weights/best.pt')
|
|
|
|
# 指定屏幕范围
|
|
# x,y,width,height 全屏None
|
|
window = None
|
|
while True:
|
|
# 获取屏幕截图
|
|
screenshot = pyautogui.screenshot()
|
|
frame = np.array(screenshot)
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
|
|
|
# 使用YOLO检测
|
|
results = yolo(frame)
|
|
|
|
# 显示结果
|
|
annotated_frame = results[0].plot()
|
|
cv2.imshow("Screen Detection", annotated_frame)
|
|
|
|
# 按'q'退出
|
|
if cv2.waitKey(1) == ord('q'):
|
|
break
|
|
|
|
|
|
|
|
cv2.destroyAllWindows() |