import cv2
from picamera2 import Picamera2
from ultralytics import YOLO

ESC = 27

# 1. 載入 YOLOv5 模型（預設權重已在 COCO 權重上訓練）
model = YOLO('yolov5n.pt')

# 2. 初始化 Picamera2，設定輸出為 640×480 RGB888
picam2 = Picamera2()
config = picam2.create_preview_configuration(main={'size': (640, 480), 'format': 'RGB888'})
picam2.configure(config)
picam2.start()

# 下面這行告訴 YOLO 只回傳 COCO 類別中 'tv' 的偵測結果 (class index = 62)
# 參考：https://docs.ultralytics.com/usage/cfg/ （參見 'classes' 選項） :contentReference[oaicite:2]{index=2}
target_classes = [62]  # COCO class 'tv' (0-based index) :contentReference[oaicite:3]{index=3}

while True:
    # 3. 擷取每一幀，並翻轉（視鏡頭實際安裝方式調整）
    frame = picam2.capture_array()
    frame = cv2.flip(frame, -1)
    if frame is None:
        print("❌ 擷取影像失敗")
        break

    # 4. 推論：加上 classes=[62] 只回傳 'tv' 類別的 bounding boxes
    results = model(frame, classes=target_classes)

    # 5. 將偵測結果繪製回影像 (plot() 只包含 'tv' 框)
    annotated = results[0].plot()

    # 6. 顯示畫面
    cv2.imshow('Detect TV Only', annotated)

    if cv2.waitKey(1) == ESC:
        break

# 7. 結束釋放
picam2.stop()
cv2.destroyAllWindows()