Update app.py

app.py

@@ -1,4 +1,4 @@
-import os, io, math, json, random, numpy as np
+import os, numpy as np
 from typing import List, Tuple, Dict, Any
 from PIL import Image
 
@@ -8,21 +8,23 @@ import torch.nn.functional as F
 import gradio as gr
 from datasets import load_dataset
 from sklearn.neighbors import NearestNeighbors
-from transformers import pipeline
 
 # =============== CONFIG ===============
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Embeddings (your chosen model)
+# Embeddings backbone
 OPENCLIP_BACKBONE = "ViT-H-14"
 OPENCLIP_PRETRAIN = "laion2B-s32B-b79K"  # laion/CLIP-ViT-H-14-laion2B-s32B-b79K
-INDEX_SIZE = int(os.getenv("INDEX_SIZE", 400))  # כמה תמונות מהדאטהסט לאינדוקס
-TOPK_NEAREST = 5
 
-# Emotion model (your chosen model)
-EMO_MODEL = "prithivMLmods/Facial-Emotion-Detection-SigLIP2"
+# Dataset (THIS IS YOUR "MODEL" SOURCE NOW)
+DATASET_NAME = "tukey/human_face_emotions_roboflow"
+DATASET_SPLIT = "train"
+
+INDEX_SIZE = int(os.getenv("INDEX_SIZE", 400))     # כמה דוגמאות מהדאטהסט לאינדוקס
+TOPK_NEAREST = 5                                   # להצגה בגלריה
+KNN_K_FOR_CLASS = 25                               # לשקלול רגשות
 
-# Optional image-variation generator (your chosen model)
+# Optional SD variations
 USE_SD_VARIATIONS = True
 SD_MODEL = "lambdalabs/sd-image-variations-diffusers"
 # =====================================
@@ -48,110 +50,126 @@ def embed_image(img: Image.Image) -> np.ndarray:
     feats = F.normalize(feats, dim=-1).squeeze(0).detach().cpu().numpy().astype(np.float32)
     return feats  # shape [D]
 
-# ---------- Dataset + index ----------
-DATASET_NAME = "tukey/human_face_emotions_roboflow"
-DATASET_SPLIT = "train"
+# ---------- Labels & stress mapping ----------
+EMO_MAP = {
+    "anger": "anger", "angry": "anger",
+    "disgust": "disgust",
+    "fear": "fear",
+    "happy": "happy", "happiness": "happy",
+    "neutral": "neutral", "calm": "neutral",
+    "sad": "sad", "sadness": "sad",
+    "surprise": "surprise",
+    "contempt": "contempt",
+}
+ALLOWED = set(EMO_MAP.values())  # whitelist קשיח
+
+STRESS_WEIGHTS = {
+    "anger": 0.95, "fear": 0.90, "disgust": 0.70, "sad": 0.80,
+    "surprise": 0.55, "neutral": 0.30, "contempt": 0.65, "happy": 0.10,
+}
+def _bucket(p: float) -> str:
+    return "Low" if p < 33 else ("Medium" if p < 66 else "High")
 
-def _load_images_for_index(n: int) -> List[Image.Image]:
+# ---------- Load dataset & build index ----------
+def _extract_label(rec: Dict[str, Any]) -> str:
+    # התאמה לשדות אפשריים בדאטהסט
+    if "label" in rec and rec["label"]:
+        raw = rec["label"]
+        if isinstance(raw, (list, tuple)): raw = raw[0]
+        return str(raw).strip().lower()
+    if "labels" in rec and rec["labels"]:
+        raw = rec["labels"][0]
+        return str(raw).strip().lower()
+    if "qa" in rec and rec["qa"] and isinstance(rec["qa"], list):
+        qa0 = rec["qa"][0]
+        if qa0 and "answer" in qa0:
+            return str(qa0["answer"]).strip().lower()
+    return ""
+
+def _map_allowed(lbl: str) -> str:
+    # ממפה לשם סטנדרטי, ומסנן החוצה לא מוכרות
+    mapped = EMO_MAP.get(lbl, lbl)
+    return mapped if mapped in ALLOWED else ""  # "" => drop
+
+def _load_images_labels_for_index(n: int) -> Tuple[List[Image.Image], List[str]]:
     ds = load_dataset(DATASET_NAME, split=DATASET_SPLIT)
+    imgs, labels = [], []
     n = min(n, len(ds))
-    imgs = []
     for i in range(n):
