app/session.py

"""
File: session.py
Author: Dr. Gordon Wright
Description: In-memory session model for the six-emotion replication
             challenge. The session is a fixed-shape dict keyed by the
             six basic emotions (no neutral). Each slot is either None
             (not yet attempted) or a Capture. Sessions live in a
             gr.State and reset on page refresh.
License: MIT License
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional, Dict, List

import numpy as np
from PIL import Image


# Six basic emotions for the replication challenge. Order is the
# canonical Ekman order; the UI grid will render in this order. Neutral
# is excluded — it isn't part of the replication ask.
BASIC_EMOTIONS: List[str] = [
    "happy",
    "sad",
    "fear",
    "disgust",
    "anger",
    "surprise",
]


# Maps the lowercase activity label onto the classifier's DICT_EMO key.
EMOTION_TO_CLASSIFIER: Dict[str, str] = {
    "happy": "Happiness",
    "sad": "Sadness",
    "fear": "Fear",
    "disgust": "Disgust",
    "anger": "Anger",
    "surprise": "Surprise",
}


@dataclass
class Capture:
    """One image-and-prediction event in a session.

    `intended` records which of the six emotions the student was
    *trying* to produce. The classifier's own answer lives in
    `emotion_probs`; the two often disagree and that disagreement is
    the lesson.
    """

    intended: str
    face: np.ndarray
    emotion_probs: Dict[str, float]
    heatmap: Optional[np.ndarray] = None
    blendshapes: Optional[Dict[str, float]] = None
    landmarks: Optional[list] = None
    bbox: Optional[tuple] = None
    image_size: Optional[tuple] = None

    def top_emotion(self) -> tuple:
        items = sorted(self.emotion_probs.items(), key=lambda x: -x[1])
        return items[0]

    def thumbnail(self) -> Image.Image:
        return Image.fromarray(self.face)

    def classifier_agrees(self) -> bool:
        top, _ = self.top_emotion()
        return top == EMOTION_TO_CLASSIFIER.get(self.intended)


def empty_session() -> Dict[str, Optional[Capture]]:
    return {emo: None for emo in BASIC_EMOTIONS}


def session_status(state: Dict[str, Optional[Capture]]) -> str:
    """One-line summary of progress across the six slots."""
    filled = sum(1 for c in state.values() if c is not None)
    if filled == 0:
        return f"0 / 6 emotions captured. Pick one from the list and submit your first attempt."
    if filled < 6:
        return f"{filled} / 6 emotions captured. Keep going."
    agree = sum(1 for c in state.values() if c is not None and c.classifier_agrees())
    return f"All 6 captured — classifier agreed on {agree} / 6."