import cv2
import dlib
import numpy as np
from PIL import Image

from predict_api import ImagePredictor

# Initialize camera and face detector
cap = cv2.VideoCapture(0)
detector = dlib.get_frontal_face_detector()

# Initialize ImagePredictor
predictor = ImagePredictor(model_path="best_model_'0.8998410174880763'.pth", class_indices_path="./class_indices.json")

while True:
    # Capture frame-by-frame
    ret, frame = cap.read()

    # Convert the image from BGR color (which OpenCV uses) to RGB color
    rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # Perform face detection
    faces = detector(rgb_image)

    # Loop through each face in this frame
    for rect in faces:
        # Get the bounding box coordinates
        x1, y1, x2, y2 = rect.left(), rect.top(), rect.right(), rect.bottom()

        # Crop the face from the frame
        face_image = rgb_image[y1:y2, x1:x2]

        # Use ImagePredictor to predict the class of this face
        result = predictor.predict(face_image)

        # Draw a rectangle around the face
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

        # Display the class name and score
        cv2.putText(frame, f"{result['result'][0]['name']}: {round(result['result'][0]['score'],4)}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)

    # Display the resulting frame
    cv2.imshow('Video', frame)

    # Exit loop if 'q' is pressed
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# When everything is done, release the capture
cap.release()
cv2.destroyAllWindows()