results = model.predict(source='./picShot/screenshot.png')
draw_box.png
完整的實(shí)時截屏檢測并畫框代碼如下:
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值显晶,用于在圖像上繪制不同類別的目標(biāo)框题画。
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):
# 通過一個循環(huán)遍歷每個邊界框徙歼。
for result in results:
box = result.boxes
# 匹配的類型數(shù)量小于1則不處理皆疹,表示未匹配到
if len(box.cls.tolist()) < 1:
continue
# 獲取置信度扎狱,粗略的使用第一個值
# conf = box.conf.tolist()[0]
# 獲取標(biāo)題映射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():
# 獲取當(dāng)前邊界框的左上角和右下角坐標(biāo),并將其轉(zhuǎn)換為整數(shù)類型
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 函數(shù)在圖像上繪制邊界框伟叛,傳入邊界框的左上角坐標(biāo)和右下角坐標(biāo)私痹,顏色值以及線寬(這里設(shè)定為3)。
cv2.rectangle(image, (x0, y0), (x1, y1), color, 3)
# 拿到標(biāo)題
text = result.names[int(box.cls.tolist()[0])]
# cv2.putText 函數(shù)在圖像上繪制標(biāo)簽文本统刮,傳入標(biāo)簽文本內(nèi)容紊遵、文本位置、字體侥蒙、字體大小暗膜、顏色值以及文本厚度
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()