import gradio as gr
from fastai.learner import load_learner
from fastai.vision.all import *
import matplotlib.pyplot as plt

# Load the trained model
def label_func(fname):
    categories = ["Happy", "Sad", "Angry"]
    category = fname.parts[-2]  # Extract the category name from the path
    return category if category in categories else "unknown"

learn = load_learner('emotion_model.pkl')

# Define function to preprocess image
def preprocess_image(image):
    # Convert image to PILImage
    image_pil = PILImage.create(image)
    # Apply same transformations as done during training
    image_resized = Resize(224)(image_pil)
    return image_resized

# Define function to predict emotion and generate confidence bar chart
def predict_emotion(image):
    plt.clf()
    # Preprocess the image
    processed_image = preprocess_image(image)
    # Predict the emotion
    pred_class, pred_idx, outputs = learn.predict(processed_image)
    # Get class names
    class_names = learn.dls.vocab
    # Generate confidence values
    confidences = [outputs[idx].item() for idx in range(len(outputs))]
    # Generate bar chart
    plt.bar(class_names, confidences)
    plt.xlabel('Emotion')
    plt.ylabel('Confidence')
    plt.title('Confidence of Predicted Emotions')
    plt.xticks(rotation=45)
    plt.tight_layout()
    return plt

# Define Gradio interface
iface = gr.Interface(
    fn=predict_emotion,
    inputs="image",
    outputs="plot",
    title="Emotion Prediction",
    description="Upload an image and get the confidence of predicted emotions"
)

# Launch the interface
iface.launch()