Update app.py

app.py

@@ -13,7 +13,7 @@ from sklearn.neighbors import NearestNeighbors
 from diffusers import StableDiffusionImageVariationPipeline
 
 # -----------------------------
-# Config (CPU-friendly defaults)
+# Config (CPU-friendly)
 # -----------------------------
 DATASET_ID = "tukey/human_face_emotions_roboflow"
 EMB_MODEL_NAME = "ViT-H-14"            # open_clip model name
@@ -25,18 +25,18 @@ EMB_MEMMAP_PATH = CACHE_DIR / "clip_vith14_laion2b.float32.memmap"
 LABELS_MEMMAP_PATH = CACHE_DIR / "labels.U32.memmap"
 KNN_META_PATH = CACHE_DIR / "knn_meta.json"
 
-# **tiny index** + light gen for CPU
-INDEX_MAX_DEFAULT = 80     # small subset → fast
+# tiny index + light generation
+INDEX_MAX_DEFAULT = 80
 BATCH_SIZE_DEFAULT = 32
 N_SYN_DEFAULT = 3
 STEPS_DEFAULT = 12
 GUIDANCE_SCALES = [2.5, 3.0, 3.5, 4.0]
-NUM_SYN_TO_SHOW = 5        # show up to 5 if you generate that many
+NUM_SYN_TO_SHOW = 5
 
 DEVICE = "cuda" if torch.cuda.is_available() else ("mps" if hasattr(torch.backends, "mps") and torch.backends.mps.is_available() else "cpu")
 
 # -----------------------------
