Update app.py

app.py

@@ -19,7 +19,6 @@ def _softmax_with_temperature(logits: torch.Tensor, temperature: float) -> torch
     if temperature <= 0:
         temperature = 1.0
     scaled = logits / float(temperature)
-    # subtract max for numerical stability
     scaled = scaled - torch.max(scaled)
     exp = torch.exp(scaled)
     return exp / torch.sum(exp)
@@ -29,64 +28,60 @@ def _entropy(probs: np.ndarray) -> float:
     p = probs[probs > 0]
     return float(-(p * np.log(p)).sum())
 
-def _make_bar(labels, probs):
-    """Return a matplotlib horizontal bar chart of top-K probabilities."""
+def _kl(p: np.ndarray, q: np.ndarray) -> float:
+    """KL divergence KL(p||q), with small epsilon for stability."""
+    eps = 1e-12
+    p = p + eps
+    q = q + eps
+    return float(np.sum(p * np.log(p / q)))
+
+def _jsd(p: np.ndarray, q: np.ndarray) -> float:
+    """Jensen–Shannon divergence (symmetric, bounded)."""
+    m = 0.5 * (p + q)
+    return 0.5 * _kl(p, m) + 0.5 * _kl(q, m)
+
+def _make_bar(labels, probs, title="Top-K predicted classes"):
+    """Return a matplotlib horizontal bar chart of probabilities."""
     fig, ax = plt.subplots(figsize=(6, 3.2))
     y = np.arange(len(labels))
-    ax.barh(y, probs)  # do not set colors (keep default)
+    ax.barh(y, probs)  # default colors only (per teaching tool rules)
     ax.set_yticks(y, labels)
     ax.invert_yaxis()
     ax.set_xlim(0, 1)
     ax.set_xlabel("Probability")
-    ax.set_title("Top-K predicted classes")
+    ax.set_title(title)
     fig.tight_layout()
     return fig
 
-def analyze(img, top_k=5, temperature=1.0):
+def _analyze_single(img, top_k=5, temperature=1.0):
     """
-    Run the image through the classifier and expose:
-      - Top-K probabilities (bar chart + table)
-      - Pre-softmax logits for those Top-K
-      - Uncertainty metrics (entropy; top-1 margin; cumulative top-K)
-      - Preprocessing info and inference time
+    Return: (quick_label_dict, bar_plot, table_rows, notes_markdown, full_probs_numpy)
     """
     if img is None:
-        return (
-            {"<no image>": 1.0},   # quick glance label block
-            None,                  # bar plot
-            [],                    # table rows
-            "Please upload an image.",  # metrics markdown
-        )
+        return ({"<no image>": 1.0}, None, [], "Please upload an image.", None)
 
     t0 = time.perf_counter()
     inputs = processor(images=img, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
     logits = outputs.logits[0]  # shape [num_labels]
-    # Temperature-scaled softmax
     probs = _softmax_with_temperature(logits, temperature)
-    # Top-K
+
     k = max(1, int(top_k))
     k = min(k, probs.shape[0])
     top_vals, top_idx = torch.topk(probs, k=k, dim=-1)
     top_idx = top_idx.tolist()
     top_vals = top_vals.tolist()
-
     labels = [LABELS[i] for i in top_idx]
     logits_top = [float(logits[i]) for i in top_idx]
 
-    # Quick glance dict for gr.Label (keeps students oriented)
     quick = {lab: float(p) for lab, p in zip(labels, top_vals)}
-
-    # Bar chart
     fig = _make_bar(labels, top_vals)
 
-    # Table rows (Rank, Label, Probability, Logit)
     rows = []
     for rank, (lab, p, lg) in enumerate(zip(labels, top_vals, logits_top), start=1):
         rows.append([rank, lab, round(float(p), 6), round(float(lg), 6)])
 
-    # Metrics
     probs_np = probs.detach().cpu().numpy()
     H = _entropy(probs_np)
     top1 = float(top_vals[0])
@@ -95,7 +90,6 @@ def analyze(img, top_k=5, temperature=1.0):
     cum_topk = float(sum(top_vals))
     infer_ms = (time.perf_counter() - t0) * 1000.0
 
-    # Preprocessing info (resize/crop) from processor config if available
     size = processor.size if hasattr(processor, "size") else {}
     target_h = size.get("height", None)
     target_w = size.get("width", None)
@@ -108,32 +102,134 @@ def analyze(img, top_k=5, temperature=1.0):
         f"- Cumulative Top-K probability: **{cum_topk:.3f}**  \n\n"
         f"**Preprocessing & Runtime**  \n"
         f"- Processor target size: **{size_str}**  \n"
-        f"- Inference time: **{infer_ms:.1f} ms** (CPU)  \n\n"
-        f"_Tip:_ Adjust **Temperature** to watch softmax sharpen (T<1) or soften (T>1) the distribution."
+        f"- Inference time: **{infer_ms:.1f} ms**  \n"
     )
 
