Update app.py

app.py

@@ -8,36 +8,35 @@ from PIL import Image
 
 import torch
 import open_clip
-from datasets import load_dataset, DatasetDict
+from datasets import load_dataset
 from sklearn.neighbors import NearestNeighbors
-
 from diffusers import StableDiffusionImageVariationPipeline
 
 # -----------------------------
 # Config
 # -----------------------------
 DATASET_ID = "tukey/human_face_emotions_roboflow"
-EMB_MODEL_NAME = "ViT-H-14"                      # open_clip model name
-EMB_PRETRAINED = "laion2b_s32b_b79k"             # maps to laion/CLIP-ViT-H-14-laion2B-s32B-b79K
+EMB_MODEL_NAME = "ViT-H-14"                   # open_clip model name
+EMB_PRETRAINED = "laion2b_s32b_b79k"          # laion/CLIP-ViT-H-14-laion2B-s32B-b79K
 GEN_MODEL_ID = "lambdalabs/sd-image-variations-diffusers"
 
-CACHE_DIR = Path("./cache")
-CACHE_DIR.mkdir(parents=True, exist_ok=True)
+CACHE_DIR = Path("./cache"); CACHE_DIR.mkdir(parents=True, exist_ok=True)
 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"
 
-# generation defaults
-N_SYN = 12                 # generate more, then keep top-5
+# Default speed settings (can be overridden by Fast mode at runtime)
+INDEX_MAX = 1000         # cap number of dataset items used for the KNN index (first run only)
+BATCH_SIZE = 32          # batch size for embedding build
+N_SYN = 6                # how many variations to generate before picking top-5
 NUM_SYN_TO_SHOW = 5
-STEPS = 35
+STEPS = 20               # diffusion steps
 GUIDANCE_SCALES = [2.5, 3.0, 3.5, 4.0]
 
-# device selection
 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 (kept from your Colab logic)
