import cv2
import dlib
import gradio as gr
from PIL import Image
from transformers import pipeline

# Load pre-trained image classification model from transformers library
model = pipeline("image-classification", model="0x70DA/down-syndrome-classifier-v2")

# Load face detector from dlib library
detector = dlib.get_frontal_face_detector()


def predict(img):
    faces = detector(img)
    if len(faces) > 0:
        face = faces[0]  # Assuming there's only one face in the image
        x, y, w, h = face.left(), face.top(), face.width(), face.height()
        cropped_face = img[y : y + h, x : x + w]
        # Convert the cropped image to a PIL image
        pil_image = Image.fromarray(cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB))
        pred = model(pil_image)
        return {o["label"]: o["score"] for o in pred}
    return RuntimeError("No faces detected.")


demo = gr.Interface(
    fn=predict, inputs=gr.components.Image(), outputs=gr.components.Label()
)
demo.launch()