import streamlit as st
from PIL import Image
import torch
from src.model import EmotionCNN
import torchvision.transforms as transforms

st.title("🧠 Emotion Detector")

uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "png", "jpeg"])
if uploaded_file:
    image = Image.open(uploaded_file).convert("L").resize((48, 48))
    st.image(image, caption="Uploaded Image", use_column_width=True)

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    input_tensor = transform(image).unsqueeze(0)

    model = EmotionCNN()
    model.load_state_dict(torch.load("emotion_cnn.pth", map_location="cpu"))
    model.eval()

    with torch.no_grad():
        output = model(input_tensor)
        prediction = torch.argmax(output, dim=1).item()

    label_map = {0: "Happy", 1: "Sad", 2: "Neutral"}
    st.success(f"Predicted Emotion: **{label_map[prediction]}**")