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

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

model = torch.hub.load('ultralytics/yolov5', 'custom', path = 'runs/train/exp/weights/best.pt')

cap = cv2.VideoCapture('test2.mp4')

while True:
    ret, frame = cap.read()
    if not ret:
        print('無法取得影像')
        break

    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    results = model(rgb_frame)

    output_frame = results.render()[0]  # 在影像上繪製標註

    bgr_frame = cv2.cvtColor(output_frame, cv2.COLOR_RGB2BGR)
    cv2.imshow('YOLOv5 Detection', bgr_frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()