+# Canonical label + stress mapping (from your Colab)
 # -----------------------------
 CANON = {"anger","disgust","fear","happy","neutral","sad","surprise","contempt"}
 CANON_MAP = {
@@ -57,17 +56,14 @@ STRESS_W = {
 }
 def _bucket(pct: float) -> str:
     return "Low" if pct < 33 else ("Medium" if pct < 66 else "High")
-
 def stress_from_top3(res: List[Dict]) -> Tuple[float, str]:
     probs = {}
     for r in res:
         lbl = CANON_MAP.get(str(r["emotion"]).lower(), str(r["emotion"]).lower())
-        if lbl not in CANON:
-            continue
+        if lbl not in CANON: continue
         probs[lbl] = probs.get(lbl, 0.0) + float(r["confidence_pct"]) / 100.0
     Z = sum(probs.values()) or 1.0
-    for k in list(probs):
-        probs[k] /= Z
+    for k in list(probs): probs[k] /= Z
     s01 = sum(probs.get(k, 0.0) * STRESS_W.get(k, 0.0) for k in probs)
     s01 = max(0.0, min(1.0, s01))
     pct = round(s01 * 100.0, 2)
@@ -89,108 +85,81 @@ _dataset_for_labels = None
 # -----------------------------
 def _load_openclip():
     global _openclip_model, _preprocess
-    if _openclip_model is not None and _preprocess is not None:
-        return _openclip_model, _preprocess
+    if _openclip_model is not None: return _openclip_model, _preprocess
     model, _, preprocess = open_clip.create_model_and_transforms(
-        model_name=EMB_MODEL_NAME,
-        pretrained=EMB_PRETRAINED,
-        device=DEVICE
+        model_name=EMB_MODEL_NAME, pretrained=EMB_PRETRAINED, device=DEVICE
     )
     model.eval()
     _openclip_model, _preprocess = model, preprocess
     return _openclip_model, _preprocess
 
-def _ensure_knn_index():
-    """Build (first run) or load a memmap + KNN index over the dataset embeddings."""
+def _fit_knn(X):
+    return NearestNeighbors(metric="cosine", algorithm="brute").fit(X)
+
+def _ensure_knn_index(index_max: int | None = None, batch_size: int | None = None):
+    """Build (first run) or load a memmap + KNN over dataset embeddings."""
     global _nn, _X, _labels_source, _dataset_for_labels
 
-    if _nn is not None and _X is not None:
+    if _nn is not None and _X is not None:  # already ready
         return
 
-    # Load dataset (train split; if missing, fallback to full)
+    index_max = index_max or INDEX_MAX
+    batch_size = batch_size or BATCH_SIZE
+
     dataset = load_dataset(DATASET_ID, split="train")
+    if index_max:
+        dataset = dataset.select(range(min(index_max, len(dataset))))
     _dataset_for_labels = dataset
-
     N = len(dataset)
 
-    # If memmaps already exist and meta is present -> load
+    # try loading existing cache if it matches N
     if EMB_MEMMAP_PATH.exists() and KNN_META_PATH.exists():
-        meta = json.load(open(KNN_META_PATH, "r"))
-        N_meta, D = int(meta["N"]), int(meta["D"])
-        if N_meta == N:
+        meta = json.load(open(KNN_META_PATH))
+        if int(meta.get("N", -1)) == N:
+            D = int(meta["D"])
             X = np.memmap(EMB_MEMMAP_PATH, mode="r", dtype="float32", shape=(N, D))
-            # labels memmap optional; if missing, we can fetch labels on the fly
-            labels = None
-            if LABELS_MEMMAP_PATH.exists():
-                labels = np.memmap(LABELS_MEMMAP_PATH, mode="r", dtype="U32", shape=(N,))
-            _fit_knn(X)
-            _X = X
-            _labels_source = labels
+            labels = np.memmap(LABELS_MEMMAP_PATH, mode="r", dtype="U32", shape=(N,)) if LABELS_MEMMAP_PATH.exists() else None
+            _X = X; _labels_source = labels; _nn = _fit_knn(X)
             return
 
-    # Build embeddings (first run)
+    # build embeddings (batched)
     model, preprocess = _load_openclip()
-    D = None
-    X_w = None
+    labels_mm = np.memmap(LABELS_MEMMAP_PATH, mode="w+", dtype="U32", shape=(N,))
+    X_w = None; D = None
 
     def _label_of(i):
-        try:
-            ans = dataset[i]["qa"][0]["answer"]
-            return str(ans) if ans is not None else ""
-        except Exception:
-            return ""
-
-    # write labels memmap
-    labels_mm = np.memmap(LABELS_MEMMAP_PATH, mode="w+", dtype="U32", shape=(N,))
+        try: return str(dataset[i]["qa"][0]["answer"] or "")
+        except Exception: return ""
 
     with torch.no_grad():
-        # do the first item to get D
-        x0 = preprocess(dataset[0]["image"]).unsqueeze(0)
-        if DEVICE in ("cuda", "mps"):
-            x0 = x0.to(DEVICE)
-        v0 = model.encode_image(x0).float()
-        v0 = v0 / v0.norm(dim=-1, keepdim=True)
-        D = v0.shape[1]
-        X_w = np.memmap(EMB_MEMMAP_PATH, mode="w+", dtype="float32", shape=(N, D))
-        X_w[0] = v0.detach().cpu().numpy().squeeze()
-        labels_mm[0] = _label_of(0)
-
-        # rest
-        for i in range(1, N):
-            xi = preprocess(dataset[i]["image"]).unsqueeze(0)
-            if DEVICE in ("cuda", "mps"):
-                xi = xi.to(DEVICE)
-            vi = model.encode_image(xi).float()
-            vi = vi / vi.norm(dim=-1, keepdim=True)
-            X_w[i] = vi.detach().cpu().numpy().squeeze()
-            labels_mm[i] = _label_of(i)
-
-    # flush to disk
-    del X_w
-    gc.collect()
-
-    # Save meta, reload read-only view, fit knn
+        for start in range(0, N, batch_size):
+            end = min(start + batch_size, N)
+            imgs = [dataset[i]["image"].convert("RGB") for i in range(start, end)]
+            x = torch.stack([preprocess(im) for im in imgs])
+            if DEVICE in ("cuda", "mps"): x = x.to(DEVICE)
+            v = model.encode_image(x).float()
+            v = v / v.norm(dim=-1, keepdim=True)
+
+            if X_w is None:
+                D = v.shape[1]
+                X_w = np.memmap(EMB_MEMMAP_PATH, mode="w+", dtype="float32", shape=(N, D))
+            X_w[start:end] = v.detach().cpu().numpy()
+
+            for i in range(start, end):
+                labels_mm[i] = _label_of(i)
+
+    del X_w; gc.collect()
     json.dump({"N": int(N), "D": int(D)}, open(KNN_META_PATH, "w"))
     X = np.memmap(EMB_MEMMAP_PATH, mode="r", dtype="float32", shape=(N, D))
     labels = np.memmap(LABELS_MEMMAP_PATH, mode="r", dtype="U32", shape=(N,))
-    _fit_knn(X)
-    _X = X
-    _labels_source = labels
-
-def _fit_knn(X):
-    global _nn
-    _nn = NearestNeighbors(metric="cosine", algorithm="brute").fit(X)
+    _X = X; _labels_source = labels; _nn = _fit_knn(X)
 
 def _label_by_idx(i: int):
     global _labels_source, _dataset_for_labels
     if _labels_source is not None:
-        lab = str(_labels_source[i])
-        return lab if lab else None
-    # fallback: live label read
-    try:
-        return _dataset_for_labels[i]["qa"][0]["answer"]
-    except Exception:
-        return None
+        lab = str(_labels_source[i]);  return lab if lab else None
+    try: return _dataset_for_labels[i]["qa"][0]["answer"]
+    except Exception: return None
 
 # -----------------------------
 # Embedding + inference utils
@@ -199,60 +168,42 @@ def embed_image(img: Image.Image) -> np.ndarray:
     model, preprocess = _load_openclip()
     with torch.no_grad():
         x = preprocess(img.convert("RGB")).unsqueeze(0)
-        if DEVICE in ("cuda", "mps"):
-            x = x.to(DEVICE)
+        if DEVICE in ("cuda", "mps"): x = x.to(DEVICE)
         v = model.encode_image(x).float()
         v = v / v.norm(dim=-1, keepdim=True)
     return v.detach().cpu().numpy().squeeze()
 
 def _top3_emotions_weighted_from_embed(q: np.ndarray,
-                                       start_k: int = 30, step: int = 30, method: str = "softmax", tau: float = 0.1):
+                                       start_k: int = 30, step: int = 30,
+                                       method: str = "softmax", tau: float = 0.1):
     _ensure_knn_index()
