import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, models
from PIL import Image
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np

# Load the pre-trained model (ensure to use the saved model checkpoint)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Model: EfficientNet-B0 with dropout added to reduce overfitting
model = models.efficientnet_b0(pretrained=True)
model.classifier = nn.Sequential(
    nn.Dropout(0.4),
    nn.Linear(model.classifier[1].in_features, 7)  # num_classes = 7 (angry, disgust, fear, happy, neutral, sad, surprise)
)
model.load_state_dict(torch.load("best_mood_classifier.pth", map_location=torch.device('cpu')))
model = model.to(device)
model.eval()

# Define the image transformations for the uploaded image
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])

# Class names (same order as in your dataset)
class_names = ['angry', 'disgust', 'fear', 'happy', 'neutral', 'sad', 'surprise']

# Function to predict the mood from the uploaded image
def predict_mood(image):
    image = Image.fromarray(image)
    image = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(image)
        _, preds = torch.max(outputs, 1)
        predicted_class = class_names[preds.item()]

    return predicted_class

# Gradio interface
iface = gr.Interface(
    fn=predict_mood,
    inputs=gr.Image(type="numpy"),
    outputs="text",
    live=True,
    title="Mood Classifier",
    description="Upload an image of a face and the model will predict the mood."
)

# Launch the app
iface.launch()