import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image # Pillow for image processing
import os

# Hugging Face Space Link: https://huggingface.co/spaces/Landhoff/ReassessmentAAI2

# Load the trained Keras model
model = tf.keras.models.load_model('facial_emotion_model.h5')

# Define the class labels explicitly, as train_generator is not available in deployment
# These labels should match the order your model was trained on.
class_labels = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16', '17', '18']

# Define target image size used during training
TARGET_SIZE = (128, 128)

def preprocess_image(image) -> np.ndarray:
    """Preprocesses the input image for model prediction."""
    # Resize image to target size
    image = image.resize(TARGET_SIZE)
    # Convert to numpy array
    image_array = np.array(image)
    # Normalize pixel values to [0, 1]
    image_array = image_array / 255.0
    # Expand dimensions to create a batch (1, height, width, channels)
    image_array = np.expand_dims(image_array, axis=0)
    return image_array

def predict_emotion(image) -> dict:
    """Predicts the emotion from a facial image and returns probabilities."""
    if image is None:
        return {label: 0.0 for label in class_labels}

    # Preprocess the image
    processed_image = preprocess_image(image)

    # Make prediction
    predictions = model.predict(processed_image)[0]

    # Create a dictionary of emotion labels and their probabilities
    results = {class_labels[i]: float(predictions[i]) for i in range(len(class_labels))}
    return results

# Create the Gradio interface
iface = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Image(type="pil", label="Upload Face Image"),
    outputs=gr.Label(num_top_classes=len(class_labels)),
    title="Facial Emotion Recognition",
    description="Upload an image of a face to get emotion predictions."
)

# Launch the interface
iface.launch()