Deploy MidasMap Gradio app; weights downloaded from model repo at runtime

README.md

@@ -12,27 +12,55 @@ license: mit
 
 # MidasMap Space
 
-This folder is a **template** for creating a [Hugging Face Space](https://huggingface.co/docs/hub/spaces-overview).
+Gradio demo for **[MidasMap](https://github.com/AnikS22/MidasMap)** (immunogold particle detection in TEM synapse images).
 
-**Why not Vercel for the model?** Vercel serverless functions have strict size and time limits; they are not meant for PyTorch + a ~100MB checkpoint and multi-second GPU/CPU inference. **Host the Gradio app + weights on a Space** (CPU free tier or GPU upgrade).
+## Deploy from your laptop
 
-## Create the Space
+From the **MidasMap** repo root:
 
-1. On Hugging Face: **New Space** → SDK **Gradio** → name e.g. `MidasMap`.
-2. Clone the Space repo locally, or connect **GitHub** and set the Space root to this monorepo with **App file** pointing to the copied `app.py`.
-3. Copy into the Space repository root:
-   - `app.py` from the main MidasMap repo (project root), **or** symlink / duplicate.
-   - `src/` (entire package)
-   - `requirements-space.txt` from this folder as **`requirements.txt`**
-4. In Space **Settings → Repository secrets** (if needed): none required for public weights.
-5. Ensure `checkpoints/final/final_model.pth` is present:
-   - Upload via **Files** tab, or
-   - Add a startup script to download from `AnikS22/MidasMap` on the Hub (see HF docs for `hf_hub_download`).
+```bash
+export HF_TOKEN=hf_...   # write token
+# Recommended: do not upload the ~100MB checkpoint into the Space (avoids LFS / size issues).
+export HF_SPACE_SKIP_CHECKPOINT=1
+./scripts/upload_hf_space.sh
+```
 
-After the Space builds, point your **Vercel** site (`vercel-site`) at it:
+If **`upload_hf_space.sh` fails**, use **git + LFS** instead (often more reliable):
 
-`https://yoursite.vercel.app/?embed=https://huggingface.co/spaces/YOUR_USER/YOUR_SPACE`
+```bash
+brew install git-lfs && git lfs install   # once
+export HF_TOKEN=hf_...
+./scripts/push_hf_space_git.sh
+```
 
----
+Full options: [docs/DEPLOY.md](../docs/DEPLOY.md) in the main repo.
+
+Create the Space once if needed (Gradio SDK required for auto-create):
+
+```bash
+huggingface-cli repo create MidasMap --type space --space_sdk gradio -y
+```
+
+Weights are loaded from the **model** repo `AnikS22/MidasMap` at `checkpoints/final/final_model.pth` when the file is not in the Space. Override with Space secrets / env: `MIDASMAP_HF_WEIGHTS_REPO`, `MIDASMAP_HF_WEIGHTS_FILE`.
+
+To bundle the checkpoint in the Space instead (larger upload):
 
-Gradio app and model logic: [github.com/AnikS22/MidasMap](https://github.com/AnikS22/MidasMap)
+```bash
+export HF_SPACE_SKIP_CHECKPOINT=0
+./scripts/upload_hf_space.sh
+```
+
+## Troubleshooting uploads
+
+| Symptom | What to do |
+|--------|----------------|
+| **401 / not logged in** | `export HF_TOKEN=hf_...` with a token that has **write** access, or `huggingface-cli login`. |
+| **LFS / authorization / upload stuck** | Use `HF_SPACE_SKIP_CHECKPOINT=1` so only code uploads; ensure the **model** repo (not the Space) contains `checkpoints/final/final_model.pth`. |
+| **Space does not exist** | Create it in the HF web UI (**New Space** → **Gradio**) or run `huggingface-cli repo create ... --type space --space_sdk gradio`. |
+| **“No space_sdk provided”** | The Space repo must be created as **Gradio** (or pass `--space_sdk gradio` when using `repo create`). |
+| **Model not found on Space** | First boot downloads weights from the Hub; public repos need no token. Private model repo: add `HF_TOKEN` as a Space **secret** (read). |
+| **Still failing** | Try `pip install hf_transfer` and `export HF_HUB_ENABLE_HF_TRANSFER=1` before upload. Or use **git** + **git lfs** clone of the Space, copy files, commit, push. |
+
+## Vercel embed
+
+`https://yoursite.vercel.app/?embed=https://huggingface.co/spaces/YOUR_USER/YOUR_SPACE`

app.py

@@ -91,7 +91,11 @@ def load_model(checkpoint_path: str):
         if torch.backends.mps.is_available()
         else "cpu"
     )
-    MODEL = ImmunogoldCenterNet(bifpn_channels=128, bifpn_rounds=2)
+    MODEL = ImmunogoldCenterNet(
+        bifpn_channels=128,
+        bifpn_rounds=2,
+        imagenet_encoder_fallback=False,
+    )
     ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
     MODEL.load_state_dict(ckpt["model_state_dict"])
     MODEL.to(DEVICE)
@@ -176,6 +180,94 @@ def _df_to_preview_html(df: pd.DataFrame) -> str:
     )
 
 
+def _numpy_image_to_uint8_rgb(img: np.ndarray) -> np.ndarray:
+    """Normalize various arrays to HxWx3 uint8 for cropping / display."""
+    if img is None:
+        return None
+    arr = np.asarray(img)
+    if arr.size == 0:
+        return None
+    if arr.ndim == 2:
+        arr = np.stack([arr, arr, arr], axis=-1)
+    elif arr.ndim == 3 and arr.shape[2] == 4:
+        arr = arr[:, :, :3]
+    if arr.dtype in (np.float32, np.float64):
+        mx = float(arr.max()) if arr.size else 1.0
+        if mx <= 1.0:
+            arr = (np.clip(arr, 0, 1) * 255.0).astype(np.uint8)
+        else:
+            arr = np.clip(arr, 0, 255).astype(np.uint8)
+    else:
+        arr = np.clip(arr, 0, 255).astype(np.uint8)
+    return arr
+
+
+def magnifier_zoom(
+    store: dict,
+    view: str,
+    center_x_pct: float,
+    center_y_pct: float,
+    zoom: float,
+    output_px: int,
+) -> np.ndarray | None:
+    """
+    Crop a square region around (center_x_pct, center_y_pct) and upscale for a loupe view.
+    zoom: 1 = see ~full width in loupe; larger = stronger magnification (smaller crop).
+    """
+    if not store or not isinstance(store, dict):
+        return None
+    key = {"Overlay": "overlay", "Heatmaps": "heatmap", "Summary": "stats"}.get(view, "overlay")
+    img = _numpy_image_to_uint8_rgb(store.get(key))
+    if img is None:
+        return None
+    h, w = img.shape[:2]
+    cx = int(np.clip(center_x_pct / 100.0 * (w - 1), 0, w - 1))
+    cy = int(np.clip(center_y_pct / 100.0 * (h - 1), 0, h - 1))
+    z = max(1.0, float(zoom))
+    half_w = max(1, int(w / (2.0 * z)))
+    half_h = max(1, int(h / (2.0 * z)))
+    x0, x1 = max(0, cx - half_w), min(w, cx + half_w)
+    y0, y1 = max(0, cy - half_h), min(h, cy + half_h)
+    if x1 <= x0 or y1 <= y0:
+        crop = img
+    else:
+        crop = img[y0:y1, x0:x1]
+    side = int(np.clip(output_px, 256, 1024))
+    try:
+        from PIL import Image as PILImage
+
+        pil = PILImage.fromarray(crop)
+        pil = pil.resize((side, side), PILImage.Resampling.LANCZOS)
+        return np.asarray(pil)
+    except Exception:
+        from skimage.transform import resize
+
+        up = resize(crop, (side, side), order=1, preserve_range=True)
+        return np.clip(up, 0, 255).astype(np.uint8)
+
+
+def run_detection(
+    image_file,
+    conf_threshold: float,
+    nms_6nm: int,
+    nms_12nm: int,
+    px_per_um: float,
+    progress=gr.Progress(track_tqdm=False),
+):
+    """Run model and return outputs plus viz state for the magnifier."""
+    out = detect_particles(
+        image_file,
+        conf_threshold,
+        nms_6nm,
+        nms_12nm,
+        px_per_um,
+        progress=progress,
+    )
+    overlay, hm, stats, csvp, table, summary = out
+    store = {"overlay": overlay, "heatmap": hm, "stats": stats}
+    return overlay, hm, stats, csvp, table, summary, store
+
+
 def detect_particles(
     image_file,
     conf_threshold: float = 0.25,
@@ -250,8 +342,27 @@ def detect_particles(
 
     from skimage.transform import resize
 
-    hm6_up = resize(hm_np[0], (h, w), order=1)
-    hm12_up = resize(hm_np[1], (h, w), order=1)
+    hm6_up = np.clip(
+        np.nan_to_num(resize(hm_np[0], (h, w), order=1), nan=0.0),
+        0.0,
+        1.0,
+    )
+    hm12_up = np.clip(
+        np.nan_to_num(resize(hm_np[1], (h, w), order=1), nan=0.0),
+        0.0,
+        1.0,
+    )
+
+    def _heatmap_vmax(hm: np.ndarray) -> float:
+        """Stable color scale: avoid invisible overlays when max is tiny or flat."""
+        flat = hm.ravel()
+        if flat.size == 0:
+            return 0.3
+        mx = float(np.max(flat))
+        if mx < 1e-6:
+            return 0.3
+        p99 = float(np.percentile(flat, 99.0))
+        return float(np.clip(max(0.12, p99 * 1.05, mx * 0.95), 0.05, 1.0))
 
     # --- Overlay (publication-style legend + scale bar) ---
     fig_overlay, ax = plt.subplots(figsize=(11, 11))
@@ -291,21 +402,54 @@ def detect_particles(
     overlay_img = np.asarray(fig_overlay.canvas.renderer.buffer_rgba())[:, :, :3]
     plt.close(fig_overlay)
 
-    # --- Heatmaps ---
-    fig_hm, axes = plt.subplots(1, 2, figsize=(14, 6.2))
-    axes[0].imshow(img, cmap="gray", aspect="equal")
-    axes[0].imshow(hm6_up, cmap="magma", alpha=0.55, vmin=0, vmax=max(0.3, float(hm6_up.max())))
-    axes[0].set_title(f"AMPA (6 nm) channel · n = {n_6nm}", fontsize=11)
-    axes[0].axis("off")
-
-    axes[1].imshow(img, cmap="gray", aspect="equal")
-    axes[1].imshow(hm12_up, cmap="inferno", alpha=0.55, vmin=0, vmax=max(0.3, float(hm12_up.max())))
-    axes[1].set_title(f"NR1 (12 nm) channel · n = {n_12nm}", fontsize=11)
-    axes[1].axis("off")
+    # --- Heatmaps: row1 = overlay on EM; row2 = model heat only (debug-friendly) ---
+    # Training uses Gaussian GT; inference heatmaps are learned sigmoid blobs, not analytic Gaussians.
+    v6, v12 = _heatmap_vmax(hm6_up), _heatmap_vmax(hm12_up)
+    fig_hm, axes = plt.subplots(2, 2, figsize=(14, 12))
+    ax00, ax01 = axes[0]
+    ax10, ax11 = axes[1]
+
+    for ax, hm, v, cmap, title in (
+        (ax00, hm6_up, v6, "magma", f"AMPA overlay · n={n_6nm} · vmax={v6:.2f}"),
+        (ax01, hm12_up, v12, "inferno", f"NR1 overlay · n={n_12nm} · vmax={v12:.2f}"),
+    ):
+        ax.imshow(img, cmap="gray", aspect="equal", interpolation="nearest")
+        ax.imshow(
+            hm,
+            cmap=cmap,
+            alpha=0.6,
+            vmin=0.0,
+            vmax=v,
+            interpolation="bilinear",
+        )
+        ax.set_title(title, fontsize=10)
+        ax.axis("off")
+
+    ax10.imshow(hm6_up, cmap="magma", vmin=0.0, vmax=v6, interpolation="nearest")
+    ax10.set_title(f"AMPA heatmap only · max={float(np.max(hm6_up)):.4f}", fontsize=10)
+    ax10.axis("off")
+
+    ax11.imshow(hm12_up, cmap="inferno", vmin=0.0, vmax=v12, interpolation="nearest")
+    ax11.set_title(f"NR1 heatmap only · max={float(np.max(hm12_up)):.4f}", fontsize=10)
+    ax11.axis("off")
+
     plt.tight_layout()
-    fig_hm.canvas.draw()
-    heatmap_img = np.asarray(fig_hm.canvas.renderer.buffer_rgba())[:, :, :3]
+    # PNG raster → uint8 RGB (reliable in Gradio vs raw canvas buffer on some setups)
+    from io import BytesIO
+
+    _buf = BytesIO()
+    fig_hm.savefig(_buf, format="png", dpi=120, bbox_inches="tight", facecolor="white")
     plt.close(fig_hm)
+    _buf.seek(0)
+    try:
+        from PIL import Image as _PILImage
+
+        heatmap_img = np.asarray(_PILImage.open(_buf).convert("RGB"))
+    except Exception:
+        import matplotlib.image as _mimg
+
+        _buf.seek(0)
+        heatmap_img = (_mimg.imread(_buf)[:, :, :3] * 255.0).clip(0, 255).astype(np.uint8)
 
     # --- Stats (µm where helpful) ---
     fig_stats, axes = plt.subplots(1, 3, figsize=(16, 4.8))
@@ -411,58 +555,73 @@ def detect_particles(
 
 
 MM_CSS = """
-.gradio-container { max-width: 1320px !important; margin: auto !important; }
+@import url("https://fonts.googleapis.com/css2?family=Libre+Baskerville:wght@700&family=Source+Sans+3:wght@400;600;700&display=swap");
+.gradio-container { max-width: 1280px !important; margin: auto !important; padding: 1rem 0.75rem 2rem !important; }
 .mm-brand-bar {
   display: flex; align-items: center; justify-content: space-between;
-  flex-wrap: wrap; gap: 0.75rem;
-  padding: 0.6rem 0 1.25rem;
-  border-bottom: 1px solid var(--border-color-primary);
-  margin-bottom: 1.25rem;
+  flex-wrap: wrap; gap: 0.5rem 1rem;
+  padding: 0 0 1rem;
+  margin-bottom: 1rem;
+  border-bottom: 1px solid rgba(148, 163, 184, 0.2);
 }
 .mm-brand-bar span {
-  font-size: 0.72rem; letter-spacing: 0.14em; text-transform: uppercase;
-  color: var(--body-text-color-subdued); font-weight: 600;
+  font-size: 0.7rem; letter-spacing: 0.06em;
+  color: var(--body-text-color-subdued); font-weight: 500;
 }
 .mm-hero {
-  padding: 1.5rem 1.35rem 1.35rem;
-  margin-bottom: 0.25rem;
-  border-radius: 10px;
-  background: linear-gradient(145deg, #0c4a6e22 0%, #0f172a 48%, #1e1b4b33 100%);
-  border: 1px solid #33415588;
+  padding: 1.35rem 1.5rem;
+  margin-bottom: 1.25rem;
+  border-radius: 16px;
+  background: linear-gradient(155deg, rgba(13, 148, 136, 0.12) 0%, rgba(15, 23, 42, 0.95) 42%, rgba(30, 27, 75, 0.15) 100%);
+  border: 1px solid rgba(148, 163, 184, 0.15);
+  box-shadow: 0 8px 32px rgba(0, 0, 0, 0.2);
 }
 .mm-hero h1 {
   font-family: "Libre Baskerville", Georgia, serif;
   font-weight: 700;
   letter-spacing: -0.02em;
-  margin: 0 0 0.4rem 0;
-  font-size: 1.65rem;
-  color: #f1f5f9;
+  margin: 0 0 0.5rem 0;
+  font-size: 1.75rem;
+  color: #f8fafc;
 }
 .mm-hero .mm-sub {
-  margin: 0 0 0.85rem 0;
-  color: #94a3b8;
-  font-size: 0.92rem;
-  line-height: 1.55;
-  max-width: 58ch;
+  margin: 0 0 1rem 0;
+  color: #cbd5e1;
+  font-size: 0.95rem;
+  line-height: 1.6;
+  max-width: 62ch;
 }
-.mm-badge-row { display: flex; flex-wrap: wrap; gap: 0.4rem; }
+.mm-badge-row { display: flex; flex-wrap: wrap; gap: 0.45rem; }
 .mm-badge {
-  font-size: 0.65rem; text-transform: uppercase; letter-spacing: 0.07em;
-  padding: 0.2rem 0.5rem; border-radius: 4px;
-  background: #0e749033; color: #99f6e4; border: 1px solid #14b8a644;
+  font-size: 0.62rem; letter-spacing: 0.05em; font-weight: 600;
+  padding: 0.28rem 0.55rem; border-radius: 999px;
+  background: rgba(45, 212, 191, 0.12); color: #5eead4;
+  border: 1px solid rgba(45, 212, 191, 0.25);
 }
-.mm-layout { display: flex; gap: 1.25rem; align-items: flex-start; flex-wrap: wrap; }
+.mm-layout { display: flex; gap: 1.5rem; align-items: flex-start; flex-wrap: wrap; }
 .mm-sidebar {
-  flex: 1 1 280px; max-width: 340px;
-  padding: 1rem 1.1rem; border-radius: 10px;
-  border: 1px solid var(--border-color-primary);
+  flex: 1 1 300px; max-width: 360px;
+  padding: 1.25rem 1.35rem; border-radius: 16px;
+  border: 1px solid rgba(148, 163, 184, 0.12);
   background: var(--block-background-fill);
+  box-shadow: 0 4px 24px rgba(0, 0, 0, 0.12);
+}
+.mm-main {
+  flex: 1 1 480px; min-width: 0;
+  padding: 0.25rem 0.15rem;
+  border-radius: 16px;
 }
-.mm-main { flex: 3 1 520px; min-width: 0; }
 .mm-panel-title {
-  font-size: 0.7rem; text-transform: uppercase; letter-spacing: 0.1em;
-  color: var(--body-text-color-subdued); font-weight: 600; margin: 0 0 0.65rem 0;
+  font-size: 0.72rem; text-transform: uppercase; letter-spacing: 0.08em;
+  color: var(--body-text-color-subdued); font-weight: 600; margin: 0 0 0.75rem 0;
 }
+.mm-loupe-help {
+  font-size: 0.82rem; line-height: 1.45; color: var(--body-text-color-subdued);
+  margin: 0 0 0.75rem 0; padding: 0.65rem 0.85rem;
+  border-radius: 10px; background: rgba(30, 41, 59, 0.45);
+  border: 1px solid rgba(148, 163, 184, 0.12);
+}
+.tabs > .tab-nav button { font-weight: 500 !important; letter-spacing: 0.01em; }
 .mm-callout {
   margin: 0; padding: 0.75rem 0.9rem; border-radius: 8px;
   background: #1e293b66; border: 1px solid var(--border-color-primary);
@@ -513,21 +672,10 @@ table.mm-table td { padding: 0.35rem 0.5rem; border-bottom: 1px solid #33415544;
 
 
 def build_app():
+    # Use named hues only (no custom Color dicts): avoids Gradio/Jinja template bugs on some stacks (e.g. HF Spaces + Py3.13).
     theme = gr.themes.Soft(
-        primary_hue=gr.themes.Color(
-            c50="#f0fdfa",
-            c100="#ccfbf1",
-            c200="#99f6e4",
-            c300="#5eead4",
-            c400="#2dd4bf",
-            c500="#14b8a6",
-            c600="#0d9488",
-            c700="#0f766e",
-            c800="#115e59",
-            c900="#134e4a",
-            c950="#042f2e",
-        ),
-        neutral_hue=gr.themes.colors.slate,
+        primary_hue="teal",
+        neutral_hue="slate",
         font=("Source Sans 3", "ui-sans-serif", "system-ui", "sans-serif"),
         font_mono=("IBM Plex Mono", "ui-monospace", "monospace"),
     ).set(
@@ -539,51 +687,47 @@ def build_app():
         block_label_text_size="*text_sm",
     )
 
-    head = """
-    <link href="https://fonts.googleapis.com/css2?family=Libre+Baskerville:wght@700&family=Source+Sans+3:wght@400;600;700&display=swap" rel="stylesheet">
-    """
-
     with gr.Blocks(
         title="MidasMap — Immunogold analysis",
         theme=theme,
         css=MM_CSS,
-        head=head,
     ) as app:
         gr.HTML(
             """
             <div class="mm-brand-bar">
-              <span>Quantitative EM · synapse immunogold</span>
-              <span>Research use · validate critical counts manually</span>
+              <span>MidasMap · immunogold on TEM synapses</span>
+              <span>For research — verify important counts by eye</span>
             </div>
             <div class="mm-hero">
               <h1>MidasMap</h1>
               <p class="mm-sub">
-                Automated particle picking for <strong>freeze-fracture replica immunolabeling (FFRIL)</strong> TEM:
-                <strong>6 nm</strong> gold (AMPA receptors) and <strong>12 nm</strong> gold (NR1 / NMDA receptors).
-                Coordinates export in <strong>µm</strong> for comparison to physiology and super-resolution data—set calibration to match your microscope.
+                Find <strong>6 nm</strong> (AMPA) and <strong>12 nm</strong> (NR1) gold particles in
+                <strong>FFRIL</strong> micrographs. Set <strong>calibration</strong> so exports are in µm.
+                Use the <strong>magnifying glass</strong> below to inspect beads and heatmaps up close.
               </p>
               <div class="mm-badge-row">
-                <span class="mm-badge">FFRIL / TEM</span>
+                <span class="mm-badge">FFRIL</span>
                 <span class="mm-badge">CenterNet</span>
-                <span class="mm-badge">CEM500K backbone</span>
-                <span class="mm-badge">LOOCV F1 ≈ 0.94</span>
+                <span class="mm-badge">CEM500K</span>
+                <span class="mm-badge">F1 ≈ 0.94 LOOCV</span>
               </div>
             </div>
             """
         )
 
+        viz_state = gr.State({"overlay": None, "heatmap": None, "stats": None})
+
         with gr.Row(elem_classes=["mm-layout"]):
             with gr.Column(elem_classes=["mm-sidebar"]):
-                gr.HTML('<p class="mm-panel-title">Micrograph & calibration</p>')
+                gr.HTML('<p class="mm-panel-title">1 · Upload & settings</p>')
                 image_input = gr.File(
-                    label="Upload image",
+                    label="Micrograph",
                     file_types=[".tif", ".tiff", ".png", ".jpg", ".jpeg"],
                 )
                 px_per_um_in = gr.Number(
                     value=DEFAULT_PX_PER_UM,
-                    label="Calibration (pixels per µm)",
-                    info=f"Default {DEFAULT_PX_PER_UM:.0f} matches the published training set. "
-                    "Update if your acquisition scale differs.",
+                    label="Pixels per µm",
+                    info=f"Default {DEFAULT_PX_PER_UM:.0f} matches the training corpus. Change if your scale differs.",
                     minimum=1,
                     maximum=1e6,
                 )
@@ -592,28 +736,64 @@ def build_app():
                     maximum=0.95,
                     value=0.25,
                     step=0.05,
-                    label="Confidence threshold",
-                    info="Higher → fewer, sharper peaks. Lower → recall with more false positives.",
+                    label="Confidence",
+                    info="Higher = stricter (fewer hits). Lower = more sensitive.",
                 )
-                with gr.Accordion("Advanced · non-max suppression", open=False):
+                with gr.Accordion("Advanced · peak spacing (NMS)", open=False):
                     nms_6nm = gr.Slider(
                         minimum=1,
                         maximum=9,
                         value=3,
                         step=2,
-                        label="NMS · 6 nm channel",
-                        info="Minimum spacing between AMPA peaks on the heatmap grid.",
+                        label="Spacing · 6 nm channel",
+                        info="Minimum gap between AMPA peaks on the model grid.",
                     )
                     nms_12nm = gr.Slider(
                         minimum=1,
                         maximum=9,
                         value=5,
                         step=2,
-                        label="NMS · 12 nm channel",
+                        label="Spacing · 12 nm channel",
                     )
                 detect_btn = gr.Button("Run detection", variant="primary", size="lg")
 
-                with gr.Accordion("For neuroscientists — interpretation", open=False):
+                with gr.Accordion("Magnifying glass", open=True):
+                    gr.HTML(
+                        """<p class="mm-loupe-help" style="margin-top:0">
+                        After you run detection, pick which result to inspect and adjust the sliders.
+                        <strong>Magnification</strong> zooms in (smaller crop, upscaled). Use the fullscreen icon on any image for a larger view.
+                        </p>"""
+                    )
+                    mag_view = gr.Radio(
+                        choices=["Overlay", "Heatmaps", "Summary"],
+                        value="Overlay",
+                        label="Source image",
+                    )
+                    mag_cx = gr.Slider(
+                        0, 100, value=50, step=0.5,
+                        label="Pan left ↔ right (%)",
+                    )
+                    mag_cy = gr.Slider(
+                        0, 100, value=50, step=0.5,
+                        label="Pan up ↔ down (%)",
+                    )
+                    mag_zoom = gr.Slider(
+                        1, 10, value=2.5, step=0.25,
+                        label="Magnification",
+                        info="Higher = stronger zoom (smaller region).",
+                    )
+                    mag_out = gr.Slider(
+                        256, 768, value=480, step=64,
+                        label="Loupe window (px)",
+                    )
+                    mag_out_img = gr.Image(
+                        label="Loupe preview",
+                        type="numpy",
+                        height=380,
+                        show_fullscreen_button=True,
+                    )
+
+                with gr.Accordion("Notes for scientists", open=False):
                     gr.Markdown(
                         """
 #### What the model outputs
@@ -634,8 +814,9 @@ Sahai, A. (2026). *MidasMap* (software). https://github.com/AnikS22/MidasMap
                     )
 
             with gr.Column(elem_classes=["mm-main"]):
+                gr.HTML('<p class="mm-panel-title">2 · Results</p>')
                 summary_md = gr.HTML(
-                    value="<p class='mm-callout'>Upload a synapse micrograph to begin. Adjust calibration before export if your scale differs from the default.</p>"
+                    value="<p class='mm-callout'>Upload a micrograph and tap <strong>Run detection</strong>. Set pixels/µm before exporting if your scale differs.</p>"
                 )
                 with gr.Tabs():
                     with gr.Tab("Overlay"):
@@ -643,18 +824,21 @@ Sahai, A. (2026). *MidasMap* (software). https://github.com/AnikS22/MidasMap
                             label="Detections + scale bar",
                             type="numpy",
                             height=540,
+                            show_fullscreen_button=True,
                         )
                     with gr.Tab("Heatmaps"):
                         heatmap_output = gr.Image(
                             label="Class-specific maps",
                             type="numpy",
                             height=540,
+                            show_fullscreen_button=True,
                         )
-                    with gr.Tab("Quant summary"):
+                    with gr.Tab("Summary"):
                         stats_output = gr.Image(
-                            label="Distributions & table",
+                            label="Counts & distributions",
                             type="numpy",
                             height=440,
+                            show_fullscreen_button=True,
                         )
                     with gr.Tab("Table & export"):
                         table_output = gr.HTML(
@@ -674,8 +858,10 @@ Sahai, A. (2026). *MidasMap* (software). https://github.com/AnikS22/MidasMap
             """
         )
 
+        mag_inputs = [viz_state, mag_view, mag_cx, mag_cy, mag_zoom, mag_out]
+
         detect_btn.click(
-            fn=detect_particles,
+            fn=run_detection,
             inputs=[image_input, conf_slider, nms_6nm, nms_12nm, px_per_um_in],
             outputs=[
                 overlay_output,
@@ -684,12 +870,52 @@ Sahai, A. (2026). *MidasMap* (software). https://github.com/AnikS22/MidasMap
                 csv_output,
                 table_output,
                 summary_md,
+                viz_state,
             ],
-        )
+        ).then(magnifier_zoom, mag_inputs, mag_out_img)
+
+        for _ctrl in (mag_view, mag_cx, mag_cy, mag_zoom, mag_out):
+            _ctrl.change(magnifier_zoom, mag_inputs, mag_out_img)
 
     return app
 
 
+def _running_on_hf_space() -> bool:
+    """Hugging Face Spaces injects these env vars; Gradio must bind 0.0.0.0 and never use share=True."""
+    return bool(
+        os.environ.get("SPACE_REPO_NAME")
+        or os.environ.get("SPACE_AUTHOR_NAME")
+        or os.environ.get("SPACE_ID")
+    )
+
+
+def _resolve_checkpoint(ckpt: Path) -> Path:
+    """Use local .pth if present; on HF Space fetch from the Hub model repo if missing (smaller Space uploads)."""
+    if ckpt.is_file():
+        return ckpt
+    if _running_on_hf_space():
+        try:
+            from huggingface_hub import hf_hub_download
+        except ImportError as e:
+            raise SystemExit(
+                "huggingface_hub is required on the Space to download weights. "
+                "Add it to requirements.txt or bundle checkpoints/final/final_model.pth in the Space."
+            ) from e
+        repo_id = os.environ.get("MIDASMAP_HF_WEIGHTS_REPO", "AnikS22/MidasMap").strip()
+        filename = os.environ.get(
+            "MIDASMAP_HF_WEIGHTS_FILE", "checkpoints/final/final_model.pth"
+        ).strip()
+        print(f"Checkpoint not found at {ckpt}; downloading {filename} from model repo {repo_id} ...")
+        cached = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")
+        return Path(cached)
+    raise SystemExit(
+        f"Checkpoint not found: {ckpt}\n"
+        "Train with train_final.py or download from Hugging Face:\n"
+        "  huggingface-cli download AnikS22/MidasMap checkpoints/final/final_model.pth "
+        "--local-dir . --repo-type model"
+    )
+
+
 def main():
     parser = argparse.ArgumentParser(description="MidasMap web dashboard")
     parser.add_argument(
@@ -712,39 +938,24 @@ def main():
     if os.environ.get("GRADIO_SHARE", "").lower() in ("1", "true", "yes"):
         args.share = True
 
-    ckpt = Path(args.checkpoint)
-    if not ckpt.is_file():
-        raise SystemExit(
-            f"Checkpoint not found: {ckpt}\n"
-            "Train with train_final.py or download from Hugging Face:\n"
-            "  huggingface-cli download AnikS22/MidasMap checkpoints/final/final_model.pth "
-            "--local-dir ."
-        )
+    if _running_on_hf_space():
+        args.share = False
+        if not args.server_name:
+            args.server_name = "0.0.0.0"
+
+    ckpt = _resolve_checkpoint(Path(args.checkpoint))
 
     load_model(str(ckpt))
     demo = build_app()
+    port = int(os.environ.get("GRADIO_SERVER_PORT", os.environ.get("PORT", str(args.port))))
     launch_kw = dict(
         share=args.share,
-        server_port=args.port,
+        server_port=port,
         server_name=args.server_name,
         show_api=False,
         inbrowser=False,
     )
-    try:
-        demo.launch(**launch_kw)
-    except ValueError as err:
-        if (
-            "localhost is not accessible" in str(err)
-            and not launch_kw.get("share")
-            and os.environ.get("GRADIO_SHARE", "").lower() not in ("1", "true", "yes")
-        ):
-            print(
-                "Localhost check failed in this environment; starting with share=True "
-                "(Gradio tunnel). Use --share next time, or set GRADIO_SHARE=1."
-            )
-            build_app().launch(**{**launch_kw, "share": True})
-        else:
-            raise
+    demo.launch(**launch_kw)
 
 
 if __name__ == "__main__":

requirements.txt

@@ -14,5 +14,7 @@ PyYAML>=6.0
 albumentations>=1.3.0
 opencv-python-headless>=4.7.0
 gradio==4.44.1
+# Avoid Jinja2 3.2+ cache key issues with some Gradio/Starlette stacks on HF Spaces.
+jinja2>=3.1.0,<3.2.0
 huggingface_hub>=0.20.0,<0.25.0
 tqdm>=4.65.0

src/ensemble.py

@@ -163,6 +163,22 @@ def ensemble_predict(
     return np.mean(all_heatmaps, axis=0), np.mean(all_offsets, axis=0)
 
 
+def _tile_origins(axis_len: int, patch: int, stride_step: int) -> list:
+    """
+    Starting indices for sliding windows along one axis so the last window
+    flush-aligns with the far edge. Plain range(0, n-patch+1, step) misses
+    the bottom/right of most image sizes (e.g. 2048 with patch 512, step 384),
+    leaving heatmap strips at zero.
+    """
+    if axis_len <= patch:
+        return [0]
+    last = axis_len - patch
+    starts = list(range(0, last + 1, stride_step))
+    if starts[-1] != last:
+        starts.append(last)
+    return starts
+
+
 def sliding_window_inference(
     model: ImmunogoldCenterNet,
     image: np.ndarray,
@@ -188,10 +204,18 @@ def sliding_window_inference(
         offsets: (2, H/2, W/2) numpy array
     """
     model.eval()
+    orig_h, orig_w = image.shape[:2]
+    # Pad bottom/right so each dim >= patch_size; otherwise range() for tiles is empty
+    # and heatmaps stay all zeros (looks like a "broken" heatmap in the UI).
+    pad_h = max(0, patch_size - orig_h)
+    pad_w = max(0, patch_size - orig_w)
+    if pad_h > 0 or pad_w > 0:
+        image = np.pad(image, ((0, pad_h), (0, pad_w)), mode="reflect")
+
     h, w = image.shape[:2]
     stride_step = patch_size - overlap
 
-    # Output dimensions at model stride
+    # Output dimensions at model stride (padded image)
     out_h = h // 2
     out_w = w // 2
     out_patch = patch_size // 2
@@ -200,8 +224,8 @@ def sliding_window_inference(
     offsets = np.zeros((2, out_h, out_w), dtype=np.float32)
     count = np.zeros((out_h, out_w), dtype=np.float32)
 
-    for y0 in range(0, h - patch_size + 1, stride_step):
-        for x0 in range(0, w - patch_size + 1, stride_step):
+    for y0 in _tile_origins(h, patch_size, stride_step):
+        for x0 in _tile_origins(w, patch_size, stride_step):
             patch = image[y0 : y0 + patch_size, x0 : x0 + patch_size]
             tensor = (
                 torch.from_numpy(patch)
@@ -233,4 +257,9 @@ def sliding_window_inference(
     count = np.maximum(count, 1)
     offsets /= count[np.newaxis, :, :]
 
+    # Crop back to original (pre-pad) spatial extent in heatmap space
+    crop_h, crop_w = orig_h // 2, orig_w // 2
+    heatmap = heatmap[:, :crop_h, :crop_w]
+    offsets = offsets[:, :crop_h, :crop_w]
+
     return heatmap, offsets

src/model.py

@@ -215,6 +215,7 @@ class ImmunogoldCenterNet(nn.Module):
         bifpn_channels: int = 128,
         bifpn_rounds: int = 2,
         num_classes: int = 2,
+        imagenet_encoder_fallback: bool = True,
     ):
         super().__init__()
         self.num_classes = num_classes
@@ -229,13 +230,14 @@ class ImmunogoldCenterNet(nn.Module):
         # Load pretrained weights
         if pretrained_path:
             self._load_pretrained(backbone, pretrained_path)
-        else:
-            # Use ImageNet weights as fallback, adapting conv1
+        elif imagenet_encoder_fallback:
+            # Training: better init when CEM500K path is missing (downloads ~100MB).
             imagenet_backbone = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
             state = imagenet_backbone.state_dict()
             # Mean-pool RGB conv1 weights → grayscale
             state["conv1.weight"] = state["conv1.weight"].mean(dim=1, keepdim=True)
             backbone.load_state_dict(state, strict=False)
+        # else: random encoder init — use when loading a full checkpoint immediately (Gradio, predict).
 
         # Extract encoder stages
         self.stem = nn.Sequential(