# 修改截圖方式
# 優(yōu)化了窗口捕獲和推理的幀數(shù)（約30幀）
# 修改為按Esc結(jié)束（請(qǐng)務(wù)必使用Esc結(jié)束窗口運(yùn)行，否則極容易出現(xiàn)“無(wú)權(quán)限訪問(wèn)”的問(wèn)題，需要焦點(diǎn)在窗口上按Esc）
# 新增多線程調(diào)用優(yōu)化
# 新增names顯示，修改數(shù)字為names里面的名字(兼容中文)
# 新增窗口置頂和大小調(diào)整
# 新增名稱(chēng)/置信度，F(xiàn)PS顯示
# forecast.py
import os
import threading
import time
from queue import Queue

import CV2
import numpy as np
import pygetwindow as gw
import win32con
import win32gui
import yaml
from mss import mss

from ultralytics import YOLO


def load_names(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        data = yaml.safe_load(f)
    return data['names']


# 替換為你的data.yaml 里面保存有nc和names的那個(gè)，可以使用中文，推理出來(lái)的就是你names的名字，不是0,1數(shù)字了
names_path = "data.yaml"
names = load_names(names_path)

# 窗口名稱(chēng)和大小 ，推理時(shí)允許手動(dòng)拖拽窗口大小
window_name = "YOLOv8 Predict Test"
display_window_width = 768
display_window_height = 432
# 27 按Esc結(jié)束
exit_code = 27


# exit_code = ord("q")  按q結(jié)束

def capture_screen(img_queue):
    with mss() as sct:
        monitor = sct.monitors[0]

        while True:
            sct_img = sct.grab(monitor)
            img = np.array(sct_img)
            img = CV2.cvtColor(img, CV2.COLOR_BGRA2BGR)
            img_queue.put(img)


class YoloThread(threading.Thread):
    def __init__(self, model, img_queue, result_queue):
        threading.Thread.__init__(self)
        self.model = model
        self.img_queue = img_queue
        self.result_queue = result_queue

    def run(self):
        while True:
            img = self.img_queue.get()
            results = self.model.predict(source=img, conf=0.25, iou=0.75)
            self.result_queue.put((img, results))


def run(model, top_most=True):
    window_flag = CV2.WINDOW_NORMAL
    fps_update_interval = 0.5  # 每0.5秒更新一次
    frame_counter = 0
    last_fps_update_time = time.time()
    fps = 0

    img_queue = Queue()
    result_queue = Queue()

    capture_thread = threading.Thread(target=capture_screen, args=(img_queue,))
    capture_thread.daemon = True
    capture_thread.start()

    yolo_thread = YoloThread(model, img_queue, result_queue)
    yolo_thread.daemon = True
    yolo_thread.start()

    # 將這兩個(gè)函數(shù)放在 while True 循環(huán)之外
    CV2.namedWindow(window_name, window_flag)
    CV2.resizeWindow(window_name, display_window_width, display_window_height)

    while True:
        current_frame_time = time.time()
        img, results = result_queue.get()
        for result in results:
            if len(result.boxes.xyxy) > 0:
                boxes_conf = np.array(result.boxes.conf.tolist())
                boxes_xyxy = result.boxes.xyxy.tolist()
                boxes_cls = result.boxes.cls.tolist()

                for i, box_xyxy in enumerate(boxes_xyxy):
                    CV2.rectangle(img, (int(box_xyxy[0]), int(box_xyxy[1])),
                                  (int(box_xyxy[2]), int(box_xyxy[3])), (0, 0, 150), 2)
                    class_name = names[int(boxes_cls[i])]
                    confidence_text = f"{class_name}: {boxes_conf[i]:.2f}"
                    CV2.putText(img, confidence_text, (int(box_xyxy[0]), int(box_xyxy[1]) - 20),
                                CV2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 150), 2, CV2.LINE_AA)

        frame_counter += 1
        elapsed_time = current_frame_time - last_fps_update_time

        fps_text = f"FPS: {fps}"
        CV2.putText(img, fps_text, (display_window_width - 700, 40),
                    CV2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 150), 2, CV2.LINE_AA)

        if elapsed_time >= fps_update_interval:
            fps = int(frame_counter / elapsed_time)
            frame_counter = 0
            last_fps_update_time = current_frame_time

        CV2.imshow(window_name, img)

        if top_most:
            window = gw.getWindowsWithTitle(window_name)[0]
            window_handle = window._hWnd
            win32gui.SetWindowPos(window_handle, win32con.HWND_TOPMOST, 0, 0, 0, 0,
                                  win32con.SWP_NOMOVE | win32con.SWP_NOSIZE)

        if CV2.waitKey(1) == exit_code:
            CV2.destroyAllWindows()
            os._exit(0)


if __name__ == '__main__':
    # 是否默認(rèn)窗口置頂
    top_most = True
    # 你的best.pt模型(生成模型之后千萬(wàn)記得替換模型，我就因?yàn)橥浱鎿Q模型還以為我數(shù)據(jù)集有問(wèn)題hh)
    model = YOLO("../best.pt")

    run(model, top_most)

效果如下：