import gradio as gr
import os
import cv2
import numpy as np
import imutils
from keras.preprocessing.image import img_to_array
from keras.models import load_model

# Load the pre-trained models and define parameters
detection_model_path = 'haarcascade_files/haarcascade_frontalface_default.xml'
emotion_model_path = 'model_2_aug_nocall_BEST/model_2_aug_nocall_entire_model.h5'
face_detection = cv2.CascadeClassifier(detection_model_path)
emotion_classifier = load_model(emotion_model_path, compile=False)
EMOTIONS = ['neutral', 'happiness', 'surprise', 'sadness', 'anger', 'disgust', 'fear', 'contempt', 'unknown']

# Function to predict emotions from a frame
def predict(frame_or_path):
    if isinstance(frame_or_path, np.ndarray):  # If input is a webcam frame
        frame = imutils.resize(frame_or_path, width=300)
    else:  # If input is a file path
        frame = cv2.imread(frame_or_path)
        if frame is None:
            return None, "Error: Unable to read image or video."
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_detection.detectMultiScale(gray, scaleFactor=1.1,
                                                minNeighbors=5, minSize=(30, 30),
                                                flags=cv2.CASCADE_SCALE_IMAGE)
        if len(faces) == 0:
            return frame, "No face detected."
        (fX, fY, fW, fH) = faces[0]
        roi = gray[fY:fY + fH, fX:fX + fW]
        roi = cv2.resize(roi, (48, 48))
        roi = roi.astype("float") / 255.0
        roi = img_to_array(roi)
        roi = np.expand_dims(roi, axis=0)
        preds = emotion_classifier.predict(roi)[0]
        label = EMOTIONS[preds.argmax()]
        cv2.putText(frame, label, (fX, fY - 10),
                    cv2.FONT_HERSHEY_DUPLEX, 1, (238, 164, 64), 1)
        cv2.rectangle(frame, (fX, fY), (fX + fW, fY + fH),
                      (238, 164, 64), 2)
        return frame, {emotion: float(prob) for emotion, prob in zip(EMOTIONS, preds)}

# Define input and output components for Gradio
image_input = gr.components.Image(sources=["webcam", "upload"], label="Your face")
image_output = gr.components.Image(label="Predicted Emotion")
label_output = gr.components.Label(num_top_classes=2, label="Top 2 Probabilities")

# Launch the Gradio interface
title = "Facial Emotion Recognition"
description = "How well can this model predict your emotions? Take a picture with your webcam, or upload an image, and it will guess if you are happy, sad, angry, disgusted, scared, surprised, or neutral."
thumbnail = "https://raw.githubusercontent.com/gradio-app/hub-emotion-recognition/master/thumbnail.png"

gr.Interface(fn=predict, inputs=image_input, outputs=[image_output, label_output],
            title=title, description=description, thumbnail=thumbnail).launch()