import cv2
from picamera2 import Picamera2
from ultralytics import YOLO

ESC = 27

# 載入 YOLO
model = YOLO('yolov5n.pt')

# 初始化 Picamera2
picam2 = Picamera2()
config = picam2.create_preview_configuration(main={'size': (640, 480), 'format': 'RGB888'})
picam2.configure(config)
picam2.start()

while True:
    frame = picam2.capture_array()
    frame = cv2.flip(frame, -1)
    if frame is None:
        print(f'幁抓取失敗')
        break

    # 做 YOLO 推論
    results = model(frame)
    annotated = results[0].plot()

    # 顯示
    cv2.imshow('preview', annotated)

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

picam2.stop()
cv2.destroyAllWindows()