-        # column is usually "image"
-        im = ds[i].get("image")
-        if isinstance(im, Image.Image):
-            imgs.append(im.copy())
-    return imgs
+        rec = ds[i]
+        im = rec.get("image")
+        if not isinstance(im, Image.Image):
+            continue
+        raw_lbl = _extract_label(rec)
+        mapped = _map_allowed(raw_lbl)
+        if not mapped:
+            continue  # זורק תוויות לא מותרות/ריקות
+        imgs.append(im.copy())
+        labels.append(mapped)
+    return imgs, labels
 
 def build_index(imgs: List[Image.Image]) -> Tuple[NearestNeighbors, np.ndarray]:
-    vecs = []
-    for im in imgs:
-        vecs.append(embed_image(im))
+    vecs = [embed_image(im) for im in imgs]
     X = np.stack(vecs, axis=0)
-    nn = NearestNeighbors(metric="cosine", n_neighbors=min(TOPK_NEAREST, len(imgs)))
+    nn = NearestNeighbors(metric="cosine", n_neighbors=min(max(TOPK_NEAREST, KNN_K_FOR_CLASS), len(imgs)))
     nn.fit(X)
     return nn, X
 
 print("Loading dataset & building index (first time only)...")
-DATASET_IMAGES: List[Image.Image] = _load_images_for_index(INDEX_SIZE)
+DATASET_IMAGES, DATASET_LABELS = _load_images_labels_for_index(INDEX_SIZE)
+if len(DATASET_IMAGES) == 0:
+    raise RuntimeError("No images with allowed labels were loaded from the dataset.")
 NN_MODEL, EMB_MATRIX = build_index(DATASET_IMAGES)
-print(f"Index ready with {len(DATASET_IMAGES)} images.")
+print(f"Index ready with {len(DATASET_IMAGES)} images (labels={sorted(set(DATASET_LABELS))}).")
 
+# ---------- Nearest & KNN-based classification ----------
 def nearest5(pil_img: Image.Image) -> List[Tuple[Image.Image, str]]:
     q = embed_image(pil_img).reshape(1, -1)
-    dists, idxs = NN_MODEL.kneighbors(q, n_neighbors=min(5, len(DATASET_IMAGES)))
-    # cosine distance -> similarity = 1 - dist
+    n = min(5, len(DATASET_IMAGES))
+    dists, idxs = NN_MODEL.kneighbors(q, n_neighbors=n)
     out = []
     for rank, (dist, idx) in enumerate(zip(dists[0], idxs[0]), start=1):
-        sim = 1.0 - float(dist)
+        sim = max(0.0, 1.0 - float(dist))  # cosine distance -> similarity
         im = DATASET_IMAGES[int(idx)]
         caption = f"#{rank}  sim={sim:.3f}  idx={int(idx)}"
         out.append((im, caption))
-    return out  # list of (PIL, caption)
-
-# ---------- Emotion & Stress ----------
-EMO_MAP = {
-    "anger": "anger", "angry": "anger",
-    "disgust": "disgust",
-    "fear": "fear",
-    "happy": "happy", "happiness": "happy",
-    "neutral": "neutral", "calm": "neutral",
-    "sad": "sad", "sadness": "sad",
-    "surprise": "surprise",
-    "contempt": "contempt",
-}
-
-# higher == more stressed
-STRESS_WEIGHTS = {
-    "anger": 0.95,
-    "fear": 0.90,
-    "disgust": 0.70,
-    "sad": 0.80,
-    "surprise": 0.55,
-    "neutral": 0.30,
-    "contempt": 0.65,
-    "happy": 0.10,
-}
-
-def _bucket(p: float) -> str:
-    return "Low" if p < 33 else ("Medium" if p < 66 else "High")
-
-emo_pipe = pipeline("image-classification", model=EMO_MODEL, device=0 if DEVICE == "cuda" else -1)
+    return out
 
-def _pipe_to_probs(res: List[Dict[str, Any]]) -> Dict[str, float]:
-    acc: Dict[str, float] = {}
-    for r in res:
-        label = (r.get("label") or r.get("emotion") or "").lower()
-        if not label:
-            continue
-        key = EMO_MAP.get(label, label)
-        score = float(r.get("score") or r.get("confidence") or r.get("confidence_pct", 0.0) / 100.0)
-        acc[key] = acc.get(key, 0.0) + score
-    Z = sum(acc.values()) or 1.0
-    for k in list(acc.keys()):
-        acc[k] = acc[k] / Z
-    return acc
+def knn_probs(pil_img: Image.Image, k: int = KNN_K_FOR_CLASS) -> Dict[str, float]:
+    q = embed_image(pil_img).reshape(1, -1)
+    k = min(k, len(DATASET_IMAGES))
+    dists, idxs = NN_MODEL.kneighbors(q, n_neighbors=k)
+    sims = 1.0 - dists[0]  # higher is better
+    sims = np.maximum(sims, 0.0)
+    votes: Dict[str, float] = {}
+    for sim, idx in zip(sims, idxs[0]):
+        lbl = DATASET_LABELS[int(idx)]
+        if lbl in ALLOWED:
+            votes[lbl] = votes.get(lbl, 0.0) + float(sim)
+    Z = sum(votes.values()) or 1.0
+    return {k: v / Z for k, v in votes.items()}
 
 def emotions_top3(pil_img: Image.Image) -> List[List[Any]]:
