# yolov5屏幕實時檢測

import mss
import CV2
import os
import threading
import time
import torch
import numpy as np
from win32gui import FindWindow, GetWindowRect
# # 通過 torch.hub.load 函數從指定路徑加載Yolov5模型，
# 使用的是自定義模型（'custom'），模型文件為'yolov5s.pt'，設備為GPU設備編號為0，源為本地。
# 加載完成后，將模型賦值給變量 yolov5 。
yolov5 = torch.hub.load('D:\\pythonProject\\202301\\yolov5master\\yolov5-master', 'custom', path='yolov5s.pt', device='0', source='local')
# 設置了 yolov5 的置信度閾值為0.3（ yolov5.conf = 0.3 ）和IoU閾值為0.4（ yolov5.iou = 0.4 ），用于篩選檢測結果。
yolov5.conf = 0.3
yolov5.iou = 0.4
# 定義了一個顏色列表 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)]

# (128, 0, 128), (0, 128, 128), (0, 0, 128)
# 定義了一個標簽列表 LABELS ，這里的類別還沒有修改成對應yolov5s.pt對應的類別
LABELS = ['tree', 'stone', 'herb', 'wolf', 'zombie', 'dear', 'bear', 'helicopter', 'berry', 'mushroom', 'cole']
# 創(chuàng)建了一個大小為(1280, 720, 3)的空圖像 img_src ，用于顯示檢測結果。
img_src = np.zeros((1280, 720, 3), np.uint8)

# 總的來說，這段代碼的作用是加載Yolov5模型，并設置模型的一些參數。同時定義了顏色列表和標簽列表，以及一個空的圖像用于顯示檢測結果。


def getScreenshot():
? ?# 這段代碼用于獲取屏幕截圖。首先，通過注釋掉的代碼可以看出，它嘗試根據窗口標題找到窗口的句柄，
? ?# 然后獲取窗口的位置信息（左上角坐標和右下角坐標）。
? ?# 但是這部分代碼被注釋掉了，所以直接給定了一個固定的窗口位置信息（左上角坐標為(0, 0)，右下角坐標為(961, 1035)）。
? ?# id = FindWindow(None, "Windows.UI.Core.CoreWindow")
? ?# x0, y0, x1, y1 = GetWindowRect(id)
? ?x0, y0, x1, y1 = 0, 0, 961, 1035
? ?# ?定義了兩個變量 mtop 和 mbot ，分別表示從截圖中去除的頂部和底部的像素行數。
? ?mtop, mbot = 30, 50
? ?# 創(chuàng)建了一個字典 monitor ，包含了屏幕截圖的位置和大小信息，其中左上角坐標為 x0 和 y0 ，寬度為 x1 - x0 ，高度為 y1 - y0 。
? ?monitor = {"left": x0, "top": y0, "width": x1 - x0, "height": y1 - y0}
? ?# 使用 mss.mss().grab(monitor) 函數獲取屏幕截圖，并將結果存儲在變量 img_src 中。
? ?# mss 是一個用于屏幕截圖的庫， grab 函數用于捕獲指定位置和大小的屏幕區(qū)域。
? ?img_src = np.array(mss.mss().grab(monitor))
? ?# 通過 time.sleep(0.1) 函數等待一段時間，以確保截圖操作完成。
? ?time.sleep(0.1)
? ?# 對截圖進行處理。將 img_src 的通道數限制為3，即去除可能存在的第四個通道（alpha通道）。
? ?img_src = img_src[:, :, :3]
? ?# 根據設定的 mtop 和 mbot 值，從截圖中去除對應的頂部和底部像素行
? ?img_src = img_src[mtop:-mbot]
? ?# 返回處理后的截圖 img_src 以及窗口位置信息的列表。
? ?return img_src, [x0, y0, x1, y1, mtop, mbot]


def getMonitor():
? ?global img_src
? ?while True:
? ? ? ?# 通過 getScreenshot() 函數獲取屏幕截圖，并將結果賦值給 img_src 變量。
? ? ? ?# 函數返回的 ? ?第二個值 ? 用下劃線，表示不需要使用該值。
? ? ? ?img_src, _ = getScreenshot()