-    max_k = _X.shape[0]
-    k = min(start_k, max_k)
-    scores: Dict[str, float] = {}
-
+    max_k = _X.shape[0]; k = min(start_k, max_k)
     while True:
         dist, idx = _nn.kneighbors(q.reshape(1, -1), n_neighbors=k)
         idx, dist = idx[0], dist[0]
-        sims = 1.0 - dist
-        sims = np.clip(sims, 0.0, None)
+        sims = np.clip(1.0 - dist, 0.0, None)
         w = np.exp(sims / tau) if method == "softmax" else sims
 
-        scores.clear()
-        total_w = 0.0
+        scores: Dict[str, float] = {}
         for i, wi in zip(idx, w):
             lab = _label_by_idx(int(i))
-            if lab is None:
-                continue
+            if lab is None: continue
             lab = CANON_MAP.get(str(lab).lower(), str(lab).lower())
             scores[lab] = scores.get(lab, 0.0) + float(wi)
-            total_w += float(wi)
 
-        # stop when we have >= 3 unique emotions
-        if len([k for k in scores.keys() if k in CANON]) >= 3 or k == max_k:
+        if len([k for k in scores if k in CANON]) >= 3 or k == max_k:
             break
         k = min(k + step, max_k)
 
-    if not scores:
-        return []
-
-    # keep only canonical keys
     scores = {k: v for k, v in scores.items() if k in CANON and v > 0}
-    if not scores:
-        return []
-
+    if not scores: return []
     top_items = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:3]
     vals = np.array([v for _, v in top_items], dtype=np.float32)
     pct = (vals / vals.sum()) * 100.0 if vals.sum() > 0 else np.zeros_like(vals)
