Upload 317 files

README.md

@@ -15,16 +15,16 @@ Iman Sabir Ezzat, Randa K Ismail, Ayden Chavez, Marisa Zallocchi, PhD, Steven Fe
 A variant of the neuron tracer that applies the **MedCLIPSeg** approach
 (Koleilat et al., *Probabilistic Vision–Language Adaptation for Data-Efficient
 and Generalizable Medical Image Segmentation*, CVPR 2026): a **text prompt**
-describes the target ("nerve fibers …"), a frozen **CLIP** backbone produces a
-**fiber-probability map**, and a probabilistic (Monte-Carlo) step yields a
-pixel-level **uncertainty map**.
+describes the target ("nerve fibers …") and a frozen **CLIP** backbone produces a
+**fiber-probability map**, which is then traced into a skeleton.
 
 ## What the app does
-- **Classical trace (reliable):** the deterministic 3D tracer produces the
-  quantification you should trust — number of fibers, total length, diameter,
-  branch points.
-- **MedCLIP (text-prompted):** shows the CLIP fiber-probability and uncertainty
-  maps for the neurofilament channel. The prompt is editable.
+- **Neurofilament (MIP):** the maximum-intensity projection of the input channel.
+- **Tracked skeleton (white on black):** the MedCLIP fiber-probability map is
+  thresholded, skeletonised and pruned into a fiber skeleton.
+- **Total trace length (µm):** the spacing-aware total length of that MedCLIP
+  skeleton, reported above the images. The text prompt and probability threshold
+  are editable.
 
 ## Few-shot training (included)
 A runnable, few-shot version of the MedCLIPSeg approach — **frozen CLIP encoders +

app.py

@@ -1,72 +1,48 @@
 """Neuron Quantification using AI — MedCLIPSeg variant.
 
-Same task as the main tracer, applying the MedCLIPSeg (CVPR 2026) vision-language
-approach: a text prompt describes the target ("nerve fibers ..."), a frozen CLIP
-backbone produces a fiber-probability map and a pixel-level UNCERTAINTY map.
+Applies the MedCLIPSeg (CVPR 2026) vision-language approach: a text prompt
+describes the target ("nerve fibers ..."), a frozen CLIP backbone produces a
+fiber-probability map, and that map is traced into a skeleton. The app reports a
+white-on-black **tracked skeleton** and its **total trace length** in microns.
 
-The reliable quantification (length, fibers, diameter, branch points) still comes
-from the deterministic tracer. The MedCLIP maps are shown alongside; zero-shot
-general CLIP is only a coarse prior on this domain — train the bundled model
-(reference_medclipseg/) on a GPU with masks for real segmentation quality.
+Zero-shot general CLIP is only a coarse prior on this domain — train the bundled
+model (reference_medclipseg/) on a GPU with masks for real segmentation quality.
 """
 import traceback
-import numpy as np
-import pandas as pd
 import gradio as gr
 
 import processing as P
 import medclipseg as MC
 
-METRIC_COLS = ["Region", "Number of fibers", "Total length (um)",
-               "Mean diameter (um)", "Branch points", "Area covered (um^2)"]
 
-
-def _cm(arr, name):
-    import matplotlib.cm as cm
-    return (getattr(cm, name)(np.clip(arr, 0, 1))[..., :3] * 255).astype(np.uint8)
-
-
-def analyze(file_obj, fg_prompt, run_medclip):
+def analyze(file_obj, fg_prompt, threshold):
     if file_obj is None:
-        return None, None, None, None, None, "Upload a CZI or TIFF z-stack."
+        return None, None, "Upload a CZI or TIFF z-stack."
     try:
         img = P.load_image(file_obj)
         nf, _ = P.guess_channels(img)
         nf_mip = P.channel_preview(img.data[nf])
 
-        # --- deterministic trace + quantification (the reliable numbers) ---
-        trace = P.trace_neurites(img.data[nf], img.voxel, sensitivity=1.0,
-                                 prune_um=3.0)
-        m = P.compute_metrics(trace, "Whole field", min_fiber_um=5.0)
-        overlay = P.overlay_on_original(nf_mip, trace.skeleton)
-        df = pd.DataFrame([{
-            "Region": "Whole field", "Number of fibers": m.n_fibers,
-            "Total length (um)": round(m.total_length_um, 2),
-            "Mean diameter (um)": round(m.mean_diameter_um, 3),
-            "Branch points": m.n_branch_points,
-            "Area covered (um^2)": round(m.area_covered_um2, 2)}],
-            columns=METRIC_COLS)
-        status = (f"**Classical trace:** {m.n_fibers} fibers / "
-                  f"{m.total_length_um:.0f} µm (the reliable quantification).")
+        if fg_prompt and fg_prompt.strip():
+            MC.FG_PROMPTS = [p.strip() for p in fg_prompt.split("|") if p.strip()]
 