def yolov5Detect():
? ?# 通過 CV2.namedWindow 函數創(chuàng)建一個名為"Window Name"的窗口，
? ?# 并指定窗口的屬性為 CV2.WINDOW_NORMAL ，即可調整窗口的大小。
? ?CV2.namedWindow("Window Name", CV2.WINDOW_NORMAL)
? ?# 使用 CV2.resizeWindow 函數將窗口大小設置為960x540像素
? ?CV2.resizeWindow("Window Name", 960, 540)
? ?# 使用 CV2.moveWindow 函數將窗口移動到屏幕上的指定位置（1560, 0）
? ?CV2.moveWindow("Window Name", 1560, 0)
? ?global img_src
? ?while True:
? ? ? ?# 通過 while True 創(chuàng)建一個無限循環(huán)。
? ? ? ?# 在每次循環(huán)中，首先將全局變量 img_src 的值復制給變量 img ，
? ? ? ?# 以確保獲取到最新的屏幕截圖。
? ? ? ?img = img_src.copy()
? ? ? ?# 調用 getDetection 函數對 img 進行目標檢測，
? ? ? ?# 返回檢測到的邊界框信息存儲在變量 bboxes 中。
? ? ? ?bboxes = getDetection(img)
? ? ? ?# 調用 drawBBox 函數在 img 上繪制檢測到的邊界框
? ? ? ?img = drawBBox(img, bboxes)
? ? ? ?# 通過 CV2.imshow 函數在窗口中顯示繪制好邊界框的圖像
? ? ? ?CV2.imshow("Window Name", img)
? ? ? ?# 如果用戶按下鍵盤上的"q"鍵，通過 CV2.waitKey 函數檢測到按鍵事件，
? ? ? ?if CV2.waitKey(1) & 0xFF == ord("q"):
? ? ? ? ? ?# 就會銷毀窗口并退出循環(huán)，否則繼續(xù)下一次循環(huán)。
? ? ? ? ? ?CV2.destroyAllWindows()
? ? ? ? ? ?break


def getLargestBox(bboxes, type):
? ?# 定義了一個變量 largest 并初始化為-1，用于記錄當前最大的面積值。
? ?largest = -1
? ?# 定義了一個空的NumPy數組 bbox_largest ，用于存儲最大的邊界框。
? ?bbox_largest = np.array([])
? ?for bbox in bboxes:
? ? ? ?# 通過一個循環(huán)遍歷每個邊界框。對于每個邊界框，首先判斷其對應的類別是否在給定的 type 類型列表中。
? ? ? ?# 這里通過 LABELS[int(bbox[5])] 來獲取邊界框的類別標簽，并判斷其是否在 type 列表中。
? ? ? ?# 如果在，則繼續(xù)執(zhí)行下面的操作；如果不在，則跳過該邊界框。
? ? ? ?if LABELS[int(bbox[5])] in type:
? ? ? ? ? ?# 獲取當前邊界框的左上角和右下角坐標。
? ? ? ? ? ?x0, y0, x1, y1 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
? ? ? ? ? ?# 計算其面積，使用的是邊界框的寬度乘以高度。
? ? ? ? ? ?area = (x1 - x0) * (y1 - y0)
? ? ? ? ? ?# 將當前邊界框的面積與 largest 進行比較。
? ? ? ? ? ?# 如果當前邊界框的面積大于 largest ，則更新 largest 為當前面積值，
? ? ? ? ? ?# 并將當前邊界框賦值給 bbox_largest 。
? ? ? ? ? ?if area > largest:
? ? ? ? ? ? ? ?largest = area
? ? ? ? ? ? ? ?bbox_largest = bbox
? ?return bbox_largest


def drawBBox(image, bboxes):
? ?# 通過一個循環(huán)遍歷每個邊界框。
? ?for bbox in bboxes:
? ? ? ?# 對于每個邊界框，首先獲取其置信度 conf 和類別ID classID 。
? ? ? ?conf = bbox[4]
? ? ? ?classID = int(bbox[5])
? ? ? ?# 如果置信度大于 yolov5.conf （即設定的置信度閾值），則執(zhí)行下面的操作。
? ? ? ?if conf > yolov5.conf:
? ? ? ? ? ?# 獲取當前邊界框的左上角和右下角坐標，并將其轉換為整數類型
? ? ? ? ? ?x0, y0, x1, y1 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
? ? ? ? ? ?# 這里根據需求設置條件，演示只取10個顏色
? ? ? ? ? ?if classID >= 10:
? ? ? ? ? ? ? ?classID = 10
? ? ? ? ? ?color = [int(c) for c in COLORS[classID]]
? ? ? ? ? ?# 使用 CV2.rectangle 函數在圖像上繪制邊界框，傳入邊界框的左上角坐標和右下角坐標，顏色值以及線寬（這里設定為3）。
? ? ? ? ? ?CV2.rectangle(image, (x0, y0), (x1, y1), color, 3)
? ? ? ? ? ?text = "{}: {:.2f}".format(LABELS[classID], conf)
? ? ? ? ? ?# CV2.putText 函數在圖像上繪制標簽文本，傳入標簽文本內容、文本位置、字體、字體大小、顏色值以及文本厚度
? ? ? ? ? ?CV2.putText(image, text, (max(0, x0), max(0, y0 - 5)), CV2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
? ? ? ? ? ?# 打印出邊界框的坐標和標簽文本
? ? ? ? ? ?print([x0, y0, x1, y1], text)
? ?return image


def getDetection(img):
? ?# 使用yolov5模型對圖像進行目標檢測
? ?# 將圖像轉換為RGB格式，并調整大小為1280
? ?bboxes = np.array(yolov5(img[:, :, ::-1], size=1280).xyxy[0].cpu())
? ?return bboxes


if __name__ == '__main__':
? ?# 創(chuàng)建了兩個線程，分別執(zhí)行 getMonitor 和 yolov5Detect 函數。
? ?t1 = threading.Thread(target=getMonitor, args=())
? ?t1.start()
? ?t2 = threading.Thread(target=yolov5Detect, args=())
? ?t2.start()