-    res = emo_pipe(pil_img)
-    probs = _pipe_to_probs(res)
+    probs = knn_probs(pil_img)
     items = sorted(probs.items(), key=lambda kv: kv[1], reverse=True)[:3]
     table = []
     for i, (emo, p) in enumerate(items, start=1):
         table.append([i, emo, round(100.0 * p, 2)])
-    return table  # [[rank, emotion, pct]]
+    # משלימים אם יש פחות מ-3
+    seen = {r[1] for r in table}
+    for fill in ["neutral", "other"]:
+        if len(table) >= 3: break
+        if fill in ALLOWED and fill not in seen:
+            table.append([len(table)+1, fill, 0.0])
+    return table
 
 def stress_index(pil_img: Image.Image) -> Tuple[str, float]:
-    res = emo_pipe(pil_img)
-    probs = _pipe_to_probs(res)
-    raw = 0.0
-    for k, v in probs.items():
-        w = STRESS_WEIGHTS.get(k, 0.5)
-        raw += v * w
+    probs = knn_probs(pil_img)
+    raw = sum(probs.get(k, 0.0) * STRESS_WEIGHTS.get(k, 0.5) for k in ALLOWED)
     pct = max(0.0, min(100.0, 100.0 * raw))
     return f"{pct:.1f}%  ({_bucket(pct)})", pct
 
-# ---------- Optional: SD image variations (1 image only) ----------
+# ---------- Optional: SD image variations ----------
 sd_pipe = None
 if USE_SD_VARIATIONS:
     try:
@@ -172,15 +190,13 @@ def generate_one_variation(pil_img: Image.Image, steps: int) -> Image.Image:
     return out
 
 # =====================  GRADIO UI  =====================
-CSS = """
-.box { border: 1px solid #e5e7eb; border-radius: 12px; padding: 10px; }
-"""
+CSS = ".box { border: 1px solid #e5e7eb; border-radius: 12px; padding: 10px; }"
 
-with gr.Blocks(title="Face Emotion & Stress Analyzer — CPU-friendly", css=CSS, fill_height=False) as demo:
+with gr.Blocks(title="Face Emotion & Stress Analyzer — KNN over tukey dataset", css=CSS, fill_height=False) as demo:
     gr.Markdown(
-        "### Face Emotion & Stress Analyzer — CPU-friendly\n"
+        "### Face Emotion & Stress Analyzer — **KNN over `tukey/human_face_emotions_roboflow`**\n"
         "- Embeddings: **laion/CLIP-ViT-H-14-laion2B-s32B-b79K** (open_clip)\n"
-        "- Emotion model: **prithivMLmods/Facial-Emotion-Detection-SigLIP2**\n"
+        "- Emotion model: **KNN using labels from `tukey/human_face_emotions_roboflow`**\n"
         "- Optional SD variations: **lambdalabs/sd-image-variations-diffusers** (1 synthetic only)\n"
         "- Right column shows nearest 5 images from the dataset (clickable)."
     )
@@ -246,9 +262,7 @@ with gr.Blocks(title="Face Emotion & Stress Analyzer — CPU-friendly", css=CSS,
         gal = nearest5(img)  # list[(PIL, caption)]
         gal_imgs = [g[0] for g in gal]
         gal_caps = [g[1] for g in gal]
-        # gr.Gallery accepts [(img, caption), ...]
         gallery = [(im, cap) for im, cap in zip(gal_imgs, gal_caps)]
-        # return
         return t3, s_label, gallery, gal_imgs, list(range(len(gal_imgs)))
 
     upload_image.change(
@@ -258,7 +272,6 @@ with gr.Blocks(title="Face Emotion & Stress Analyzer — CPU-friendly", css=CSS,
     )
 
     def on_gallery_select(evt: gr.SelectData, imgs: List[Image.Image], idxs: List[int]):
-        # evt.index is the clicked cell
         if imgs is None or not imgs:
             return [], ""
         i = int(evt.index) if evt is not None else 0
@@ -291,4 +304,4 @@ with gr.Blocks(title="Face Emotion & Stress Analyzer — CPU-friendly", css=CSS,
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()