-    return [
-        {"rank": i + 1, "emotion": lab, "confidence_pct": int(round(p))}
-        for i, ((lab, _), p) in enumerate(zip(top_items, pct))
-    ]
+    return [{"rank": i+1, "emotion": lab, "confidence_pct": int(round(p))}
+            for i, ((lab, _), p) in enumerate(zip(top_items, pct))]
 
 def analyze_face(image: Image.Image):
-    """Return top-3 emotions + stress for the original image."""
-    _ensure_knn_index()
     q = embed_image(image)
     top3 = _top3_emotions_weighted_from_embed(q)
     stress_pct, stress_lbl = stress_from_top3(top3)
@@ -263,55 +214,31 @@ def analyze_face(image: Image.Image):
 # -----------------------------
 def _get_gen_pipe():
     global _gen_pipe
-    if _gen_pipe is not None:
-        return _gen_pipe
+    if _gen_pipe is not None: return _gen_pipe
     gen_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
     pipe = StableDiffusionImageVariationPipeline.from_pretrained(
-        GEN_MODEL_ID,
-        revision="v2.0",
-        torch_dtype=gen_dtype
-    )
-    pipe = pipe.to(DEVICE)
+        GEN_MODEL_ID, revision="v2.0", torch_dtype=gen_dtype
+    ).to(DEVICE)
     _gen_pipe = pipe
     return _gen_pipe
 
-def generate_synthetics(base_image: Image.Image, base_embed: np.ndarray):
-    """Generate N_SYN variations, compute similarity to original embedding, keep top-5."""
+def generate_synthetics(base_image: Image.Image, base_embed: np.ndarray, n_syn: int, steps: int):
     pipe = _get_gen_pipe()
-
-    # Deterministic seed stream
     base_gen = torch.Generator(device="cpu").manual_seed(42)
-
     records = []
-    for i in range(N_SYN):
+    for _ in range(n_syn):
         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)
-
-        out = pipe(
-            image=base_image.convert("RGB"),
-            guidance_scale=gs,
-            num_inference_steps=STEPS,
-            generator=g
-        )
+        out = pipe(image=base_image.convert("RGB"),
+                   guidance_scale=gs, num_inference_steps=steps, generator=g)
         img = out.images[0]
-
-        # embed and compute similarity to the original
         emb = embed_image(img)
-        sim = float(np.dot(emb, base_embed))  # both normalized
-
-        # top3 + stress for each synthetic
+        sim = float(np.dot(emb, base_embed))
         top3_syn = _top3_emotions_weighted_from_embed(emb)
         stress_pct, stress_lbl = stress_from_top3(top3_syn)
-
-        records.append({
-            "image": img,
-            "similarity": sim,
-            "top3": top3_syn,
-            "stress": f"{stress_pct}% ({stress_lbl})"
-        })
-
-    # keep best NUM_SYN_TO_SHOW by similarity
+        records.append({"image": img, "similarity": sim, "top3": top3_syn,
+                        "stress": f"{stress_pct}% ({stress_lbl})"})
     records.sort(key=lambda r: r["similarity"], reverse=True)
     return records[:NUM_SYN_TO_SHOW]
 
@@ -319,21 +246,20 @@ def generate_synthetics(base_image: Image.Image, base_embed: np.ndarray):
 # Gradio app
 # -----------------------------
 def _format_top3_for_table(top3: List[Dict]) -> List[List]:
-    rows = []
-    for r in top3:
-        rows.append([r["rank"], r["emotion"], r["confidence_pct"]])
-    return rows
+    return [[r["rank"], r["emotion"], r["confidence_pct"]] for r in top3]
 
