import torch
from torch.serialization import add_safe_globals
from models.yolo import Model
import cv2

add_safe_globals([Model])    # 允許 PyTorch 安全地載入 YOLOv5 模型類別

# 載入自行訓練的 YOLOv5 模型
model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/exp/weights/best.pt')
model.conf = 0.5  # 設定信心門檻值

# 初始化攝影機
cap = cv2.VideoCapture(0)    # 0 表示第一個 USB 攝影機
if not cap.isOpened():
    print("無法打開攝影機")
    exit()

while True:
    ret, frame = cap.read()    # 讀取影格
    if not ret:
        print("無法獲取影格")
        break

    # 使用模型進行推論
    results = model(frame)

    # 繪製結果
    annotated_frame = results.render()[0]    # 獲取渲染後的影格
    cv2.imshow('YOLOv5 Detection', annotated_frame)    # 顯示影格

    # 按下 'q' 鍵退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 釋放資源
cap.release()
cv2.destroyAllWindows()