# 優(yōu)化了up的腳本并添加了新功能
# 修改截圖方式
# 優(yōu)化了窗口捕獲和推理的幀數(shù)（約30幀）
# 修改為按Esc結(jié)束（請務(wù)必使用Esc結(jié)束窗口運行，否則極容易出現(xiàn)“無權(quán)限訪問”的問題，需要焦點在窗口上按Esc）
# 新增多線程調(diào)用優(yōu)化
# 新增names顯示，修改數(shù)字為names里面的名字(兼容中文)
# 新增窗口置頂和大小調(diào)整
# 新增名稱/置信度，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的那個，可以使用中文，推理出來的就是你names的名字，不是0,1數(shù)字了
names_path = "data.yaml"
names = load_names(names_path)

# 窗口名稱和大小 ，推理時允許手動拖拽窗口大小
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()

    # 將這兩個函數(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 = *******indowsWithTitle(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模型(生成模型之后千萬記得替換模型，我就因為忘記替換模型還以為我數(shù)據(jù)集有問題hh)
    model = YOLO("../best.pt")

    run(model, top_most)

效果如下：