-with gr.Blocks(title="Face Emotions + Stress (CLIP ViT-H-14 + SD Variations)") as demo:
+with gr.Blocks(title="Face Emotions + Stress (Fast)") as demo:
     gr.Markdown(
-        "## Face Emotion & Stress Analyzer\n"
-        "- Embeddings: **laion/CLIP-ViT-H-14-laion2B-s32B-b79K** (via `open_clip`)\n"
+        "## Face Emotion & Stress Analyzer (Fast)\n"
+        "- Embeddings: **laion/CLIP-ViT-H-14-laion2B-s32B-b79K** via `open_clip`\n"
         "- Synthetic variations: **lambdalabs/sd-image-variations-diffusers**\n"
         "- KNN labels from: **tukey/human_face_emotions_roboflow**\n"
+        "- First run builds a cached index (capped by `INDEX_MAX`).\n"
     )
 
     with gr.Row():
         inp = gr.Image(type="pil", label="Upload a face image", sources=["upload", "webcam"])
+        fast_mode = gr.Checkbox(value=True, label="Fast mode (smaller index & fewer synthetics)")
 
     analyze_btn = gr.Button("Analyze & Generate Synthetics")
 
@@ -342,59 +268,54 @@ with gr.Blocks(title="Face Emotions + Stress (CLIP ViT-H-14 + SD Variations)") a
             top3_tbl = gr.Dataframe(
                 headers=["Rank", "Emotion", "Confidence (%)"],
                 datatype=["number", "str", "number"],
-                interactive=False,
-                row_count=(3, "fixed"),
-                col_count=(3, "fixed"),
+                interactive=False, row_count=(3, "fixed"), col_count=(3, "fixed"),
                 label="Top-3 emotions (original image)"
             )
             stress_txt = gr.Label(label="Stress index (original)")
         with gr.Column():
             gal = gr.Gallery(
                 label="Top 5 synthetic variations (click one)",
-                columns=[5], height=200, preview=True
+                columns=[5], height=200, preview=True  # no selectable kwarg
             )
             syn_stress = gr.Label(label="Stress index (selected synthetic)")
             syn_top3 = gr.JSON(label="Top-3 emotions (selected synthetic)")
 
     status = gr.Markdown(visible=False)
+    syn_state = gr.State([])
 
-    # State to pass around generated records
-    syn_state = gr.State([])  # list of dicts: {image, similarity, top3, stress}
-
-    def run_pipeline(image: Image.Image):
+    def run_pipeline(image: Image.Image, fast: bool):
         try:
-            # Step 1: top3 + stress (original)
-            top3, stress, q = analyze_face(image)
+            # Tune runtime knobs
+            idx_max = 600 if fast else INDEX_MAX
+            bs = 32 if fast else BATCH_SIZE
+            n_syn = 4 if fast else N_SYN
+            steps = 16 if fast else STEPS
 
-            # Step 2: synthetics + pick top 5
-            syn = generate_synthetics(image, q)
+            # Ensure (or build) index with chosen cap/batch
+            _ensure_knn_index(index_max=idx_max, batch_size=bs)
 
-            # gallery expects a list of images or (image, caption)
+            # Original image analysis
+            top3, stress, q = analyze_face(image)
+
+            # Synthetics
+            syn = generate_synthetics(image, q, n_syn=n_syn, steps=steps)
             items = [(r["image"], f"sim={r['similarity']:.3f}") for r in syn]
-            top3_rows = _format_top3_for_table(top3)
-            return top3_rows, stress, items, syn, gr.update(visible=False), None
+
+            return _format_top3_for_table(top3), stress, items, syn, gr.update(visible=False), None
         except Exception as e:
             return None, None, None, [], gr.update(visible=True, value=f"**Error:** {e}"), None
 
     analyze_btn.click(
-        run_pipeline,
-        inputs=[inp],
+        run_pipeline, inputs=[inp, fast_mode],
         outputs=[top3_tbl, stress_txt, gal, syn_state, status, syn_top3]
     )
 
     def on_gallery_select(evt: gr.SelectData, syn_records: List[Dict]):
-        # evt.index is the clicked item index in gallery
-        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]
+        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]
         return gr.update(value=rec["stress"]), gr.update(value=rec["top3"])
 
-    gal.select(
-        fn=on_gallery_select,
-        inputs=[syn_state],
-        outputs=[syn_stress, syn_top3]
-    )
+    gal.select(fn=on_gallery_select, inputs=[syn_state], outputs=[syn_stress, syn_top3])
 
 if __name__ == "__main__":
     demo.launch()