+    return quick, fig, rows, md, probs_np
+
+def _align_topk(labelsA, probsA, labelsB, probsB, K=5):
+    """Make a unified label set of size up to K using union-of-top labels then rank by max(prob)."""
+    dA = dict(zip(labelsA, probsA))
+    dB = dict(zip(labelsB, probsB))
+    union = set(labelsA) | set(labelsB)
+    # rank by max(prob from A, prob from B)
+    ranked = sorted(list(union), key=lambda x: max(dA.get(x, 0.0), dB.get(x, 0.0)), reverse=True)
+    chosen = ranked[:K]
+    a = [float(dA.get(l, 0.0)) for l in chosen]
+    b = [float(dB.get(l, 0.0)) for l in chosen]
+    return chosen, a, b
+
+def analyze_single(img, top_k=5, temperature=1.0):
+    quick, fig, rows, md, _ = _analyze_single(img, top_k, temperature)
     return quick, fig, rows, md
 
+def analyze_pair(imgA, imgB, top_k=5, temperature=1.0):
+    """
+    A/B analysis:
+      - show per-image quick dict, bar chart, table, notes
+      - show aligned Top-K delta bar and divergence metrics
+    """
+    # Analyze each side
+    qa, figa, rowsa, mda, pa = _analyze_single(imgA, top_k, temperature)
+    qb, figb, rowsb, mdb, pb = _analyze_single(imgB, top_k, temperature)
+
+    # If either missing, return as-is
+    if pa is None or pb is None:
+        return qa, figa, rowsa, mda, qb, figb, rowsb, mdb, None, "Upload both images for delta metrics."
+
+    # Build aligned top-K over labels
+    # We need label sets and probs for both to compute aligned bars
+    # Recover top-K labels directly from rows (rank, label, prob, logit)
+    labelsA = [r[1] for r in rowsa]
+    probsA  = [r[2] for r in rowsa]
+    labelsB = [r[1] for r in rowsb]
+    probsB  = [r[2] for r in rowsb]
+    chosen, a, b = _align_topk(labelsA, probsA, labelsB, probsB, K=max(int(top_k), 1))
+
+    # Delta bar (A−B)
+    deltas = [float(x - y) for x, y in zip(a, b)]
+    fig_delta = _make_bar([f"{lbl} (Δ)" for lbl in chosen], deltas, title="Aligned Top-K Δ Probabilities (A − B)")
+
+    # Distribution-level differences (full softmax vectors)
+    # Ensure same length and normalize to prob distributions
+    pa = pa / (pa.sum() + 1e-12)
+    pb = pb / (pb.sum() + 1e-12)
+    H_a = _entropy(pa)
+    H_b = _entropy(pb)
+    jsd = _jsd(pa, pb)
+
+    # Top-1 labels for each side
+    top1_a_idx = int(np.argmax(pa))
+    top1_b_idx = int(np.argmax(pb))
+    top1_a = LABELS[top1_a_idx]
+    top1_b = LABELS[top1_b_idx]
+
+    diff_md = (
+        f"**A/B Divergence**  \n"
+        f"- Jensen–Shannon divergence: **{jsd:.4f}** (0=same, higher=more different)  \n"
+        f"- Entropy A / B: **{H_a:.3f} / {H_b:.3f}** nats  \n"
+        f"- Top-1 A / B: **{top1_a} / {top1_b}**  \n"
+        f"- Aligned Top-K shown above is ranked by max(prob_A, prob_B).  \n"
+        f"_Tip:_ Try different crops/lighting or adjust **Temperature** to watch distributions change."
+    )
+
+    return qa, figa, rowsa, mda, qb, figb, rowsb, mdb, fig_delta, diff_md
+
 with gr.Blocks(fill_height=True, analytics_enabled=False) as demo:
