scratch/smoke_session.py

"""Smoke test the R1 session model.

Builds a Capture from a real classifier pass on a sample face, prints
the resolved caption, the top blendshape activations, and the gallery
item that would be handed to gr.Gallery. Also exercises remove-last
and clear semantics on a small session.

Run from repo root: `cd ~/Code/research/totes-emosh && uv run python scratch/smoke_session.py`
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

from PIL import Image

ROOT = Path(__file__).parent.parent
os.chdir(ROOT)
sys.path.insert(0, str(ROOT))


def main() -> None:
    from app.app_utils import preprocess_image_and_predict
    from app.session import (
        Capture,
        LABEL_PRESETS,
        gallery_items,
        resolve_label,
        session_status,
    )

    print(f"LABEL_PRESETS: {len(LABEL_PRESETS)} entries")
    print(f"  first 3: {LABEL_PRESETS[:3]}")
    print(f"  last:    {LABEL_PRESETS[-1]}")

    print(f"\nresolve_label('warm-up:happy', '') = {resolve_label('warm-up:happy', '')!r}")
    print(f"resolve_label('(custom)', 'AU1+AU6 attempt') = {resolve_label('(custom)', 'AU1+AU6 attempt')!r}")
    print(f"resolve_label('(custom)', '') = {resolve_label('(custom)', '')!r}")

    print("\nRunning pipeline on Happy.png and building a Capture...")
    img = Image.open(ROOT / "images" / "Happy.png").convert("RGB")
    face, heatmap, probs, blendshapes = preprocess_image_and_predict(img)

    cap = Capture(
        label="warm-up:happy",
        face=face,
        heatmap=heatmap,
        emotion_probs=probs,
        blendshapes=blendshapes,
    )

    print(f"  caption: {cap.caption()}")
    print(f"  top emotion: {cap.top_emotion()}")
    print(f"  thumbnail size: {cap.thumbnail().size}")

    print("\n  top 5 blendshapes:")
    if blendshapes:
        for name, score in sorted(blendshapes.items(), key=lambda x: -x[1])[:5]:
            print(f"    {name:<24} {score:.3f}")

    print("\nExercising session operations:")
    session = [cap]
    print(f"  after add: {session_status(session)}")
    items = gallery_items(session)
    print(f"  gallery_items returns {len(items)} item(s); first caption: {items[0][1]!r}")

    session = session + [cap, cap]
    print(f"  after two more adds: {session_status(session)}")

    session = session[:-1]
    print(f"  after remove-last: {session_status(session)}")

    session = []
    print(f"  after clear: {session_status(session)}")


if __name__ == "__main__":
    main()