import cv2
import dlib
import streamlit as st
from PIL import Image
from transformers import pipeline
import numpy as np

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

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

# Define the prediction function
def predict(image):
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)  # Convert PIL Image to NumPy array
    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 {"No Face Detected": 0.0}

# Create the Streamlit app interface
st.title("Down Syndrome Classifier")

uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

if uploaded_image is not None:
    image = Image.open(uploaded_image)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    st.write("Classifying...")

    result = predict(image)

    st.write("Classification Results:")
    for label, score in result.items():
        st.write(f"{label}: {score:.4f}")