Spaces:
Sleeping
Sleeping
File size: 3,217 Bytes
f32d9c8 dfc458b f32d9c8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """
Emotion detection using a PyTorch vision transformer (no TensorFlow needed).
Model: dima806/face_emotions_image_detection (ViT-tiny, ~28MB)
Labels: angry, disgust, fear, happy, sad, surprise, neutral
Downloaded from HuggingFace Hub on first use.
"""
from __future__ import annotations
import os
# Must be set before any transformers import to block TF (crashes without AVX on Rosetta)
os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("USE_JAX", "0")
os.environ.setdefault("USE_TORCH", "1")
os.environ.setdefault("TRANSFORMERS_NO_TF", "1")
os.environ.setdefault("TRANSFORMERS_NO_JAX", "1")
from typing import Dict, Tuple
import numpy as np
import torch
from PIL import Image
EMOTION_LABELS = ["angry", "disgust", "fear", "happy", "sad", "surprise", "neutral"]
DISPLAY_LABELS = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]
# Map model output labels → display labels (capitalized)
_LABEL_MAP = {l.lower(): l.capitalize() for l in DISPLAY_LABELS}
# Also map variants
_LABEL_MAP.update({"happiness": "Happy", "sadness": "Sad",
"anger": "Angry", "neutral": "Neutral",
"disgust": "Disgust", "fear": "Fear", "surprise": "Surprise"})
MODEL_ID = "dima806/face_emotions_image_detection"
class EmotionDetector:
"""
Lightweight ViT-based facial expression recognizer.
Input : face crop as RGB numpy array (any size).
Output: dict mapping emotion label → probability.
"""
def __init__(self) -> None:
self._pipe = None # lazy load
def _load(self):
import os
# Prevent transformers from loading TensorFlow (crashes on machines without AVX)
os.environ["USE_TF"] = "0"
os.environ["USE_JAX"] = "0"
os.environ["USE_TORCH"] = "1"
os.environ["TRANSFORMERS_NO_TF"] = "1"
os.environ["TRANSFORMERS_NO_JAX"] = "1"
from transformers import pipeline
print("[emotion] Loading emotion model (first run)…")
pipe = pipeline(
"image-classification",
model = MODEL_ID,
top_k = None,
device = -1, # CPU
)
print("[emotion] Loaded.")
return pipe
# ── public ────────────────────────────────────────────────────────────
def predict(self, face_rgb: np.ndarray) -> Dict[str, float]:
"""
Returns dict: { 'Happy': 0.87, 'Sad': 0.05, … }
"""
if self._pipe is None:
self._pipe = self._load()
pil = Image.fromarray(face_rgb).convert("RGB")
results = self._pipe(pil)
# Normalise labels → Title Case and fill any missing emotions
out: Dict[str, float] = {lbl: 0.0 for lbl in DISPLAY_LABELS}
for r in results:
label = _LABEL_MAP.get(r["label"].lower(), r["label"].capitalize())
out[label] = float(r["score"])
return out
def top_emotion(self, face_rgb: np.ndarray) -> "Tuple[str, float]":
probs = self.predict(face_rgb)
label = max(probs, key=probs.get)
return label, probs[label]
|