-# Canonical label + stress mapping (your logic)
+# Canonical labels + stress
 # -----------------------------
 CANON = {"anger","disgust","fear","happy","neutral","sad","surprise","contempt"}
 CANON_MAP = {
@@ -87,7 +87,7 @@ def _load_openclip():
 def _fit_knn(X): return NearestNeighbors(metric="cosine", algorithm="brute").fit(X)
 
 def _ensure_knn_index(index_max: int, batch_size: int, progress: gr.Progress | None = None):
-    """Build (first run) or load a memmap + KNN over a tiny dataset subset."""
+    """Build (first run) or load a tiny memmap + KNN over a subset of the dataset."""
     global _nn, _X, _labels_source, _dataset_for_labels
 
     if _nn is not None and _X is not None:
@@ -99,7 +99,6 @@ def _ensure_knn_index(index_max: int, batch_size: int, progress: gr.Progress | N
     _dataset_for_labels = dataset
     N = len(dataset)
 
-    # load cache if exists with same N
     if EMB_MEMMAP_PATH.exists() and KNN_META_PATH.exists():
         meta = json.load(open(KNN_META_PATH))
         if int(meta.get("N", -1)) == N:
@@ -109,12 +108,10 @@ def _ensure_knn_index(index_max: int, batch_size: int, progress: gr.Progress | N
             _X = X; _labels_source = labels; _nn = _fit_knn(X)
             return
 
-    # build tiny embedding index (batched)
     model, preprocess = _load_openclip()
     labels_mm = np.memmap(LABELS_MEMMAP_PATH, mode="w+", dtype="U32", shape=(N,))
     X_w = None; D = None
 
-    step = 0
     with torch.no_grad():
         for start in range(0, N, batch_size):
             end = min(start + batch_size, N)
@@ -130,7 +127,6 @@ def _ensure_knn_index(index_max: int, batch_size: int, progress: gr.Progress | N
             for i in range(start, end):
                 try: labels_mm[i] = str(dataset[i]["qa"][0]["answer"] or "")
                 except Exception: labels_mm[i] = ""
-            step += 1
             if progress: progress(((end)/N), desc=f"Building index {end}/{N}")
 
     del X_w; gc.collect()
@@ -190,8 +186,22 @@ def analyze_face(image: Image.Image):
     stress_pct, stress_lbl = stress_from_top3(top3)
     return top3, f"{stress_pct}% ({stress_lbl})", q
 
+# ----- Nearest neighbors images from dataset -----
+def _get_dataset_image(i: int) -> Image.Image:
+    return _dataset_for_labels[int(i)]["image"].convert("RGB")
+
+def nearest_k_images_from_dataset(q_emb: np.ndarray, k: int = 5):
+    dist, idx = _nn.kneighbors(q_emb.reshape(1, -1), n_neighbors=k)
+    dist, idx = dist[0], idx[0]
+    sims = (1.0 - dist).tolist()
+    out = []
+    for i, s in zip(idx, sims):
+        img = _get_dataset_image(int(i))
+        out.append((img, float(s), int(i)))
+    return out
+
 # -----------------------------
-# Generator (optional, on click)
+# Generator (optional)
 # -----------------------------
 def _get_gen_pipe():
     global _gen_pipe
@@ -207,7 +217,7 @@ def generate_synthetics(base_image: Image.Image, base_embed: np.ndarray, n_syn:
     pipe = _get_gen_pipe()
     base_gen = torch.Generator(device="cpu").manual_seed(42)
     records = []
-    for i in progress.tqdm(range(n_syn), desc="Generating"):
+    for _ in progress.tqdm(range(n_syn), desc="Generating"):
         seed = int(torch.randint(0, 2**31 - 1, (1,), generator=base_gen).item())
         gs = random.choice(GUIDANCE_SCALES)
         g = torch.Generator(device="cpu").manual_seed(seed)
@@ -222,7 +232,7 @@ def generate_synthetics(base_image: Image.Image, base_embed: np.ndarray, n_syn:
     return records[:NUM_SYN_TO_SHOW]
 
 # -----------------------------
-# Gradio app (two-step: Analyze → (optional) Generate)
+# UI
 # -----------------------------
 def _format_top3_for_table(top3: List[Dict]) -> List[List]:
     return [[r["rank"], r["emotion"], r["confidence_pct"]] for r in top3]
@@ -231,8 +241,8 @@ with gr.Blocks(title="Face Emotions + Stress (CPU Fast)") as demo:
     gr.Markdown(
         "## Face Emotion & Stress Analyzer — CPU-friendly\n"
         "- Embeddings: **laion/CLIP-ViT-H-14-laion2B-s32B-b79K** (open_clip)\n"
-        "- Synthetic variations (optional): **lambdalabs/sd-image-variations-diffusers**\n"
-        "- Uses a tiny cached index for speed on free CPU.\n"
+        "- Optional SD variations: **lambdalabs/sd-image-variations-diffusers**\n"
+        "- Also shows **nearest 5 images from the dataset** for 1-click results.\n"
     )
 
     with gr.Row():
@@ -252,35 +262,66 @@ with gr.Blocks(title="Face Emotions + Stress (CPU Fast)") as demo:
             )
             stress_txt = gr.Label(label="Stress index (original)")
         with gr.Column():
-            # generation controls (optional)
-            n_syn = gr.Slider(0, 5, value=N_SYN_DEFAULT, step=1, label="How many variations to generate")
+            # Nearest 5 from dataset (one-click examples)
+            nn_gal = gr.Gallery(
+                label="Nearest 5 from dataset (click one)",
+                columns=[5], height=220, preview=True
+            )
+            nn_stress = gr.Label(label="Stress (nearest image)")
+            nn_top3 = gr.JSON(label="Top-3 emotions (nearest image)")
+
+            # Optional generator
+            n_syn = gr.Slider(0, 5, value=N_SYN_DEFAULT, step=1, label="How many SD variations to generate")
             steps = gr.Slider(8, 30, value=STEPS_DEFAULT, step=2, label="Diffusion steps (higher = slower/better)")
             gen_btn = gr.Button("Generate variations (optional)")
             gal = gr.Gallery(label="Synthetic variations (click one)", columns=[5], height=220, preview=True)
-            syn_stress = gr.Label(label="Stress index (selected synthetic)")
+            syn_stress = gr.Label(label="Stress (selected synthetic)")
             syn_top3 = gr.JSON(label="Top-3 emotions (selected synthetic)")
 
     status = gr.Markdown(visible=False)
 
-    # state we pass between steps
-    syn_state = gr.State([])          # list of generated records
-    q_state = gr.State(None)          # original embedding
-    img_state = gr.State(None)        # original image (for gen step)
+    # State
+    syn_state = gr.State([])    # generated variations
+    q_state = gr.State(None)    # embedding of original image
+    img_state = gr.State(None)  # original image
 
+    # ---- Analyze ----
     def do_analyze(image: Image.Image, cap: int, batch: int, progress=gr.Progress(track_tqdm=True)):
         try:
             _ensure_knn_index(index_max=int(cap), batch_size=int(batch), progress=progress)
             top3, stress, q = analyze_face(image)
-            return _format_top3_for_table(top3), stress, [], [], q, image, gr.update(visible=False)
+
+            # nearest 5 images from dataset
+            neigh = nearest_k_images_from_dataset(np.array(q, dtype=np.float32), k=5)
+            nn_items = [(im, f"sim={sim:.3f} • idx={idx}") for im, sim, idx in neigh]
+
+            # return: top3, stress, nn gallery, (empty SD gallery), syn_state, q, img, status
+            return (_format_top3_for_table(top3), stress,
+                    nn_items, [], [], q, image, gr.update(visible=False))
         except Exception as e:
-            return None, None, [], [], None, None, gr.update(visible=True, value=f"**Error:** {e}")
+            return None, None, [], [], [], None, None, gr.update(visible=True, value=f"**Error:** {e}")
 
     analyze_btn.click(
         do_analyze,
         inputs=[inp, idx_cap, bs],
-        outputs=[top3_tbl, stress_txt, gal, syn_state, q_state, img_state, status]
+        outputs=[top3_tbl, stress_txt, nn_gal, gal, syn_state, q_state, img_state, status]
     )
 
+    # ---- One-click on a nearest image ----
+    def on_nn_select(evt: gr.SelectData, q):
+        if q is None:
+            return gr.update(value="Analyze first"), None
+        neigh = nearest_k_images_from_dataset(np.array(q, dtype=np.float32), k=5)
+        i = max(0, min(int(evt.index), len(neigh)-1))
+        img, _, _ = neigh[i]
+        emb = embed_image(img)
+        top3 = _top3_emotions_weighted_from_embed(emb)
+        stress_pct, stress_lbl = stress_from_top3(top3)
+        return f"{stress_pct}% ({stress_lbl})", top3
+
+    nn_gal.select(fn=on_nn_select, inputs=[q_state], outputs=[nn_stress, nn_top3])
+
+    # ---- Optional: generate SD variations ----
     def do_generate(n: int, s: int, q, img, progress=gr.Progress()):
         if q is None or img is None:
             return [], [], gr.update(visible=True, value="**Error:** Analyze first."), None
@@ -297,6 +338,7 @@ with gr.Blocks(title="Face Emotions + Stress (CPU Fast)") as demo:
         outputs=[gal, syn_state, status, syn_top3]
     )
 
+    # select from generated synthetics
     def on_gallery_select(evt: gr.SelectData, syn_records: List[Dict]):
         if not syn_records or evt is None: return gr.update(value=None), gr.update(value=None)
         i = int(evt.index); rec = syn_records[i]