-        prob_img = unc_img = None
-        if run_medclip:
-            if fg_prompt and fg_prompt.strip():
-                MC.FG_PROMPTS = [p.strip() for p in fg_prompt.split("|") if p.strip()]
-            prob, unc = MC.segment_best(nf_mip)
-            prob_img, unc_img = _cm(prob, "viridis"), _cm(unc, "magma")
-            if MC.has_trained_model():
-                status += (f"\n\n**MedCLIP (few-shot trained, text-prompted):** "
-                           f"fiber-probability + uncertainty maps (frozen CLIP + "
-                           f"trained decoder; held-out Dice ≈ 0.53). Prompt(s): "
-                           f"*{', '.join(MC.FG_PROMPTS)}*.")
-            else:
-                status += (f"\n\n**MedCLIP (zero-shot, text-prompted):** coarse "
-                           f"prior — train with train_medclip_fewshot.py for the "
-                           f"few-shot model.")
-        return overlay, nf_mip, prob_img, unc_img, df, status
+        # --- MedCLIP text-prompted fiber-probability map -> tracked skeleton ---
+        prob, _ = MC.segment_best(nf_mip)
+        trace = P.trace_from_probability(prob, img.voxel,
+                                         threshold=float(threshold), prune_um=3.0)
+        m = P.compute_metrics(trace, "MedCLIP", min_fiber_um=5.0)
+        skel_img = P.skeleton_image(trace.skeleton, dilate=1)   # white on black
+
+        model = ("few-shot trained decoder" if MC.has_trained_model()
+                 else "zero-shot CLIP prior")
+        status = (f"**MedCLIP total trace length: "
+                  f"{m.total_length_um:,.1f} µm.**\n\n"
+                  f"Tracked from the MedCLIP fiber-probability map "
+                  f"({model}, text-prompted) at threshold {float(threshold):.2f}. "
+                  f"Prompt(s): *{', '.join(MC.FG_PROMPTS)}*.")
+        return nf_mip, skel_img, status
     except Exception as e:  # noqa: BLE001
-        return None, None, None, None, None, f"Error:\n{e}\n{traceback.format_exc()}"
+        return None, None, f"Error:\n{e}\n{traceback.format_exc()}"
 
 
 HEADER_HTML = """
@@ -77,7 +53,7 @@ HEADER_HTML = """
     Iman Sabir Ezzat, Randa K Ismail, Ayden Chavez, Marisa Zallocchi, PhD, Steven Fernandes, PhD
   </div>
   <div style="font-weight:600; font-size:0.95rem; color:#4b5563; margin-top:0.25rem;">
-    MedCLIPSeg variant — text-prompted vision-language segmentation
+    MedCLIPSeg variant — text-prompted vision-language fiber tracking
   </div>
 </div>
 """
@@ -125,20 +101,17 @@ with gr.Blocks(title="Neuron Quantification using AI — MedCLIPSeg",
                 label="Text prompt(s) for the target (separate with | )",
                 value="a fluorescence microscopy image of nerve fibers | "
                       "neurofilament nerve fibers and axons")
-            use_mc = gr.Checkbox(value=True,
-                                 label="Run MedCLIP text-prompted segmentation")
+            thresh = gr.Slider(0.1, 0.9, value=0.5, step=0.05,
+                               label="Fiber-probability threshold")
             btn = gr.Button("Analyze", variant="primary")
         with gr.Column(scale=2):
             status = gr.Markdown()
             with gr.Row():
-                out_overlay = gr.Image(label="Traced neurons (classical)", height=250)
-                out_orig = gr.Image(label="Neurofilament (MIP)", height=250)
-            with gr.Row():
-                out_prob = gr.Image(label="MedCLIP fiber probability", height=250)
-                out_unc = gr.Image(label="MedCLIP uncertainty", height=250)
-            out_table = gr.Dataframe(label="Quantification", wrap=True)
-    btn.click(analyze, [file_in, prompt, use_mc],
-              [out_overlay, out_orig, out_prob, out_unc, out_table, status])
+                out_orig = gr.Image(label="Neurofilament (MIP)", height=280)
+                out_skel = gr.Image(label="Tracked skeleton (white on black)",
+                                    height=280)
+    btn.click(analyze, [file_in, prompt, thresh],
+              [out_orig, out_skel, status])
 
 if __name__ == "__main__":
     demo.launch()

processing.py

@@ -568,6 +568,36 @@ def trace_neurites(nf_vol: np.ndarray, voxel: tuple,
     return TraceResult(mask, skel, dist, voxel)
 
 
+def trace_from_probability(prob: np.ndarray, voxel: tuple,
+                           threshold: float = 0.5,
+                           min_object_px: int = 32,
+                           prune_um: float = 3.0) -> TraceResult:
+    """Trace a 2D fiber-probability map (e.g. a MedCLIP map) into a skeleton.
+
+    ``prob`` is a 2D float map in [0, 1] at the MIP resolution; ``threshold``
+    binarises it into a fiber mask which is cleaned, skeletonised and pruned.
+    The result is wrapped as a 3D (Z=1) ``TraceResult`` so the standard
+    ``compute_metrics`` / ``skeleton_image`` helpers apply unchanged and the
+    total length is measured with the same (dy, dx) spacing-aware code path.
+    """
+    dz, dy, dx = voxel
+    prob = np.asarray(prob, dtype=np.float32)
+    mask = prob >= float(threshold)
+    if mask.any():
+        mask = remove_small_objects(mask, int(min_object_px))
+        mask = ndi.binary_closing(mask, structure=np.ones((3, 3), bool))
+    if not mask.any():
+        z = np.zeros((1,) + prob.shape, bool)
+        return TraceResult(z, z.copy(), np.zeros(z.shape, np.float32), voxel)
+    # Wrap the 2D skeleton as a single-plane (Z=1) volume so the 3D-voxel
+    # pruning / metrics helpers apply unchanged (dz never affects in-plane steps).
+    skel = skeletonize(mask)[None]
+    if prune_um and prune_um > 0:
+        skel = prune_skeleton(skel, voxel, spur_um=float(prune_um))
+    dist = ndi.distance_transform_edt(mask, sampling=(dy, dx)).astype(np.float32)
+    return TraceResult(mask[None], skel, dist[None], voxel)
+
+
 # --------------------------------------------------------------------------- #
 #  Region definition (IHC vs OHC) from Myo7a
 # --------------------------------------------------------------------------- #