from ultralytics import YOLO
import cv2

if __name__ == "__main__":

    corner_model_path = "../assets/models/corner.pt"
    pieces_model_path = "../assets/models/unified-nano-refined.pt"

    print("Initializing model...")
    corner_model = YOLO(corner_model_path)
    pieces_model = YOLO(pieces_model_path)

    print("Initializing camera...")
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    cap.set(cv2.CAP_PROP_FPS, 30)

    print("Initialized")
    if not cap.isOpened():
        print("Error: Could not open camera")
        exit()

    cv2.namedWindow("Predictions", cv2.WINDOW_NORMAL)

    while True:
        ret, frame = cap.read()
        if not ret:
            print("Error: Failed to grab frame")
            break

        # Optional: resize frame to improve YOLO performance
        # frame = cv2.resize(frame, (416, 416))

        corner_result = corner_model.predict(source=frame, conf=0.6)
        pieces_result = pieces_model.predict(source=frame, conf=0.6)

        corner_annotated_frame = corner_result[0].plot()
        pieces_annotated_frame = pieces_result[0].plot(img=corner_annotated_frame)

        cv2.imshow("Predictions", pieces_annotated_frame)
        cv2.resizeWindow("Predictions", 640, 640)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()