import gradio as gr
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.optimizers import Adam
from PIL import Image
import numpy as np

# Define the list of emotion class names (from previous K-Fold execution)
sorted_unique_class_names = [
    '0', '1', '10', '11', '12', '13', '14', '15', '16', '17', '18',
    '2', '3', '4', '5', '6', '7', '8', '9'
]

# --- Load the trained model ---
model_path = 'facial_emotion_model.keras'
model = load_model(model_path)
# Compile the model after loading to avoid the 'No training configuration found' warning
model.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])

# --- Preprocessing function ---
def preprocess_image(image: Image.Image):
    image = image.resize((128, 128))
    image_array = np.array(image)
    image_array = np.expand_dims(image_array, axis=0)
    image_array = image_array / 255.0
    return image_array

# --- Prediction function ---
def predict_emotion(image: Image.Image):
    if image is None:
        return "No image provided.", {},

    processed_image = preprocess_image(image)
    predictions = model.predict(processed_image)
    predicted_class_idx = np.argmax(predictions, axis=1)[0]
    predicted_emotion = sorted_unique_class_names[predicted_class_idx]
    
    confidence_scores = {sorted_unique_class_names[i]: float(predictions[0][i]) for i in range(len(sorted_unique_class_names))}
    
    return predicted_emotion, confidence_scores

# --- Gradio Interface ---
iface = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Image(type="pil", label="Upload an image of a face"),
    outputs=[
        gr.Textbox(label="Predicted Emotion"),
        gr.Label(num_top_classes=len(sorted_unique_class_names), label="Confidence Scores")
    ],
    title="Facial Emotion Recognition",
    description="Upload a facial image to get the predicted emotion. The model was trained on a custom dataset with 19 emotion classes."
)

iface.launch(debug=True)