-    gr.Markdown("# 🖼️ Image Classification — ‘Watch the Decision’\nUpload an image, then watch the classifier’s probabilities, logits, and uncertainty metrics update.")
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            img = gr.Image(type="pil", label="Upload an image (JPEG/PNG)", height=360)
-            topk = gr.Slider(1, 10, value=5, step=1, label="Top-K predictions")
-            temp = gr.Slider(0.25, 2.0, value=1.0, step=0.05, label="Softmax Temperature")
-            run = gr.Button("Analyze", variant="primary")
-        with gr.Column(scale=1):
-            gr.Markdown("### Quick glance")
-            glance = gr.Label(num_top_classes=10)
-            gr.Markdown("### Probabilities (Top-K)")
-            plot = gr.Plot()
-            gr.Markdown("### Details (Top-K)")
-            table = gr.Dataframe(headers=["Rank", "Label", "Probability", "Logit"], datatype=["number", "str", "number", "number"], row_count=5)
-            gr.Markdown("### Metrics & Notes")
-            notes = gr.Markdown()
-
-    run.click(analyze, [img, topk, temp], [glance, plot, table, notes])
+    gr.Markdown("# 🖼️ Image Classification — ‘Watch the Decision’\nVisualize probabilities, logits, entropy, and A/B deltas.\n\n"
+                "_Notes:_ Predictions reflect the ImageNet‑1k label space; unusual objects or logos may be misclassified. "
+                "Do not use for identity or sensitive inferences.")
+
+    with gr.Tab("Single Image"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                img = gr.Image(type="pil", label="Upload image (JPEG/PNG)", height=340)
+                topk = gr.Slider(1, 10, value=5, step=1, label="Top‑K predictions")
+                temp = gr.Slider(0.25, 2.0, value=1.0, step=0.05, label="Softmax Temperature")
+                run = gr.Button("Analyze", variant="primary")
+            with gr.Column(scale=1):
+                gr.Markdown("### Quick glance")
+                glance = gr.Label(num_top_classes=10)
+                gr.Markdown("### Probabilities (Top‑K)")
+                plot = gr.Plot()
+                gr.Markdown("### Details (Top‑K)")
+                table = gr.Dataframe(headers=["Rank", "Label", "Probability", "Logit"], datatype=["number", "str", "number", "number"], row_count=5)
+                gr.Markdown("### Metrics & Notes")
+                notes = gr.Markdown()
+        run.click(analyze_single, [img, topk, temp], [glance, plot, table, notes])
+
+    with gr.Tab("A/B Compare"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                imgA = gr.Image(type="pil", label="Image A", height=300)
+                imgB = gr.Image(type="pil", label="Image B", height=300)
+                topkAB = gr.Slider(1, 10, value=5, step=1, label="Aligned Top‑K")
+                tempAB = gr.Slider(0.25, 2.0, value=1.0, step=0.05, label="Softmax Temperature")
+                runAB = gr.Button("Analyze A/B", variant="primary")
+            with gr.Column(scale=1):
+                gr.Markdown("### A — Quick glance")
+                glanceA = gr.Label(num_top_classes=10)
+                gr.Markdown("### A — Probabilities (Top‑K)")
+                plotA = gr.Plot()
+                gr.Markdown("### A — Details (Top‑K)")
+                tableA = gr.Dataframe(headers=["Rank", "Label", "Probability", "Logit"], datatype=["number", "str", "number", "number"], row_count=5)
+                gr.Markdown("### A — Notes")
+                notesA = gr.Markdown()
+            with gr.Column(scale=1):
+                gr.Markdown("### B — Quick glance")
+                glanceB = gr.Label(num_top_classes=10)
+                gr.Markdown("### B — Probabilities (Top‑K)")
+                plotB = gr.Plot()
+                gr.Markdown("### B — Details (Top‑K)")
+                tableB = gr.Dataframe(headers=["Rank", "Label", "Probability", "Logit"], datatype=["number", "str", "number", "number"], row_count=5)
+                gr.Markdown("### B — Notes")
+                notesB = gr.Markdown()
+        with gr.Row():
+            gr.Markdown("### Aligned Top‑K Δ (A − B) & Divergence")
+        with gr.Row():
+            deltaPlot = gr.Plot()
+            deltaNotes = gr.Markdown()
+        runAB.click(analyze_pair, [imgA, imgB, topkAB, tempAB], [glanceA, plotA, tableA, notesA, glanceB, plotB, tableB, notesB, deltaPlot, deltaNotes])
 
 if __name__ == "__main__":
     demo.launch()