feat: vis heatmap in --vis (no new arg)

Add get_vit_heatmap() per-patch logits 21x21->336 jet overlay where ViT looks (MoE 5 prefix tokens). --vis now 1x3: image | heatmap overlay | bar (was 1x2). Uses core.head on patch tokens, top class mean. Tested test.png 1.2MB with heatmap, noise 62K. Autocast bf16, fallback to 1x2 if heatmap fails. No new CLI arg, enhanced existing --vis.

plant_cli.py

@@ -73,6 +73,42 @@ def log(msg): print(msg, flush=True)
 def get_transform(img_size=336):
     return T.Compose([T.Resize(int(img_size*1.14)), T.CenterCrop(img_size), T.ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
 
+def get_vit_heatmap(visual_model, image_path, device="cuda", img_size=336):
+    """Feature heatmap where the ViT looks - per-patch logits for top class, 21x21 -> 336."""
+    try:
+        from pathlib import Path as _P
+        im = open_image_pil(image_path)
+        tf = get_transform(img_size)
+        x = tf(im).unsqueeze(0).to(device)
+        core = visual_model.core
+        with torch.no_grad():
+            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device=="cuda"):
+                tokens = core.stem.forward_features(x)  # [1, T, D]
+                if core.input_proj is not None:
+                    tokens = core.input_proj(tokens.float())
+                else:
+                    tokens = tokens.float()
+                for b in core.blocks:
+                    tokens, _, _, _ = b(tokens)
+                tokens = core.exit_norm(tokens)  # [1, T, D]
+                n_prefix = int(getattr(core, "n_prefix", 1))
+                patch_tokens = tokens[:, n_prefix:, :]  # [1, P, D]
+                P = patch_tokens.shape[1]
+                h = w = int(round(P ** 0.5))
+                pt = patch_tokens.squeeze(0)  # [P, D]
+                logits_patch = core.head(pt)  # [P, 999]
+                top_idx = int(logits_patch.mean(0).argmax())
+                heat = logits_patch[:, top_idx]
+                heat = heat - heat.min()
+                heat = heat / (heat.max() - heat.min() + 1e-8)
+                heat_np = heat.float().cpu().numpy().reshape(h, w)
+                heat_img = Image.fromarray((heat_np * 255).astype(np.uint8)).resize((img_size, img_size), Image.BILINEAR)
+                heat_arr = np.array(heat_img).astype(np.float32) / 255.0
+                return heat_arr, top_idx, heat_np
+    except Exception as e:
+        log(f"heatmap failed {e}")
+        return None, None, None
+
 def load_species_map(manifest=MANIFEST):
     # Public release: species_labels.json (999 WA species) - preferred
     if SPECIES_JSON.exists():
@@ -438,25 +474,58 @@ def cmd_identify(args):
         pd.DataFrame(meta["topk"]).to_csv(args.output, index=False)
         print(f"wrote {args.output}")
     if args.vis:
-        # visualize topk bar + image
+        # visualize topk bar + image + feature heatmap (where ViT looks)
         try:
             import matplotlib
             matplotlib.use("Agg")
             import matplotlib.pyplot as plt
             n=len(image_paths)
-            fig, axes = plt.subplots(1, 2, figsize=(14,5), gridspec_kw={"width_ratios":[1,1.2]})
-            ax_img, ax_bar = axes
-            # show first image
-            im=open_image_pil(image_paths[0]).resize((336,336))
-            ax_img.imshow(im)
-            title=f"{meta['topk'][0]['species']}\n{meta['topk'][0]['prob']:.1%} conf {meta['confidence']:.1%} {dt*1000:.0f}ms"
-            if meta["is_unknown"]:
-                title="UNKNOWN\n"+title
-            if len(image_paths)>1: title+=f" {n}-view"
-            ax_img.set_title(title, fontsize=9, color="red" if meta["is_unknown"] else "black")
-            ax_img.axis("off")
-            probs_bar=[r["prob"] for r in meta["topk"]]
-            species=[r["species"] for r in meta["topk"]]
+            # --- generate heatmap for first image (where feature extractor finds features) ---
+            heat_arr = None
+            try:
+                vm, _, _ = load_visual_model(device=device)
+                heat_arr, hm_idx, _ = get_vit_heatmap(vm, image_paths[0], device=device, img_size=args.size)
+            except Exception as e:
+                log(f"heatmap gen failed {e}")
+                heat_arr = None
+            # 1x3 layout: image | heatmap overlay | bar (if heatmap available, else 1x2)
+            if heat_arr is not None:
+                fig, axes = plt.subplots(1, 3, figsize=(18,5), gridspec_kw={"width_ratios":[1,1,1.2]})
+                ax_img, ax_heat, ax_bar = axes
+                # original image
+                im_raw=open_image_pil(image_paths[0]).resize((args.size,args.size))
+                ax_img.imshow(im_raw)
+                title=f"{meta['topk'][0]['species']}\n{meta['topk'][0]['prob']:.1%} conf {meta['confidence']:.1%} {dt*1000:.0f}ms"
+                if meta["is_unknown"]:
+                    title="UNKNOWN\n"+title
+                if len(image_paths)>1: title+=f" {n}-view"
+                ax_img.set_title(title, fontsize=9, color="red" if meta["is_unknown"] else "black")
+                ax_img.axis("off")
+                # heatmap overlay (jet)
+                ax_heat.imshow(im_raw)
+                ax_heat.imshow(heat_arr, cmap="jet", alpha=0.55, vmin=0, vmax=1)
+                ax_heat.set_title(f"Feature heatmap\nTop patch {hm_idx} {heat_arr.max():.2f}", fontsize=9, color="red" if meta["is_unknown"] else "black")
+                ax_heat.axis("off")
+            else:
+                fig, axes = plt.subplots(1, 2, figsize=(14,5), gridspec_kw={"width_ratios":[1,1.2]})
+                ax_img, ax_bar = axes
+                im=open_image_pil(image_paths[0]).resize((336,336))
+                ax_img.imshow(im)
+                title=f"{meta['topk'][0]['species']}\n{meta['topk'][0]['prob']:.1%} conf {meta['confidence']:.1%} {dt*1000:.0f}ms"
+                if meta["is_unknown"]:
+                    title="UNKNOWN\n"+title
+                if len(image_paths)>1: title+=f" {n}-view"
+                ax_img.set_title(title, fontsize=9, color="red" if meta["is_unknown"] else "black")
+                ax_img.axis("off")
+                # heat axis not present, will reuse ax_bar below
+            # bar chart (common)
+            if heat_arr is not None:
+                probs_bar=[r["prob"] for r in meta["topk"]]
+                species=[r["species"] for r in meta["topk"]]
+            else:
+                probs_bar=[r["prob"] for r in meta["topk"]]
+                species=[r["species"] for r in meta["topk"]]
+            # ax_bar is defined in both branches (for heat case it's third axis)
             colors=["red" if meta["is_unknown"] and i==0 else "green" if i==0 else "steelblue" for i in range(len(probs_bar))]
             bars=ax_bar.barh(range(len(probs_bar))[::-1], probs_bar[::-1], color=colors[::-1])
             ax_bar.set_yticks(range(len(probs_bar))[::-1])
@@ -465,13 +534,13 @@ def cmd_identify(args):
             ax_bar.set_title(f"WA 999 spp Top-{args.topk} {'UNKNOWN' if meta['is_unknown'] else ''}", fontsize=10, color="red" if meta["is_unknown"] else "black")
             for p, bar in zip(probs_bar[::-1], bars):
                 ax_bar.text(p+0.01, bar.get_y()+bar.get_height()/2, f"{p:.1%}", va="center", fontsize=7)
-            # add unknown text if unknown
             if meta["is_unknown"]:
                 ax_bar.text(0.5, -1, f"sim {meta['unknown_info']['sim_max']:.2f} < {meta['thr_cos']:.2f}", ha="center", fontsize=8, color="red", transform=ax_bar.transAxes)
             plt.tight_layout()
             plt.savefig(args.vis, dpi=150)
-            print(f"wrote {args.vis}")
+            print(f"wrote {args.vis} {'with heatmap' if heat_arr is not None else ''}")
         except Exception as e:
+            import traceback; traceback.print_exc()
             log(f"vis failed {e}")
 
 def cmd_batch(args):