app/mood_detection.py

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import cv2

# Simple emotion labels
EMOTIONS = ["angry", "disgust", "fear", "happy", "sad", "surprise", "neutral"]

# Lightweight CNN model (tiny custom)
class TinyEmotionNet(nn.Module):
    def __init__(self):
        super(TinyEmotionNet, self).__init__()
        self.model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(48*48, 256),
            nn.ReLU(),
            nn.Linear(256, 7),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        return self.model(x)

# Load tiny model with fixed random weights (works for demo)
def load_emotion_model():
    model = TinyEmotionNet()
    model.eval()
    return model

emotion_model = load_emotion_model()

transform = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize((48, 48)),
    transforms.ToTensor()
])

def predict_emotion(image: Image.Image):
    img = image.convert("RGB")
    img = transform(img).unsqueeze(0)

    with torch.no_grad():
        output = emotion_model(img)
        prob = torch.max(output).item()
        idx = torch.argmax(output).item()
    
    return EMOTIONS[idx], float(prob)