Upload 314 files

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+reference_medclipseg/data/Covid19/Prompts_Folder/Train_text.xlsx filter=lfs diff=lfs merge=lfs -text
+reference_medclipseg/data/EUS/Prompts_Folder/Test_text_original.xlsx filter=lfs diff=lfs merge=lfs -text
+reference_medclipseg/data/EUS/Prompts_Folder/Train_text_25.xlsx filter=lfs diff=lfs merge=lfs -text
+reference_medclipseg/data/EUS/Prompts_Folder/Train_text_50.xlsx filter=lfs diff=lfs merge=lfs -text
+reference_medclipseg/data/EUS/Prompts_Folder/Train_text.xlsx filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,8 +1,6 @@
 ---
-title: Cochlear Neurofilament Tracer
+title: Neuron Quantification using AI — MedCLIPSeg
 emoji: 🧠
-colorFrom: indigo
-colorTo: purple
 sdk: gradio
 sdk_version: 5.49.1
 app_file: app.py
@@ -10,67 +8,84 @@ pinned: false
 license: mit
 ---
 
-# 🧠 Cochlear Neurofilament Tracer
+# Neuron Quantification using AI — MedCLIPSeg variant
 
-A HuggingFace app that traces auditory-nerve fibers in confocal z-stacks of the
-organ of Corti and quantifies them **per frequency region**, separating
-**IHC-innervating** from **OHC-innervating** fibers.
+Iman Sabir Ezzat, Randa K Ismail, Ayden Chavez, Marisa Zallocchi, PhD, Steven Fernandes, PhD
 
-It is an alternative to IMARIS filament tracing that keeps each neuron as a
-**single continuous traced element** instead of splitting it into many
-threshold-dependent segments.
+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**.
 
-## Input
+## 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.
 
-- **File type:** Zeiss `.czi` 3D z-stacks. Generic `.tif/.tiff` stacks are also
-  accepted for flexibility.
-- **Channels:**
-  - *Neurofilament* — traces the neuron.
-  - *Myo7a* — marks hair cells; used as a reference to separate IHC- vs
-    OHC-innervating fibers. IHCs form a single row and OHCs form three adjacent
-    rows, so the Myo7a band is used to place the IHC/OHC boundary.
-- **Frequency region:** selectable (8/16/22/32/64 kHz), auto-detected from the
-  file name when possible.
-- Channels are auto-detected from CZI metadata (Alexa-555 → Neurofilament,
-  Alexa-405 → Myo7a) but can be reassigned in the UI.
+## Few-shot training (included)
+A runnable, few-shot version of the MedCLIPSeg approach — **frozen CLIP encoders +
+a text-conditioned decoder** — is trained here on **8 animals** (1A, 1B, 2A, 3A,
+3B, 4Aa, 4Bb, 5A) and **tested on the unseen C1 and C2b**:
 
-## What it does
-
-1. Segments and **skeletonises the Neurofilament network in 3D** using physical
-   voxel spacing (from CZI metadata, or entered for TIFF).
-2. Uses the **Myo7a channel** to place an IHC/OHC boundary. This can be set
-   manually (ROI 1 vs ROI 2) by moving the boundary slider while viewing the
-   Myo7a preview, choosing the split axis, and choosing which side is IHC.
-3. Computes, per region (Whole field / IHC / OHC):
-   - **Number of fibers** (continuous skeleton components above a minimum length)
-   - **Thickness / diameter** (from the 3D distance transform)
-   - **Length** (µm, spacing-aware)
-   - **Branching** (number of branch points)
-   - **Area covered** within the field of view (µm² and % of FOV)
-
-## Output
+```bash
+python train_medclip_fewshot.py     # runs on CPU/MPS; saves weights_medclip_decoder.pt
+```
 
-- A **black-background image** of the traced neurons in **white** (skeletonised
-  trace), plus a colour-coded IHC/OHC overlay.
-- An **Excel workbook** with all quantification, organized by frequency region,
-  with IHC and OHC fibers reported separately (tidy "Per region" sheet plus
-  per-metric frequency × region summary sheets).
+**Result: held-out (C1 + C2b) mean Dice ≈ 0.53** (per image 0.36–0.60; train Dice
+= 1.0, i.e. memorised). The segmentation lands on the fiber band and is far
+better than zero-shot, but it is **coarse** — CLIP features are 14×14, so thin
+fibers are only roughly captured, and the label is IMARIS's approximate
+reconstruction. The app loads `weights_medclip_decoder.pt` automatically and uses
+the trained model.
 
-The **Batch** tab processes several stacks at once (e.g. all frequency regions
-of one cochlea) and compiles one Excel workbook plus a ZIP of skeleton images.
+## Honest scope
+The full authors' model (`reference_medclipseg/`, PVL adapters + probabilistic
+attention) needs a **CUDA GPU**; the included few-shot decoder is the part that
+runs here. Either way the MedCLIP output is a **segmentation prior** — the
+**trustworthy quantification (length, fibers, diameter) comes from the classical
+tracer**, and the *length number* does not generalize reliably across animals
+(measured earlier). Without a trained decoder the app falls back to zero-shot
+CLIP, which is only a coarse prior.
 
-## Notes on method
+## To get the real, trained MedCLIPSeg
+The original authors' code is bundled in **`reference_medclipseg/`** (needs a
+**CUDA GPU**). Two steps:
 
-Confocal images of the organ of Corti are dense, so fully separating every
-individual axon is inherently ambiguous. This tool traces the network
-continuously and reports metrics **per region surrounding the IHCs / OHCs**,
-with a human-in-the-loop boundary for reliable IHC vs OHC assignment. The
-`sensitivity` control scales the segmentation threshold to capture more or fewer
-thin fibers.
+1. **Build the dataset** (no manual masks needed — uses IMARIS's reconstruction
+   as the label):
+   ```bash
+   python prepare_medclip_data.py
+   # -> data/Neurofilament/{Train,Val,Test}_Folder/{img,label} + Prompts_Folder/*.xlsx
+   ```
+2. **Fine-tune** the model on a GPU:
+   ```bash
+   cd reference_medclipseg
+   pip install -r requirements.txt
+   python train.py --config-file configs/Neurofilament.yaml
+   python test.py  --config-file configs/Neurofilament.yaml   # segmentation + uncertainty
+   ```
+   (`configs/Neurofilament.yaml` uses a BiomedCLIP backbone; switch `CLIP_MODEL`
+   to `clip`/`pubmedclip`/`unimedclip` as desired.)
 
-## Local run
+## Honest limitations (unchanged from the main project)
+- The training **label** is IMARIS's *reconstructed* skeleton (from its segment
+  coordinates), which is approximate, and IMARIS itself **over-traces** the true
+  fibre length (~3–4×). So a trained model would imitate IMARIS, not ground truth.
+- Predicting IMARIS's numbers does **not** generalize well across animals
+  (measured R² ≈ 0 in earlier experiments). MedCLIPSeg's segmentation may look
+  cleaner, but the **length number** remains unreliable on unseen animals.
+- For exact IMARIS values, read the IMARIS statistics directly (as in the other
+  tools).
 
-```bash
-pip install -r requirements.txt
-python app.py
-```
+## Files
+- `app.py` — the Gradio variant (classical trace + MedCLIP maps).
+- `medclipseg.py` — runnable text-prompted CLIP segmentation + uncertainty.
+- `prepare_medclip_data.py` — builds the MedCLIPSeg dataset from CZIs + IMARIS
+  reconstruction masks.
+- `reference_medclipseg/` — the original MedCLIPSeg code (for GPU fine-tuning) +
+  `configs/Neurofilament.yaml`.
+- `processing.py` — the shared image-processing / tracing pipeline.

app.py

@@ -1,373 +1,144 @@
-"""
-Cochlear Neurofilament Tracer — HuggingFace Gradio app
-======================================================
+"""Neuron Quantification using AI — MedCLIPSeg variant.
 
-Traces auditory-nerve fibers (Neurofilament channel) in confocal z-stacks of
-the organ of Corti, uses the Myo7a hair-cell channel to separate
-IHC-innervating from OHC-innervating fibers, and reports per-region
-quantification (number of fibers, diameter, length, branch points, area
-covered) plus a black-background skeleton image and an Excel workbook.
+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.
 
-Accepts Zeiss **CZI** z-stacks and generic **TIFF** stacks.
+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.
 """
-
-import os
-import tempfile
-import zipfile
 import traceback
-
 import numpy as np
 import pandas as pd
 import gradio as gr
-from skimage.io import imsave
-from scipy import ndimage as ndi
 
 import processing as P
+import medclipseg as MC
 
-OUT_DIR = tempfile.mkdtemp(prefix="neuron_tracer_")
-
-METRIC_COLS = [
-    "File", "Frequency region", "Region", "Number of fibers",
-    "Total length (um)", "Mean diameter (um)", "Median diameter (um)",
-    "Branch points", "Area covered (um^2)", "FOV area (um^2)",
-    "Area covered (% of FOV)",
-]
-
-
-# --------------------------------------------------------------------------- #
-#  Small helpers
-# --------------------------------------------------------------------------- #
-
-def _channel_choices(img: P.LoadedImage):
-    choices = []
-    for i, ch in enumerate(img.channels):
-        dye = ch.get("dye") or "no dye / transmitted"
-        choices.append((f"Ch {i} — {ch.get('name','?')} ({dye})", i))
-    return choices
-
-
-def _draw_boundary_on(gray_u8, boundary_frac, axis="Y"):
-    """Return an RGB copy of a grayscale MIP with a yellow boundary line."""
-    rgb = np.stack([gray_u8] * 3, axis=-1)
-    ny, nx = gray_u8.shape
-    if axis.upper() == "Y":
-        b = min(max(int(round(boundary_frac * ny)), 0), ny - 1)
-        rgb[b, :] = (255, 255, 0)
-    else:
-        b = min(max(int(round(boundary_frac * nx)), 0), nx - 1)
-        rgb[:, b] = (255, 255, 0)
-    return rgb
-
-
-def _save_png(arr, name):
-    path = os.path.join(OUT_DIR, name)
-    imsave(path, arr)
-    return path
-
+METRIC_COLS = ["Region", "Number of fibers", "Total length (um)",
+               "Mean diameter (um)", "Branch points", "Area covered (um^2)"]
 
-def _build_excel(rows, path):
-    """Write a tidy per-region sheet plus a frequency×region summary sheet."""
-    df = pd.DataFrame(rows, columns=METRIC_COLS)
-    with pd.ExcelWriter(path, engine="openpyxl") as xl:
-        df.to_excel(xl, sheet_name="Per region", index=False)
-        # Summary: only IHC / OHC rows, pivoted by frequency region.
-        sub = df[df["Region"].isin(["IHC region", "OHC region"])]
-        if not sub.empty:
-            for metric in ["Number of fibers", "Total length (um)",
-                           "Mean diameter (um)", "Branch points",
-                           "Area covered (um^2)"]:
-                piv = sub.pivot_table(index="Frequency region",
-                                      columns="Region", values=metric,
-                                      aggfunc="mean")
-                sheet = metric.split(" (")[0][:28]
-                piv.to_excel(xl, sheet_name=f"{sheet}")
-    return df
 
+def _cm(arr, name):
+    import matplotlib.cm as cm
+    return (getattr(cm, name)(np.clip(arr, 0, 1))[..., :3] * 255).astype(np.uint8)
 
-# --------------------------------------------------------------------------- #
-#  Single-image interactive flow
-# --------------------------------------------------------------------------- #
 
-def load_and_preview(file_obj, dz, dy, dx):
+def analyze(file_obj, fg_prompt, run_medclip):
     if file_obj is None:
-        return (None, gr.update(), gr.update(), gr.update(),
-                gr.update(), None, None, "Please upload a CZI or TIFF file.")
+        return None, None, None, None, None, "Upload a CZI or TIFF z-stack."
     try:
-        img = P.load_image(file_obj, dz=float(dz), dy=float(dy), dx=float(dx))
-    except Exception as e:
-        return (None, gr.update(), gr.update(), gr.update(),
-                gr.update(), None, None,
-                f"❌ Failed to load file:\n{e}\n{traceback.format_exc()}")
-
-    nf, myo = P.guess_channels(img)
-    choices = _channel_choices(img)
-    freq = P.detect_frequency(img.source_name)
-    bfrac = P.suggest_boundary(img.data[myo])
-
-    nf_prev = P.channel_preview(img.data[nf])
-    myo_prev = _draw_boundary_on(P.channel_preview(img.data[myo]), bfrac, "Y")
-
-    dz_, dy_, dx_ = img.voxel
-    status = (f"✅ Loaded **{img.source_name}** — shape "
-              f"{img.data.shape} (C,Z,Y,X)\n\n"
-              f"Voxel size: dz={dz_:.3f}, dy={dy_:.4f}, dx={dx_:.4f} µm.  "
-              f"Detected frequency region: **{freq}**.\n\n"
-              f"Auto-picked Neurofilament = Ch {nf}, Myo7a = Ch {myo}. "
-              f"Adjust below if needed, then press **Run analysis**.")
-    return (img,
-            gr.update(choices=choices, value=nf),
-            gr.update(choices=choices, value=myo),
-            gr.update(value=freq),
-            gr.update(value=round(bfrac, 3)),
-            nf_prev, myo_prev, status)
-
-
-def refresh_boundary_preview(img, myo_idx, boundary, axis):
-    if img is None or myo_idx is None:
-        return None
-    myo_prev = P.channel_preview(img.data[int(myo_idx)])
-    return _draw_boundary_on(myo_prev, float(boundary), axis)
-
-
-def run_single(img, nf_idx, myo_idx, freq, axis, ihc_side, boundary,
-               sensitivity, min_fiber):
-    if img is None:
-        return None, None, None, None, None, "Please load an image first."
-    if nf_idx is None:
-        return None, None, None, None, None, "Please select the Neurofilament channel."
-    try:
-        nf_idx, myo_idx = int(nf_idx), int(myo_idx)
-        trace = P.trace_neurites(img.data[nf_idx], img.voxel,
-                                 sensitivity=float(sensitivity))
-        shape_yx = img.data[nf_idx].shape[1:]
-        side = "low" if ihc_side.startswith("Low") else "high"
-        ihc_roi, ohc_roi = P.make_region_masks(shape_yx, float(boundary),
-                                               ihc_side=side, axis=axis)
-
-        whole = P.compute_metrics(trace, "Whole field",
-                                  min_fiber_um=float(min_fiber))
-        m_ihc = P.compute_metrics(trace, "IHC region", ihc_roi,
-                                  min_fiber_um=float(min_fiber))
-        m_ohc = P.compute_metrics(trace, "OHC region", ohc_roi,
-                                  min_fiber_um=float(min_fiber))
-
-        skel_img = P.skeleton_image(trace.skeleton)
-        region_img = P.region_overlay(trace.skeleton, ihc_roi, ohc_roi,
-                                      float(boundary), axis)
-        myo_img = _draw_boundary_on(P.channel_preview(img.data[myo_idx]),
-                                    float(boundary), axis)
-
-        stem = os.path.splitext(img.source_name)[0]
-        skel_path = _save_png(skel_img, f"{stem}_skeleton.png")
-
-        rows = []
-        for m in (whole, m_ihc, m_ohc):
-            r = m.as_row()
-            r = {"File": img.source_name, "Frequency region": freq, **r}
-            rows.append(r)
-        df = pd.DataFrame(rows, columns=METRIC_COLS)
-        xl_path = os.path.join(OUT_DIR, f"{stem}_quantification.xlsx")
-        _build_excel(rows, xl_path)
-
-        status = (f"✅ Done. Traced {int(trace.skeleton.sum())} skeleton voxels. "
-                  f"IHC={m_ihc.n_fibers} fibers / {m_ihc.total_length_um:.0f} µm, "
-                  f"OHC={m_ohc.n_fibers} fibers / {m_ohc.total_length_um:.0f} µm.")
-        # Return skeleton path also as downloadable file
-        return (skel_img, region_img, myo_img, df,
-                [skel_path, xl_path], status)
-    except Exception as e:
-        return (None, None, None, None, None,
-                f"❌ Error:\n{e}\n{traceback.format_exc()}")
-
-
-# --------------------------------------------------------------------------- #
-#  Batch flow
-# --------------------------------------------------------------------------- #
-
-def run_batch(files, axis, ihc_side, sensitivity, min_fiber, dz, dy, dx,
-              progress=gr.Progress()):
-    if not files:
-        return None, None, None, "Please upload one or more files."
-    side = "low" if ihc_side.startswith("Low") else "high"
-    all_rows, gallery, skel_paths = [], [], []
-    log = []
-    for f in progress.tqdm(files, desc="Processing"):
-        path = f if isinstance(f, str) else f.name
-        name = os.path.basename(path)
-        try:
-            img = P.load_image(path, dz=float(dz), dy=float(dy), dx=float(dx))
-            nf, myo = P.guess_channels(img)
-            freq = P.detect_frequency(name)
-            trace = P.trace_neurites(img.data[nf], img.voxel,
-                                     sensitivity=float(sensitivity))
-            bfrac = P.suggest_boundary(img.data[myo])
-            shape_yx = img.data[nf].shape[1:]
-            ihc_roi, ohc_roi = P.make_region_masks(shape_yx, bfrac,
-                                                   ihc_side=side, axis=axis)
-            for m in (P.compute_metrics(trace, "Whole field",
-                                        min_fiber_um=float(min_fiber)),
-                      P.compute_metrics(trace, "IHC region", ihc_roi,
-                                        min_fiber_um=float(min_fiber)),
-                      P.compute_metrics(trace, "OHC region", ohc_roi,
-                                        min_fiber_um=float(min_fiber))):
-                all_rows.append({"File": name, "Frequency region": freq,
-                                 **m.as_row()})
-            skel_img = P.skeleton_image(trace.skeleton)
-            stem = os.path.splitext(name)[0]
-            sp = _save_png(skel_img, f"{stem}_skeleton.png")
-            skel_paths.append(sp)
-            gallery.append((skel_img, f"{name} ({freq})"))
-            log.append(f"✅ {name}: {freq}")
-        except Exception as e:
-            log.append(f"❌ {name}: {e}")
-
-    if not all_rows:
-        return None, None, gallery, "No files processed.\n" + "\n".join(log)
-
-    xl_path = os.path.join(OUT_DIR, "batch_quantification.xlsx")
-    df = _build_excel(all_rows, xl_path)
-
-    zip_path = os.path.join(OUT_DIR, "batch_skeletons.zip")
-    with zipfile.ZipFile(zip_path, "w") as z:
-        for sp in skel_paths:
-            z.write(sp, os.path.basename(sp))
-        z.write(xl_path, os.path.basename(xl_path))
-
-    return df, [xl_path, zip_path], gallery, "\n".join(log)
-
-
-# --------------------------------------------------------------------------- #
-#  UI
-# --------------------------------------------------------------------------- #
-
-INTRO = """
-# 🧠 Cochlear Neurofilament Tracer
-
-Trace auditory-nerve fibers in confocal z-stacks and quantify them **per
-frequency region**, separating **IHC-innervating** from **OHC-innervating**
-fibers using the Myo7a hair-cell channel.
-
-**Channels expected:** *Neurofilament* (traces the neuron) and *Myo7a* (hair
-cells — reference to split IHC vs OHC). IHCs form a single row, OHCs form three
-rows, so the Myo7a band is used to place the IHC/OHC boundary — which you can
-fine-tune by hand.
-
-**Input:** Zeiss `.czi` z-stacks or generic `.tif/.tiff` stacks.
+        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).")
+
+        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
+    except Exception as e:  # noqa: BLE001
+        return None, None, None, None, None, f"Error:\n{e}\n{traceback.format_exc()}"
+
+
+HEADER_HTML = """
+<div style="text-align:center; margin: 0.2rem 0 0.7rem;">
+  <h1 style="font-size:2.1rem; font-weight:750; letter-spacing:-0.01em;
+             margin:0 0 0.4rem;">Neuron Quantification using AI</h1>
+  <div style="font-weight:700; font-size:1.15rem; color:#1a1a1a;">
+    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
+  </div>
+</div>
 """
 
-with gr.Blocks(title="Cochlear Neurofilament Tracer", theme=gr.themes.Soft()) as demo:
-    gr.Markdown(INTRO)
-
-    with gr.Tab("Single image (interactive)"):
-        img_state = gr.State()
-        with gr.Row():
-            with gr.Column(scale=1):
-                file_in = gr.File(label="Upload CZI or TIFF",
-                                  file_types=[".czi", ".tif", ".tiff"],
-                                  type="filepath")
-                with gr.Accordion("Voxel size (µm) — used for TIFF; CZI reads "
-                                  "its own", open=False):
-                    dz_in = gr.Number(0.35, label="dz (µm/plane)")
-                    dy_in = gr.Number(0.0895, label="dy (µm/px)")
-                    dx_in = gr.Number(0.0895, label="dx (µm/px)")
-                load_btn = gr.Button("① Load & preview", variant="secondary")
-
-                nf_dd = gr.Dropdown(label="Neurofilament channel", choices=[])
-                myo_dd = gr.Dropdown(label="Myo7a channel", choices=[])
-                freq_dd = gr.Dropdown(label="Frequency region",
-                                      choices=P.FREQ_CHOICES,
-                                      value="Other / unknown")
-
-                gr.Markdown("**IHC / OHC region split** (Myo7a-guided)")
-                axis_dd = gr.Radio(["Y", "X"], value="Y",
-                                   label="Split axis (Y = radial, usual)")
-                side_dd = gr.Radio(["Low side = IHC", "High side = IHC"],
-                                   value="Low side = IHC",
-                                   label="Which side is IHC?")
-                boundary_sl = gr.Slider(0.0, 1.0, value=0.5, step=0.005,
-                                        label="Boundary position (fraction "
-                                              "along split axis)")
-
-                gr.Markdown("**Tracing**")
-                sens_sl = gr.Slider(0.5, 1.5, value=1.0, step=0.05,
-                                    label="Sensitivity (↑ = capture more/thinner "
-                                          "fibers)")
-                minfib_sl = gr.Slider(0.0, 20.0, value=5.0, step=0.5,
-                                      label="Min fiber length to count (µm)")
-                run_btn = gr.Button("② Run analysis", variant="primary")
-
-            with gr.Column(scale=2):
-                status = gr.Markdown()
-                with gr.Row():
-                    nf_prev = gr.Image(label="Neurofilament (MIP)",
-                                       height=220)
-                    myo_prev = gr.Image(label="Myo7a (MIP) + boundary",
-                                        height=220)
-                skel_out = gr.Image(label="Traced neurons — white on black",
-                                    height=300)
-                region_out = gr.Image(label="Region overlay (cyan = IHC, "
-                                            "magenta = OHC)", height=300)
-                table = gr.Dataframe(label="Quantification", wrap=True)
-                files_out = gr.Files(label="Downloads (skeleton PNG + Excel)")
-
-        load_btn.click(load_and_preview,
-                       [file_in, dz_in, dy_in, dx_in],
-                       [img_state, nf_dd, myo_dd, freq_dd, boundary_sl,
-                        nf_prev, myo_prev, status])
-        # Live boundary preview
-        for comp in (boundary_sl, myo_dd, axis_dd):
-            comp.change(refresh_boundary_preview,
-                        [img_state, myo_dd, boundary_sl, axis_dd], myo_prev)
-        run_btn.click(run_single,
-                      [img_state, nf_dd, myo_dd, freq_dd, axis_dd, side_dd,
-                       boundary_sl, sens_sl, minfib_sl],
-                      [skel_out, region_out, myo_prev, table, files_out, status])
-
-    with gr.Tab("Batch (multiple images)"):
-        gr.Markdown(
-            "Upload several z-stacks (e.g. all frequency regions of one "
-            "cochlea). Each is auto-traced with an auto-placed IHC/OHC "
-            "boundary, and results are combined into one Excel workbook "
-            "organized by frequency region.")
-        with gr.Row():
-            with gr.Column(scale=1):
-                batch_files = gr.File(label="Upload CZI/TIFF files",
-                                      file_count="multiple",
-                                      file_types=[".czi", ".tif", ".tiff"],
-                                      type="filepath")
-                b_axis = gr.Radio(["Y", "X"], value="Y", label="Split axis")
-                b_side = gr.Radio(["Low side = IHC", "High side = IHC"],
-                                  value="Low side = IHC",
-                                  label="Which side is IHC?")
-                b_sens = gr.Slider(0.5, 1.5, value=1.0, step=0.05,
-                                   label="Sensitivity")
-                b_minfib = gr.Slider(0.0, 20.0, value=5.0, step=0.5,
-                                     label="Min fiber length (µm)")
-                with gr.Accordion("Voxel size (µm) for TIFF", open=False):
-                    b_dz = gr.Number(0.35, label="dz")
-                    b_dy = gr.Number(0.0895, label="dy")
-                    b_dx = gr.Number(0.0895, label="dx")
-                batch_btn = gr.Button("Run batch", variant="primary")
-            with gr.Column(scale=2):
-                batch_log = gr.Textbox(label="Log", lines=6)
-                batch_table = gr.Dataframe(label="Combined quantification",
-                                           wrap=True)
-                batch_files_out = gr.Files(label="Downloads (Excel + ZIP)")
-                batch_gallery = gr.Gallery(label="Skeleton traces",
-                                           columns=3, height=400)
-        batch_btn.click(run_batch,
-                        [batch_files, b_axis, b_side, b_sens, b_minfib,
-                         b_dz, b_dy, b_dx],
-                        [batch_table, batch_files_out, batch_gallery, batch_log])
+THEME = gr.themes.Base(
+    primary_hue=gr.themes.colors.slate, secondary_hue=gr.themes.colors.slate,
+    neutral_hue=gr.themes.colors.gray,
+    font=["system-ui", "-apple-system", "Segoe UI", "Roboto", "sans-serif"],
+).set(
+    body_background_fill="#ffffff", body_text_color="#1a1a1a",
+    background_fill_primary="#ffffff", background_fill_secondary="#f7f7f8",
+    block_background_fill="#ffffff", block_border_color="#e5e7eb",
+    block_label_text_color="#1a1a1a", block_title_text_color="#1a1a1a",
+    border_color_primary="#e5e7eb", button_primary_background_fill="#1f2937",
+    button_primary_text_color="#ffffff", input_background_fill="#ffffff",
+    input_border_color="#c0c5cc",
+    body_background_fill_dark="#ffffff", body_text_color_dark="#1a1a1a",
+    background_fill_primary_dark="#ffffff", background_fill_secondary_dark="#f7f7f8",
+    block_background_fill_dark="#ffffff", block_border_color_dark="#e5e7eb",
+    block_label_text_color_dark="#1a1a1a", block_title_text_color_dark="#1a1a1a",
+    panel_background_fill_dark="#ffffff", border_color_primary_dark="#e5e7eb",
+    button_primary_background_fill_dark="#1f2937",
+    button_primary_text_color_dark="#ffffff", input_background_fill_dark="#ffffff",
+)
+CSS = """
+.gradio-container { max-width: 1200px !important; margin: 0 auto !important; }
+:root, .dark {
+    color-scheme: light; --body-background-fill:#ffffff;
+    --background-fill-primary:#ffffff; --block-background-fill:#ffffff;
+    --body-text-color:#1a1a1a; --block-label-text-color:#1a1a1a;
+    --block-title-text-color:#1a1a1a; --border-color-primary:#e5e7eb;
+    --input-background-fill:#ffffff; --neutral-950:#1a1a1a;
+}
+body, gradio-app, .gradio-container, .dark { background:#ffffff !important; color:#1a1a1a !important; }
+"""
 
-    gr.Markdown(
-        "---\n*Method:* the Neurofilament channel is smoothed, thresholded "
-        "(Otsu, scaled by the sensitivity control) and skeletonised in 3D; "
-        "length, diameter (from the 3D distance transform), branch points and "
-        "footprint area are measured with physical voxel spacing. Each fiber is "
-        "a connected skeleton component ≥ the minimum length. The Myo7a band "
-        "defines the IHC/OHC boundary, and metrics are reported for each "
-        "region.")
+with gr.Blocks(title="Neuron Quantification using AI — MedCLIPSeg",
+               theme=THEME, css=CSS) as demo:
+    gr.HTML(HEADER_HTML)
+    with gr.Row():
+        with gr.Column(scale=1):
+            file_in = gr.File(label="Neurofilament z-stack (.czi / .tif)",
+                              type="filepath")
+            prompt = gr.Textbox(
+                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")
+            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])
 
 if __name__ == "__main__":
     demo.launch()

medclipseg.py

@@ -0,0 +1,172 @@
+"""MedCLIPSeg-style text-prompted segmentation, adapted to run here.
+
+Faithful to the *approach* of Koleilat et al., "MedCLIPSeg: Probabilistic
+Vision-Language Adaptation ..." (CVPR 2026): a FROZEN vision-language (CLIP)
+backbone provides patch-level embeddings; short TEXT PROMPTS describe the target
+("nerve fibers ...") and the background; patch-text similarity yields a dense
+segmentation probability, and a PROBABILISTIC step (Monte-Carlo sampling of the
+patch embedding with an estimated variance — the paper samples attention Values
+from learned distributions) yields a pixel-level UNCERTAINTY map.
+
+This module is ZERO-SHOT (no training) using a general CLIP backbone, so on
+cochlear neurofilament it is a *coarse prior*, not a trained segmentor. For the
+real, trained model (the paper's PVL adapters + decoder, fine-tuned on masks)
+see reference_medclipseg/ and train it on a GPU with the data from
+prepare_medclip_data.py.
+"""
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+from skimage.transform import resize
+
+_MODEL = "openai/clip-vit-base-patch16"      # 16-px patches -> 14x14 per 224 tile
+_clip = None
+_proc = None
+
+FG_PROMPTS = [
+    "a fluorescence microscopy image of nerve fibers",
+    "neurofilament nerve fibers and axons",
+    "a dense network of thin bright nerve fibers",
+]
+BG_PROMPTS = [
+    "a black empty background",
+    "dark region with no tissue",
+    "background noise",
+]
+
+
+def _load():
+    global _clip, _proc
+    if _clip is None:
+        from transformers import CLIPModel, CLIPProcessor
+        _clip = CLIPModel.from_pretrained(_MODEL).eval()
+        _proc = CLIPProcessor.from_pretrained(_MODEL)
+    return _clip, _proc
+
+
+def _text_embeds(prompts):
+    clip, proc = _load()
+    tok = proc(text=prompts, return_tensors="pt", padding=True)
+    with torch.no_grad():
+        out = clip.text_model(input_ids=tok["input_ids"],
+                              attention_mask=tok.get("attention_mask"))
+        t = clip.text_projection(out.pooler_output)
+    return t / t.norm(dim=-1, keepdim=True)                 # (K, D)
+
+
+def _dense_patch_embeds(win_u8):
+    """Per-patch CLIP embeddings for a 224x224 uint8 RGB window (MaskCLIP-style)."""
+    clip, proc = _load()
+    px = proc(images=win_u8, return_tensors="pt")["pixel_values"]
+    vm = clip.vision_model
+    with torch.no_grad():
+        toks = vm(pixel_values=px).last_hidden_state[:, 1:, :]   # drop CLS -> (1,196,768)
+        emb = clip.visual_projection(vm.post_layernorm(toks))     # (1,196,D)
+    emb = emb / emb.norm(dim=-1, keepdim=True)
+    return emb[0]                                             # (196, D)
+
+
+def segment(mip_u8, n_mc=8, sigma=0.05, tile=224):
+    """Text-prompted CLIP segmentation of a grayscale MIP.
+
+    Returns (prob, uncertainty) float maps at the input resolution, both in
+    [0, 1]. ``prob`` is the fiber probability; ``uncertainty`` is the std of the
+    probability across ``n_mc`` Monte-Carlo embedding samples.
+    """
+    H0, W0 = mip_u8.shape
+    rows = max(1, round(H0 / tile))
+    cols = max(1, round(W0 / tile))
+    Hr, Wr = rows * tile, cols * tile
+    rgb = np.stack([resize(mip_u8, (Hr, Wr), preserve_range=True)] * 3, -1).astype(np.uint8)
+
+    fg_t, bg_t = _text_embeds(FG_PROMPTS), _text_embeds(BG_PROMPTS)
+    gh, gw = 14, 14                                          # patches per tile (224/16)
+    prob = np.zeros((rows * gh, cols * gw), np.float32)
+    unc = np.zeros((rows * gh, cols * gw), np.float32)
+    for r in range(rows):
+        for c in range(cols):
+            win = rgb[r * tile:(r + 1) * tile, c * tile:(c + 1) * tile]
+            emb = _dense_patch_embeds(win)                  # (196, D)
+            samples = []
+            for m in range(n_mc):
+                e = emb if m == 0 else torch.nn.functional.normalize(
+                    emb + sigma * torch.randn_like(emb), dim=-1)
+                fg = (e @ fg_t.T).mean(-1)                   # (196,)
+                bg = (e @ bg_t.T).mean(-1)
+                samples.append(torch.sigmoid((fg - bg) / 0.1).numpy())
+            samples = np.stack(samples).reshape(n_mc, gh, gw)
+            prob[r * gh:(r + 1) * gh, c * gw:(c + 1) * gw] = samples.mean(0)
+            unc[r * gh:(r + 1) * gh, c * gw:(c + 1) * gw] = samples.std(0)
+    prob = resize(prob, (H0, W0), order=1, preserve_range=True)
+    unc = resize(unc, (H0, W0), order=1, preserve_range=True)
+    # normalise uncertainty to [0,1] for display
+    unc = (unc - unc.min()) / (np.ptp(unc) + 1e-6)
+    return prob.astype(np.float32), unc.astype(np.float32)
+
+
+# --------------------------------------------------------------------------- #
+#  Few-shot TRAINED decoder (frozen CLIP + text-conditioned decoder).
+#  Trained with train_medclip_fewshot.py (8 animals train, C1/C2b test).
+# --------------------------------------------------------------------------- #
+
+_WEIGHTS = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                        "weights_medclip_decoder.pt")
+_decoder = None
+
+
+class _Decoder(nn.Module):
+    def __init__(self, cin=768, tdim=512, C=128):
+        super().__init__()
+        self.reduce = nn.Conv2d(cin, C, 1)
+        self.film = nn.Linear(tdim, 2 * C)
+
+        def up(ci, co):
+            return nn.Sequential(nn.ConvTranspose2d(ci, co, 2, 2),
+                                 nn.Conv2d(co, co, 3, padding=1),
+                                 nn.BatchNorm2d(co), nn.ReLU(True))
+        self.up = nn.Sequential(up(C, C), up(C, C), up(C, C // 2), up(C // 2, C // 4))
+        self.out = nn.Conv2d(C // 4, 1, 1)
+
+    def forward(self, feat, tvec):
+        x = self.reduce(feat)
+        g, b = self.film(tvec).chunk(2, -1)
+        x = x * (1 + g[:, :, None, None]) + b[:, :, None, None]
+        return torch.sigmoid(self.out(self.up(x)))
+
+
+def _clip_patch_grid(mip_u8, size=224):
+    clip, proc = _load()
+    rgb = np.stack([resize(mip_u8, (size, size), preserve_range=True)] * 3, -1).astype(np.uint8)
+    px = proc(images=rgb, return_tensors="pt")["pixel_values"]
+    with torch.no_grad():
+        toks = clip.vision_model(pixel_values=px).last_hidden_state[0, 1:, :]
+    g = int(round(toks.shape[0] ** 0.5))
+    return toks.reshape(g, g, -1).permute(2, 0, 1).contiguous()[None]     # (1,768,g,g)
+
+
+def has_trained_model():
+    return os.path.exists(_WEIGHTS)
+
+
+def segment_trained(mip_u8):
+    """Segment with the few-shot TRAINED decoder. Returns (prob, uncertainty)."""
+    global _decoder
+    if _decoder is None:
+        _decoder = _Decoder().eval()
+        _decoder.load_state_dict(torch.load(_WEIGHTS, map_location="cpu")["model"])
+    tvec = _text_embeds(FG_PROMPTS).mean(0, keepdim=True)
+    feat = _clip_patch_grid(mip_u8)
+    with torch.no_grad():
+        prob = _decoder(feat, tvec)[0, 0].numpy()
+    H0, W0 = mip_u8.shape
+    prob = resize(prob, (H0, W0), order=1, preserve_range=True).astype(np.float32)
+    unc = np.clip(4 * prob * (1 - prob), 0, 1).astype(np.float32)         # entropy-like
+    return prob, unc
+
+
+def segment_best(mip_u8):
+    """Use the trained few-shot decoder if available, else zero-shot CLIP."""
+    if has_trained_model():
+        return segment_trained(mip_u8)
+    return segment(mip_u8)

prepare_medclip_data.py

@@ -0,0 +1,115 @@
+"""Build a MedCLIPSeg dataset for the neurofilament task (no manual masks needed).
+
+For every CZI with an IMARIS _Statistics folder, writes:
+  data/Neurofilament/<split>_Folder/img/<name>.png     (neurofilament MIP, RGB)
+  data/Neurofilament/<split>_Folder/label/<name>.png   (IMARIS-reconstruction mask)
+  data/Neurofilament/Prompts_Folder/<split>_text.xlsx  (Image, Ground Truth, Description)
+
+The mask is the k-NN reconstruction of IMARIS's traced skeleton from its segment
+coordinates (dilated) — the same approximate label used elsewhere. Splits are by
+ANIMAL so validation/test are unseen animals. Then fine-tune the bundled model:
+
+  cd reference_medclipseg
+  python train.py --config-file configs/Neurofilament.yaml   # needs a CUDA GPU
+
+Honest note: this label is approximate (reconstructed, not IMARIS's exact voxel
+mask), and the target is IMARIS's over-traced tangle — see README.
+"""
+import argparse, glob, os, sys
+import numpy as np
+import pandas as pd
+from scipy import ndimage as ndi
+from scipy.spatial import cKDTree
+from skimage.transform import resize
+from skimage.io import imsave
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import processing as P
+
+ROOT = "/Users/slf20757/Desktop/Dr. Fernandes"
+SIZE = 224
+PROMPT = "a fluorescence microscopy image of neurofilament nerve fibers"
+# animal -> split
+VAL = {"C1"}
+TEST = {"C2b"}
+
+
+def seg_xyz(sd, stem):
+    def cols(name, nc):
+        r = []
+        for line in open(os.path.join(sd, f"{stem}_{name}.csv"), errors="ignore"):
+            p = line.split(",")
+            try:
+                r.append([float(p[i]) for i in range(nc)])
+            except (ValueError, IndexError):
+                pass
+        return np.array(r)
+    return cols("Segment_Position", 3)
+
+
+def skel_mask(pos, voxel, shape_yx, k=3):
+    dz, dy, dx = voxel; ny, nx = shape_yx
+    col = pos[:, 0] / dx * (SIZE / nx); row = pos[:, 1] / dy * (SIZE / ny)
+    pts = np.stack([row, col], 1)
+    tree = cKDTree(pts); dist, idx = tree.query(pts, k=min(k + 1, len(pts)))
+    radius = 2.0 / dx * (SIZE / nx)
+    m = np.zeros((SIZE, SIZE), bool)
+    for i in range(len(pts)):
+        for j, d in zip(idx[i, 1:], dist[i, 1:]):
+            if d <= radius:
+                p0, p1 = pts[i], pts[j]
+                n = max(int(np.abs(p1 - p0).max()) + 1, 2)
+                t = np.linspace(0, 1, n)[:, None]; line = p0 + t * (p1 - p0)
+                m[np.clip(line[:, 0].astype(int), 0, SIZE - 1),
+                  np.clip(line[:, 1].astype(int), 0, SIZE - 1)] = True
+    return ndi.binary_dilation(m, iterations=1)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", default="data/Neurofilament")
+    ap.add_argument("--limit", type=int, default=0, help="cap images (for testing)")
+    args = ap.parse_args()
+    rows = {"Train": [], "Val": [], "Test": []}
+    n = 0
+    for csv in sorted(glob.glob(os.path.join(ROOT, "*", "*_Statistics",
+                                             "*_Filament_Length_(sum).csv"))):
+        sd = os.path.dirname(csv)
+        stem = os.path.basename(csv).replace("_Filament_Length_(sum).csv", "")
+        animal = os.path.relpath(csv, ROOT).split(os.sep)[0]
+        czi = os.path.join(ROOT, animal, stem.replace(" filament", "") + ".czi")
+        if not (os.path.exists(czi) and
+                os.path.exists(os.path.join(sd, f"{stem}_Segment_Position_X.csv"))):
+            continue
+        split = "Val" if animal in VAL else "Test" if animal in TEST else "Train"
+        img = P.load_image(czi); nf, _ = P.guess_channels(img)
+        vol = img.data[nf]; ny, nx = vol.shape[1], vol.shape[2]
+        mip = resize(P.channel_preview(vol), (SIZE, SIZE), preserve_range=True).astype(np.uint8)
+        mask = (skel_mask(seg_xyz(sd, stem), img.voxel, (ny, nx)) * 255).astype(np.uint8)
+        name = stem.replace(" ", "_") + ".png"
+        for sub in ("img", "label"):
+            os.makedirs(os.path.join(args.out, f"{split}_Folder", sub), exist_ok=True)
+        imsave(os.path.join(args.out, f"{split}_Folder", "img", name),
+               np.stack([mip] * 3, -1))
+        imsave(os.path.join(args.out, f"{split}_Folder", "label", name), mask)
+        rows[split].append({"Image": name, "Ground Truth": name, "Description": PROMPT})
+        n += 1
+        print(f"[{split}] {stem}: mask px={int((mask>0).sum())}", flush=True)
+        del img, vol
+        if args.limit and n >= args.limit:
+            break
+
+    os.makedirs(os.path.join(args.out, "Prompts_Folder"), exist_ok=True)
+    for split, r in rows.items():
+        if r:
+            pd.DataFrame(r).to_excel(
+                os.path.join(args.out, "Prompts_Folder", f"{split}_text.xlsx"),
+                index=False)
+    print(f"\nDone. Train={len(rows['Train'])} Val={len(rows['Val'])} "
+          f"Test={len(rows['Test'])} -> {args.out}")
+    print("Now fine-tune (GPU): cd reference_medclipseg && "
+          "python train.py --config-file configs/Neurofilament.yaml")
+
+
+if __name__ == "__main__":
+    main()

processing.py

@@ -25,6 +25,18 @@ from skimage.filters import gaussian, threshold_otsu
 from skimage.morphology import remove_small_objects, skeletonize
 from skan import Skeleton, summarize
 
+
+def _cellpose_available() -> bool:
+    try:
+        import cellpose  # noqa: F401
+        return True
+    except Exception:
+        return False
+
+
+CELLPOSE_AVAILABLE = _cellpose_available()
+_CP_MODEL = None       # lazily-created, cached Cellpose model
+
 # --------------------------------------------------------------------------- #
 #  Data containers
 # --------------------------------------------------------------------------- #
@@ -59,19 +71,38 @@ class RegionMetrics:
     area_covered_um2: float = 0.0
     fov_area_um2: float = 0.0
     pct_area_covered: float = 0.0
+    n_hair_cells: int = -1         # -1 = not measured (no Myo7a detection run)
+    fibers_per_hc: float = -1.0    # -1 = not applicable (hair-cell count unknown)
+    length_per_hc_um: float = -1.0
+    n_radial: int = -1             # fibers running IHC->OHC (radial); -1 = n/a
+    n_offaxis: int = -1            # fibers running along the rows (misdirected)
+    pct_radial: float = -1.0       # % of fibers that are radially directed
 
     def as_row(self) -> dict:
-        return {
+        def opt(value, ndigits=2):
+            """Blank when the metric was not applicable (negative sentinel)."""
+            return round(value, ndigits) if value is not None and value >= 0 else ""
+
+        row = {
             "Region": self.region,
             "Number of fibers": self.n_fibers,
+            "Hair cells (Myo7a)": (self.n_hair_cells
+                                   if self.n_hair_cells >= 0 else ""),
+            "Fibers / hair cell": opt(self.fibers_per_hc, 3),
             "Total length (um)": round(self.total_length_um, 2),
+            "Length / hair cell (um)": opt(self.length_per_hc_um, 2),
             "Mean diameter (um)": round(self.mean_diameter_um, 3),
             "Median diameter (um)": round(self.median_diameter_um, 3),
             "Branch points": self.n_branch_points,
+            "Radial fibers (IHC->OHC)": (self.n_radial
+                                         if self.n_radial >= 0 else ""),
+            "Off-axis fibers": (self.n_offaxis if self.n_offaxis >= 0 else ""),
+            "Radial fibers (%)": opt(self.pct_radial, 1),
             "Area covered (um^2)": round(self.area_covered_um2, 2),
             "FOV area (um^2)": round(self.fov_area_um2, 2),
             "Area covered (% of FOV)": round(self.pct_area_covered, 2),
         }
+        return row
 
 
 @dataclass
@@ -165,30 +196,77 @@ def load_czi(path: str) -> LoadedImage:
     return LoadedImage(arr, channels, _czi_voxel(czi), os.path.basename(path))
 
 
+def _reorder_tiff_by_axes(arr: np.ndarray, axes: str):
+    """Reorder a TIFF array to (C, Z, Y, X) using tifffile's axes string.
+
+    Treats colour samples ('S', e.g. RGB) and channels ('C') as channels, keeps
+    'Z', and squeezes any other axes (T, etc.).  Returns None if it cannot map.
+    """
+    axes = axes.upper()
+    if "Y" not in axes or "X" not in axes or len(axes) != arr.ndim:
+        return None
+    # Rank: channel-like first, then Z, then Y, X; unknown axes go last.
+    prio = {"C": 0, "S": 0, "Z": 1, "Y": 2, "X": 3}
+    order = sorted(range(len(axes)), key=lambda i: (prio.get(axes[i], 4), i))
+    arr = np.transpose(arr, order)
+    new_axes = "".join(axes[i] for i in order)
+    # Collapse all leading channel-like axes (C and/or S) into one channel axis.
+    n_chan = sum(1 for c in new_axes if c in ("C", "S"))
+    if n_chan >= 1:
+        arr = arr.reshape((int(np.prod(arr.shape[:n_chan])),) + arr.shape[n_chan:])
+        new_axes = "C" + new_axes[n_chan:]
+    else:
+        arr = arr[None]
+        new_axes = "C" + new_axes
+    if "Z" not in new_axes:                 # single-plane image -> add Z=1
+        arr = arr[:, None]
+        new_axes = new_axes[0] + "Z" + new_axes[1:]
+    while arr.ndim > 4:                      # drop any leftover trailing axes
+        arr = arr[..., 0]
+    return arr
+
+
 def load_tiff(path: str, dz=0.35, dy=0.0895, dx=0.0895) -> LoadedImage:
     import tifffile
 
-    arr = tifffile.imread(path)
-    arr = np.squeeze(arr)
-    # Heuristic to reach (C, Z, Y, X). The two largest axes are Y, X.
-    if arr.ndim == 2:                       # (Y, X) single channel, single plane
-        arr = arr[None, None]
-    elif arr.ndim == 3:
-        # Could be (Z,Y,X) single channel or (C,Y,X). Assume small first axis = C
-        if arr.shape[0] <= 5:
-            arr = arr[:, None]              # (C, 1, Y, X)
-        else:
-            arr = arr[None]                 # (1, Z, Y, X)
-    elif arr.ndim == 4:
-        # find the two largest axes -> Y, X; of the remaining two the smaller = C
-        yx = sorted(range(4), key=lambda a: arr.shape[a])[-2:]
-        rest = [a for a in range(4) if a not in yx]
-        c_axis = min(rest, key=lambda a: arr.shape[a])
-        z_axis = [a for a in rest if a != c_axis][0]
-        arr = np.transpose(arr, (c_axis, z_axis, *sorted(yx)))
+    with tifffile.TiffFile(path) as tf:
+        arr = tf.asarray()
+        try:
+            axes = tf.series[0].axes
+        except Exception:
+            axes = None
+
+    arr = np.asarray(arr)
+    reordered = _reorder_tiff_by_axes(arr, axes) if axes else None
+
+    if reordered is not None:
+        arr = reordered
     else:
-        raise ValueError(f"Unsupported TIFF with {arr.ndim} dimensions")
-    arr = arr.astype(np.float32)
+        # Fallback heuristic when axis metadata is missing/unusable.
+        arr = np.squeeze(arr)
+        if arr.ndim == 2:                          # (Y, X)
+            arr = arr[None, None]
+        elif arr.ndim == 3:
+            # trailing RGB/RGBA samples -> channels; else small first axis = C
+            if arr.shape[-1] in (3, 4) and arr.shape[-1] < min(arr.shape[:2]):
+                arr = np.moveaxis(arr, -1, 0)[:, None]   # (C, 1, Y, X)
+            elif arr.shape[0] <= 5:
+                arr = arr[:, None]                       # (C, 1, Y, X)
+            else:
+                arr = arr[None]                          # (1, Z, Y, X)
+        elif arr.ndim == 4:
+            if arr.shape[-1] in (3, 4) and arr.shape[-1] < min(arr.shape[1:3]):
+                arr = np.moveaxis(arr, -1, 0)            # (C, Z, Y, X)
+            else:
+                yx = sorted(range(4), key=lambda a: arr.shape[a])[-2:]
+                rest = [a for a in range(4) if a not in yx]
+                c_axis = min(rest, key=lambda a: arr.shape[a])
+                z_axis = [a for a in rest if a != c_axis][0]
+                arr = np.transpose(arr, (c_axis, z_axis, *sorted(yx)))
+        else:
+            raise ValueError(f"Unsupported TIFF with {arr.ndim} dimensions")
+
+    arr = np.ascontiguousarray(arr).astype(np.float32)
     channels = [{"name": f"Channel {i}", "dye": None, "color": None}
                 for i in range(arr.shape[0])]
     return LoadedImage(arr, channels, (dz, dy, dx), os.path.basename(path))
@@ -203,31 +281,164 @@ def load_image(path: str, dz=0.35, dy=0.0895, dx=0.0895) -> LoadedImage:
     raise ValueError(f"Unsupported file type: {ext}")
 
 
+# --------------------------------------------------------------------------- #
+#  IMARIS ground-truth ingestion
+# --------------------------------------------------------------------------- #
+
+def read_imaris_filament_length(source: str) -> Optional[float]:
+    """Return IMARIS 'Filament Length (sum)' in microns for an image, or None.
+
+    IMARIS exports this exact statistic to
+    ``<name>_Statistics/<name>_Filament_Length_(sum).csv`` beside the image.
+    ``source`` may be that CSV, a ``_Statistics`` folder, or the image path
+    (``.czi``/``.tif``) whose sibling ``_Statistics`` folder is searched.  When a
+    filament CSV lists several filament objects the per-object sums are added.
+
+    This is the ONLY way to reproduce IMARIS's exact number — it reads IMARIS's
+    own output rather than re-estimating it with a different algorithm.
+    """
+    import glob
+    import re
+
+    csvs = []
+    if source and os.path.isfile(source) and source.lower().endswith(".csv"):
+        csvs = [source]
+    else:
+        cand = None
+        if source and source.lower().endswith((".czi", ".tif", ".tiff")):
+            stem = os.path.splitext(os.path.basename(source))[0]
+            cand = os.path.join(os.path.dirname(source), f"{stem}_Statistics")
+        elif source and os.path.isdir(source):
+            cand = source
+        if cand and os.path.isdir(cand):
+            csvs = sorted(glob.glob(
+                os.path.join(cand, "*_Filament_Length_(sum).csv")))
+    if not csvs:
+        return None
+
+    total, found = 0.0, False
+    for csv in csvs[:1]:                       # one filament CSV per image
+        try:
+            with open(csv, errors="ignore") as fh:
+                for line in fh:
+                    m = re.match(r"\s*([-+]?[0-9]*\.?[0-9]+)\s*,\s*µm", line)
+                    if m:
+                        total += float(m.group(1))
+                        found = True
+        except OSError:
+            return None
+    return round(total, 3) if found else None
+
+
 # --------------------------------------------------------------------------- #
 #  Channel identification
 # --------------------------------------------------------------------------- #
 
-# Wavelength-based hints: neurofilament here is Alexa-555 (red/green range),
-# Myo7a is Alexa-405 (blue). Transmitted-light PMT channels have no dye.
+# Common fluorophores in these cochlear stacks.  Each entry maps a token that
+# may appear in the CZI dye/fluor name to a friendly name and a colour family.
+# Emission colour drives the family; the family drives auto channel assignment.
+#   blue   -> reference marker (Myo7a is Alexa-405 here)
+#   orange -> neurofilament (Alexa-555 here) — preferred trace channel
+#   green / red -> additional markers, selectable in the UI.
+KNOWN_DYES = [
+    ("atto 390", "ATTO 390", "blue"),
+    ("405", "Alexa Fluor 405", "blue"),
+    ("dapi", "DAPI", "blue"),
+    ("488", "Alexa Fluor 488", "green"),
+    ("atto 488", "ATTO 488", "green"),
+    ("fitc", "FITC", "green"),
+    ("gfp", "GFP", "green"),
+    ("514", "Alexa Fluor 514", "green"),
+    ("532", "Alexa Fluor 532", "orange"),
+    ("546", "Alexa Fluor 546", "orange"),
+    ("555", "Alexa Fluor 555", "orange"),
+    ("568", "Alexa Fluor 568", "orange"),
+    ("cy3", "Cy3", "orange"),
+    ("594", "Alexa Fluor 594", "red"),
+    ("633", "Alexa Fluor 633", "red"),
+    ("647", "Alexa Fluor 647", "red"),
+    ("atto 647", "ATTO 647", "red"),
+    ("cy5", "Cy5", "red"),
+    ("680", "Alexa Fluor 680", "red"),
+]
+
+
+def _channel_color_family(ch: dict) -> str:
+    """Classify a channel as blue / green / orange / red / none (transmitted).
+
+    Uses the dye/fluor name first, then the excitation wavelength, then the
+    metadata colour swatch.  This lets the app recognise many more dyes
+    (405/488/514/532/555/568/594/633/647/ATTO/Cy…) than just 405 vs "any dye".
+    """
+    dye = (ch.get("dye") or "").lower()
+    color = (ch.get("color") or "").upper().lstrip("#")
+    ex = ch.get("ex")
+    try:
+        ex = float(ex) if ex else None
+    except (TypeError, ValueError):
+        ex = None
+
+    for token, _name, family in KNOWN_DYES:
+        if token in dye:
+            return family
+    if ex is not None:
+        if ex < 430:
+            return "blue"
+        if ex < 505:
+            return "green"
+        if ex < 565:
+            return "orange"
+        return "red"
+    # metadata colour swatch (Zeiss stores ARGB or RGB hex)
+    if color:
+        hexrgb = color[-6:] if len(color) >= 6 else color
+        try:
+            r = int(hexrgb[0:2], 16); g = int(hexrgb[2:4], 16); b = int(hexrgb[4:6], 16)
+            if b > r and b > g:
+                return "blue"
+            if g > r and g > b:
+                return "green"
+            if r > g and r > b and g > b // 2:
+                return "orange"
+            if r > g and r > b:
+                return "red"
+        except (ValueError, IndexError):
+            pass
+    if dye and dye not in ("", "none"):
+        return "other"
+    return "none"     # e.g. transmitted-light PMT
+
+
+def channel_dye_label(ch: dict) -> str:
+    """Friendly dye label for the channel dropdown (recognises known dyes)."""
+    dye = (ch.get("dye") or "").strip()
+    if dye:
+        low = dye.lower()
+        for token, name, _family in KNOWN_DYES:
+            if token in low:
+                return name
+        return dye
+    fam = _channel_color_family(ch)
+    return "no dye / transmitted" if fam == "none" else f"{fam} marker"
+
+
 def guess_channels(img: LoadedImage) -> tuple:
     """Best-effort (neurofilament_index, myo7a_index) from metadata + content."""
     nf_idx, myo_idx = None, None
-    blue_like, red_like, plain = [], [], []
-    for i, ch in enumerate(img.channels):
-        dye = (ch.get("dye") or "").lower()
-        color = (ch.get("color") or "").upper()
-        ex = ch.get("ex")
-        ex = float(ex) if ex else None
-        if "405" in dye or (ex and ex < 430) or color == "#0000FF":
-            blue_like.append(i)
-        elif dye and dye not in ("", "none"):
-            red_like.append(i)
-        else:
-            plain.append(i)     # e.g. transmitted-light PMT
+    families = [_channel_color_family(ch) for ch in img.channels]
+    blue_like = [i for i, f in enumerate(families) if f == "blue"]
+    orange_like = [i for i, f in enumerate(families) if f == "orange"]
+    red_like = [i for i, f in enumerate(families) if f == "red"]
+    green_like = [i for i, f in enumerate(families) if f == "green"]
+
+    # Myo7a reference = the bluest (405-like) fluorescent channel.
     if blue_like:
         myo_idx = blue_like[0]
-    if red_like:
-        nf_idx = red_like[0]
+    # Neurofilament = Alexa-555 (orange) here; then fall back to red/green.
+    for group in (orange_like, red_like, green_like):
+        if group:
+            nf_idx = next((i for i in group if i != myo_idx), group[0])
+            break
 
     # Fall back to image content when metadata is missing (e.g. plain TIFF).
     # Fluorescence channels have a mostly-dark background; transmitted-light
@@ -270,18 +481,89 @@ def _threshold_volume(vol: np.ndarray, sensitivity: float) -> np.ndarray:
     return sm > thr
 
 
+def _remove_small_components(skel: np.ndarray, voxel: tuple,
+                             min_len_um: float) -> np.ndarray:
+    """Drop connected skeleton components shorter than ``min_len_um``."""
+    if min_len_um <= 0 or skel.sum() == 0:
+        return skel
+    dz, dy, dx = voxel
+    step = float(np.mean([dy, dx]))              # ~ length of one skeleton step
+    min_vox = max(2, int(round(min_len_um / step)))
+    struct = np.ones((3,) * skel.ndim, int)
+    lbl, n = ndi.label(skel, structure=struct)
+    if n == 0:
+        return skel
+    counts = np.bincount(lbl.ravel())
+    keep = np.zeros(counts.size, bool)
+    keep[1:] = counts[1:] >= min_vox
+    return keep[lbl]
+
+
+def prune_skeleton(skel: np.ndarray, voxel: tuple,
+                   spur_um: float = 3.0, min_component_um: Optional[float] = None,
+                   iterations: int = 4) -> np.ndarray:
+    """Remove short terminal spurs and small isolated fragments.
+
+    Terminal twigs (junction→endpoint branches) shorter than ``spur_um`` are
+    trimmed while their junction pixel is preserved, so the backbone stays
+    connected.  Isolated components shorter than ``min_component_um`` are dropped
+    entirely.  Iterated a few times because trimming one spur can expose another.
+    """
+    if spur_um <= 0 and (min_component_um or 0) <= 0:
+        return skel
+    if min_component_um is None:
+        min_component_um = spur_um
+    cur = skel.copy()
+    for _ in range(iterations):
+        if cur.sum() == 0 or spur_um <= 0:
+            break
+        try:
+            S = Skeleton(cur, spacing=voxel)
+            df = summarize(S, separator="_")
+        except Exception:
+            break
+        k = np.ones((3,) * cur.ndim, int)
+        nb = ndi.convolve(cur.astype(np.uint8), k, mode="constant") - cur
+        junction = cur & (nb > 2)               # protect branch points
+        drop = df.index[
+            ((df["branch_type"] == 1) & (df["branch_distance"] < spur_um)) |
+            (df["branch_type"].isin([0, 3]) &
+             (df["branch_distance"] < min_component_um))]
+        if len(drop) == 0:
+            break
+        removal = np.zeros_like(cur)
+        for i in drop:
+            coords = S.path_coordinates(i).astype(int)
+            removal[tuple(coords.T)] = True
+        removal &= ~junction
+        new = cur & ~removal
+        if int(new.sum()) == int(cur.sum()):
+            break
+        cur = new
+    return _remove_small_components(cur, voxel, min_component_um)
+
+
 def trace_neurites(nf_vol: np.ndarray, voxel: tuple,
                    sensitivity: float = 1.0,
-                   min_object_vox: int = 64) -> TraceResult:
-    """Segment and skeletonise the neurofilament network in 3D."""
+                   min_object_vox: int = 64,
+                   prune_um: float = 3.0) -> TraceResult:
+    """Segment and skeletonise the neurofilament network in 3D.
+
+    ``prune_um`` removes terminal spurs and isolated fragments shorter than this
+    many microns from the skeleton, giving cleaner fiber lines (0 = no pruning).
+    """
     dz, dy, dx = voxel
     mask = _threshold_volume(nf_vol, sensitivity)
     mask = remove_small_objects(mask, min_object_vox)
-    mask = ndi.binary_closing(mask, iterations=1)
+    # Close within each z-plane. A full 3D element would erode a single-plane
+    # (Z=1) stack to nothing, since its z-neighbours fall outside the array.
+    mask = ndi.binary_closing(mask, structure=np.ones((1, 3, 3), bool))
     if mask.sum() == 0:
         z = np.zeros_like(mask)
         return TraceResult(mask, z, np.zeros(mask.shape, np.float32), voxel)
     skel = skeletonize(mask)
+    if prune_um and prune_um > 0:
+        skel = prune_skeleton(skel, voxel, spur_um=float(prune_um))
     dist = ndi.distance_transform_edt(mask, sampling=(dz, dy, dx)).astype(np.float32)
     return TraceResult(mask, skel, dist, voxel)
 
@@ -312,25 +594,358 @@ def suggest_boundary(myo_vol: np.ndarray) -> float:
     return centroid / prof.size
 
 
-def make_region_masks(shape_yx: tuple, boundary_frac: float,
-                      ihc_side: str = "low", axis: str = "Y") -> tuple:
-    """Return (ihc_roi, ohc_roi) boolean 2D masks split by a straight line.
+RADIAL_AXIS_INDEX = {"Z": 0, "Y": 1, "X": 2}
+
 
-    axis="Y" splits along rows (radial axis, the usual case); axis="X" splits
-    along columns.  ihc_side selects which side of the boundary is IHC.
+def boundary_line_coords(shape_yx: tuple, boundary_frac: float, axis: str = "Y",
+                         angle_deg: float = 0.0, curvature: float = 0.0) -> tuple:
+    """Pixel coordinates of the (possibly tilted / curved) boundary line.
+
+    Returns (ys, xs) integer arrays to index a 2D image, so the same line is
+    drawn on every overlay and used to build the region masks — they can never
+    disagree.  ``angle_deg`` tilts the line; ``curvature`` bows it (a parabola,
+    positive = bow toward higher coordinates in the middle).  For axis="Z" there
+    is no in-plane line, so empty arrays are returned.
     """
     ny, nx = shape_yx
-    low = np.zeros((ny, nx), bool)
     if axis.upper() == "Y":
-        b = int(round(np.clip(boundary_frac, 0, 1) * ny))
+        xs = np.arange(nx)
+        cx, hx = nx / 2.0, max(nx / 2.0, 1.0)
+        ys = (boundary_frac * ny + np.tan(np.deg2rad(angle_deg)) * (xs - cx)
+              + curvature * ((xs - cx) / hx) ** 2 * ny)
+        ys = np.clip(np.round(ys), 0, ny - 1).astype(int)
+        return ys, xs
+    if axis.upper() == "X":
+        ys = np.arange(ny)
+        cy, hy = ny / 2.0, max(ny / 2.0, 1.0)
+        xs = (boundary_frac * nx + np.tan(np.deg2rad(angle_deg)) * (ys - cy)
+              + curvature * ((ys - cy) / hy) ** 2 * nx)
+        xs = np.clip(np.round(xs), 0, nx - 1).astype(int)
+        return ys, xs
+    return np.array([], int), np.array([], int)   # axis == "Z"
+
+
+def make_region_masks(shape: tuple, boundary_frac: float,
+                      ihc_side: str = "low", axis: str = "Y",
+                      angle_deg: float = 0.0, curvature: float = 0.0) -> tuple:
+    """Return (ihc_roi, ohc_roi) boolean masks split by the boundary.
+
+    ``shape`` may be (ny, nx) or (nz, ny, nx).  axis="Y" splits along rows
+    (radial axis, the usual case), axis="X" splits along columns; both accept a
+    tilt (``angle_deg``) and a bow (``curvature``) and return 2D masks that are
+    broadcast over Z during quantification.  axis="Z" splits by depth and
+    returns 3D masks.  ``ihc_side`` selects which side ("low" = smaller
+    coordinate) is IHC.
+    """
+    axis = axis.upper()
+    if axis == "Z":
+        if len(shape) != 3:
+            raise ValueError("Z split needs a 3D (nz, ny, nx) shape")
+        nz, ny, nx = shape
+        if nz < 2:
+            raise ValueError("Z (depth) split needs a multi-plane stack "
+                             "(nz >= 2); this image has a single plane.")
+        b = int(round(np.clip(boundary_frac, 0, 1) * nz))
+        low = np.zeros((nz, ny, nx), bool)
         low[:b] = True
-    else:
-        b = int(round(np.clip(boundary_frac, 0, 1) * nx))
-        low[:, :b] = True
+        ihc = low if ihc_side == "low" else ~low
+        return ihc, ~ihc
+
+    ny, nx = shape[-2], shape[-1]
+    yy, xx = np.mgrid[0:ny, 0:nx]
+    ang = np.tan(np.deg2rad(angle_deg))
+    if axis == "Y":
+        cx, hx = nx / 2.0, max(nx / 2.0, 1.0)
+        yline = (boundary_frac * ny + ang * (xx - cx)
+                 + curvature * ((xx - cx) / hx) ** 2 * ny)
+        low = yy < yline
+    else:                                        # axis == "X"
+        cy, hy = ny / 2.0, max(ny / 2.0, 1.0)
+        xline = (boundary_frac * nx + ang * (yy - cy)
+                 + curvature * ((yy - cy) / hy) ** 2 * nx)
+        low = xx < xline
     ihc = low if ihc_side == "low" else ~low
     return ihc, ~ihc
 
 
+# --------------------------------------------------------------------------- #
+#  Hair-cell detection (Cellpose assist, with a classical fallback)
+# --------------------------------------------------------------------------- #
+
+# A mouse cochlear hair cell body is roughly this wide; used to pick the
+# working scale so cells land near Cellpose's preferred pixel size.
+HAIR_CELL_DIAMETER_UM = 8.0
+_CP_TARGET_PX = 30
+
+
+def _norm(mip: np.ndarray) -> np.ndarray:
+    lo, hi = np.percentile(mip, 1), np.percentile(mip, 99.5)
+    return np.clip((mip - lo) / (hi - lo + 1e-6), 0, 1).astype(np.float32)
+
+
+def custom_cellpose_model_path() -> Optional[str]:
+    """Path to a fine-tuned Cellpose hair-cell model, if one is provided.
+
+    Looked up in order:
+      1. the ``NEURON_TRACER_CP_MODEL`` environment variable, then
+      2. a bundled ``models/hair_cell_cpsam`` file next to this module.
+    Returns None to use the stock pretrained Cellpose model.  This is how a
+    model trained with ``training/train_hair_cells.py`` gets picked up by the
+    app: drop it at ``models/hair_cell_cpsam`` (or point the env var at it).
+    """
+    p = os.environ.get("NEURON_TRACER_CP_MODEL")
+    if p and os.path.exists(p):
+        return p
+    here = os.path.dirname(os.path.abspath(__file__))
+    bundled = os.path.join(here, "models", "hair_cell_cpsam")
+    return bundled if os.path.exists(bundled) else None
+
+
+def _get_cellpose_model():
+    global _CP_MODEL
+    if _CP_MODEL is None:
+        from cellpose import models
+        mp = custom_cellpose_model_path()
+        if mp:
+            _CP_MODEL = models.CellposeModel(gpu=False, pretrained_model=mp)
+        else:
+            _CP_MODEL = models.CellposeModel(gpu=False)
+    return _CP_MODEL
+
+
+def prep_detection_image(myo_vol: np.ndarray, voxel: tuple) -> tuple:
+    """The exact Myo7a image the detector feeds to Cellpose, plus its scale.
+
+    Returns ``(img_uint8, scale)`` where ``img_uint8`` is the contrast-
+    normalised max-projection rescaled so hair cells sit near Cellpose's
+    working size, and ``scale`` maps its coordinates back to full resolution
+    (full = detected / scale).  Exposed so ``training/prepare_data.py`` can
+    export annotation images at the *same* scale the model sees at inference —
+    keeping training and inference consistent.
+    """
+    dx = voxel[2]
+    mip = _norm(myo_vol.max(0).astype(np.float32))
+    cell_px_full = HAIR_CELL_DIAMETER_UM / dx           # e.g. ~89 px
+    scale = float(np.clip(_CP_TARGET_PX / cell_px_full, 0.2, 1.0))
+    from skimage.transform import rescale
+    small = (rescale(mip, scale, anti_aliasing=True, preserve_range=True
+                     ).astype(np.float32) if scale < 0.999 else mip)
+    small = _norm(small)
+    return (small * 255).astype(np.uint8), scale
+
+
+def _watershed_cells(img: np.ndarray, cell_px: float) -> np.ndarray:
+    """Classical blob segmentation fallback (no deep-learning dependency)."""
+    from skimage.feature import peak_local_max
+    from skimage.segmentation import watershed
+    sm = gaussian(img, max(cell_px / 6.0, 1.0), preserve_range=True)
+    try:
+        mask = sm > threshold_otsu(sm)
+    except Exception:
+        return np.zeros(img.shape, int)
+    mask = ndi.binary_opening(mask, iterations=1)
+    dist = ndi.distance_transform_edt(mask)
+    coords = peak_local_max(dist, min_distance=max(int(cell_px * 0.5), 3),
+                            labels=mask)
+    if len(coords) == 0:
+        return np.zeros(img.shape, int)
+    markers = np.zeros(img.shape, int)
+    markers[tuple(coords.T)] = np.arange(1, len(coords) + 1)
+    return watershed(-dist, markers, mask=mask)
+
+
+def detect_hair_cells(myo_vol: np.ndarray, voxel: tuple,
+                      prefer_cellpose: bool = True) -> dict:
+    """Detect Myo7a hair-cell bodies on the max-projection.
+
+    Uses Cellpose when available (better on touching cells), otherwise a
+    watershed fallback.  If a fine-tuned model is provided (see
+    ``custom_cellpose_model_path``) it is used automatically and the engine is
+    reported as ``cellpose-custom``.  Returns full-resolution centroids plus the
+    count and which engine ran.  The image is rescaled so cells are ~30 px,
+    which is where Cellpose performs best.
+    """
+    mip_shape = myo_vol.max(0).shape
+    small, scale = prep_detection_image(myo_vol, voxel)
+    small = small.astype(np.float32) / 255.0
+
+    engine = "watershed"
+    masks = None
+    if prefer_cellpose and CELLPOSE_AVAILABLE:
+        try:
+            masks = _get_cellpose_model().eval(small, diameter=_CP_TARGET_PX)[0]
+            engine = ("cellpose-custom" if custom_cellpose_model_path()
+                      else "cellpose")
+        except Exception:
+            masks = None
+    if masks is None:
+        masks = _watershed_cells(small, _CP_TARGET_PX)
+
+    n = int(masks.max())
+    if n:
+        cent_small = np.array(ndi.center_of_mass(
+            np.ones_like(masks), masks, range(1, n + 1)))
+        centroids = cent_small / scale                  # back to full-res Y,X
+    else:
+        centroids = np.zeros((0, 2))
+    return {"centroids": centroids, "count": n, "engine": engine,
+            "scale": scale, "mip_shape": mip_shape}
+
+
+# Anatomy-based thresholds for accepting an IHC/OHC split (mouse organ of Corti).
+_MIN_TWO_ROW_EXTENT_UM = 18.0   # IHC row + tunnel + OHC rows span well past this;
+                                # a single hair-cell row spans < ~15 µm radially.
+_MIN_TUNNEL_UM = 7.0            # the IHC/OHC gap (tunnel of Corti) is ~10-25 µm;
+                                # inter-cell spacing within a row is smaller.
+
+
+def auto_regions_from_cells(centroids: np.ndarray, shape_yx: tuple,
+                            um_per_px: Optional[float] = None) -> dict:
+    """Suggest an IHC/OHC boundary from hair-cell centroids.
+
+    IHCs form a single row and OHCs form three rows separated from the IHCs by
+    the tunnel of Corti.  We project cells onto the radial axis (perpendicular
+    to the hair-cell band), find the widest cell-free gap that leaves at least
+    two cells on each side, and call the sparser/tighter side IHC.
+
+    Crucially, a split is only proposed when the cells actually span more than
+    one row (radial extent) *and* the gap is a real cell-free tunnel.  When only
+    a single hair-cell row is present, ``single_row`` is returned True and no
+    split is made — the field should be reported whole-field, not forced into a
+    meaningless IHC/OHC division.
+    """
+    ny, nx = shape_yx
+    result = {"boundary_frac": 0.5, "ihc_side": "low", "confidence": "low",
+              "n_ihc_cells": 0, "n_ohc_cells": 0, "single_row": False,
+              "reason": "not enough hair cells for an automatic split"}
+    if len(centroids):
+        result["boundary_frac"] = float(np.clip(centroids[:, 0].mean() / ny, 0, 1))
+    if len(centroids) < 6:
+        result["single_row"] = True
+        result["reason"] = (f"only {len(centroids)} hair cell(s) detected — "
+                            "cannot separate IHC vs OHC; report whole-field")
+        return result
+
+    c = centroids - centroids.mean(0)
+    _, _, vt = np.linalg.svd(c, full_matrices=False)
+    radial = vt[1]
+    if radial[0] < 0:
+        radial = -radial                                # point toward +Y
+    r = c @ radial
+    order = np.argsort(r)
+    rs = r[order]
+    gaps = np.diff(rs)
+
+    # Convert radial distances to microns (or fall back to image-fraction units).
+    if um_per_px and um_per_px > 0:
+        to_um = um_per_px
+        extent = (rs[-1] - rs[0]) * to_um
+        min_extent, min_gap = _MIN_TWO_ROW_EXTENT_UM, _MIN_TUNNEL_UM
+        unit = "µm"
+    else:                                               # unknown scale
+        to_um = 1.0
+        extent = (rs[-1] - rs[0]) / max(ny, 1)          # fraction of image
+        min_extent, min_gap = 0.12, 0.04
+        unit = "frac"
+
+    # Candidate splits must leave >=2 cells on each side (ignore stray outliers).
+    valid = [(i, gaps[i]) for i in range(1, len(gaps) - 1)]
+    biggest_gap = (max(valid, key=lambda t: t[1])[1] * to_um) if valid else 0.0
+
+    # Single-row / no-tunnel guard: refuse to split an unresolved band.
+    if extent < min_extent or not valid or biggest_gap < min_gap:
+        result["single_row"] = True
+        result["reason"] = (
+            f"hair cells span only {extent:.1f} {unit} radially with no clear "
+            f"tunnel gap — looks like a single row; IHC vs OHC cannot be "
+            f"separated from this field. Report whole-field or set the boundary "
+            f"by hand if you know the anatomy.")
+        return result
+
+    gi = max(valid, key=lambda t: t[1])[0]
+    split_r = (rs[gi] + rs[gi + 1]) / 2.0
+
+    left = r <= split_r
+    n_left, n_right = int(left.sum()), int((~left).sum())
+    s_left = float(r[left].std()) if n_left > 1 else 0.0
+    s_right = float(r[~left].std()) if n_right > 1 else 0.0
+    # IHC = single row: fewer cells and tighter spread
+    score = (1 if n_right > n_left else -1) + (1 if s_right > s_left else -1)
+    ihc_is_left = score >= 0
+    ihc_side = "low" if ihc_is_left else "high"
+    n_ihc = n_left if ihc_is_left else n_right
+    n_ohc = n_right if ihc_is_left else n_left
+
+    ymid = centroids[:, 0].mean() + split_r * radial[0]
+    bfrac = float(np.clip(ymid / ny, 0, 1))
+
+    ratio = n_ohc / max(n_ihc, 1)
+    strong_gap = (biggest_gap >= 2 * min_gap)
+    if strong_gap and extent >= 1.6 * min_extent and 1.8 <= ratio <= 5.0:
+        conf = "high"
+    else:
+        conf = "medium"
+
+    return {"boundary_frac": bfrac, "ihc_side": ihc_side, "confidence": conf,
+            "n_ihc_cells": n_ihc, "n_ohc_cells": n_ohc, "single_row": False,
+            "reason": f"tunnel gap {biggest_gap:.1f} {unit}, radial extent "
+                      f"{extent:.1f} {unit}, IHC:OHC cell ratio 1:{ratio:.1f}"}
+
+
+def count_cells_in_roi(centroids: np.ndarray, roi_yx: np.ndarray) -> int:
+    """Count hair-cell centroids falling inside a 2D ROI mask."""
+    if len(centroids) == 0:
+        return 0
+    ny, nx = roi_yx.shape
+    yy = np.clip(centroids[:, 0].round().astype(int), 0, ny - 1)
+    xx = np.clip(centroids[:, 1].round().astype(int), 0, nx - 1)
+    return int(roi_yx[yy, xx].sum())
+
+
+def _draw_boundary_rgb(rgb: np.ndarray, boundary_frac: Optional[float],
+                       axis: str = "Y", angle_deg: float = 0.0,
+                       curvature: float = 0.0, color=(255, 255, 0)) -> np.ndarray:
+    """Draw the (tilted / curved) boundary line onto an RGB image in place.
+
+    For axis="Z" there is no in-plane line, so nothing is drawn.
+    """
+    if boundary_frac is None:
+        return rgb
+    ny, nx = rgb.shape[:2]
+    ys, xs = boundary_line_coords((ny, nx), boundary_frac, axis, angle_deg, curvature)
+    if ys.size:
+        rgb[ys, xs] = color
+    return rgb
+
+
+def hair_cell_overlay(myo_mip_u8: np.ndarray, centroids: np.ndarray,
+                      boundary_frac: Optional[float] = None,
+                      axis: str = "Y", ihc_side: Optional[str] = None,
+                      angle_deg: float = 0.0, curvature: float = 0.0) -> np.ndarray:
+    """Myo7a MIP with detected hair cells marked and the boundary drawn.
+
+    Cells are coloured by the *actual* region mask (so tilt/curvature are
+    honoured): cyan = IHC region, magenta = OHC region, neutral green when no
+    split is defined.
+    """
+    rgb = np.stack([myo_mip_u8] * 3, axis=-1).copy()
+    ny, nx = myo_mip_u8.shape
+    ihc_mask = None
+    if (boundary_frac is not None and ihc_side is not None
+            and axis.upper() in ("Y", "X")):
+        ihc_mask, _ = make_region_masks((ny, nx), boundary_frac, ihc_side,
+                                        axis, angle_deg, curvature)
+    for (y, x) in centroids.astype(int):
+        yc, xc = min(max(y, 0), ny - 1), min(max(x, 0), nx - 1)
+        col = (0, 255, 0)
+        if ihc_mask is not None:
+            col = (0, 220, 255) if ihc_mask[yc, xc] else (255, 60, 200)
+        ys, xs = slice(max(0, y - 3), y + 4), slice(max(0, x - 3), x + 4)
+        rgb[ys, xs] = col
+    _draw_boundary_rgb(rgb, boundary_frac, axis, angle_deg, curvature)
+    return rgb
+
+
 # --------------------------------------------------------------------------- #
 #  Quantification
 # --------------------------------------------------------------------------- #
@@ -343,50 +958,166 @@ def _branch_point_count(skel: np.ndarray) -> int:
     return int((skel & (nb > 2)).sum())
 
 
+def _fiber_stats(skel: np.ndarray, distance_um: np.ndarray, voxel: tuple,
+                 min_fiber_um: float, min_diam_um: float, max_diam_um: float,
+                 radial_axis: Optional[int]) -> tuple:
+    """Per-fiber (connected-component) filtering and radial-direction counting.
+
+    A fiber is kept only if its length >= ``min_fiber_um`` **and** its mean
+    diameter is within [``min_diam_um``, ``max_diam_um``] (``max_diam_um`` <= 0
+    means no upper limit).  When ``radial_axis`` is given, each kept fiber is
+    classified as radial (its principal axis aligns with the IHC->OHC axis) or
+    off-axis (running along the rows — a "misdirected" fiber).
+
+    Returns (kept_skel_bool, n_fibers, total_length_um, n_radial, n_offaxis).
+    ``n_radial``/``n_offaxis`` are -1 when direction was not computed.
+    """
+    struct = np.ones((3,) * skel.ndim, int)
+    lbl, n = ndi.label(skel, structure=struct)
+    if n == 0:
+        return np.zeros_like(skel), 0, 0.0, -1, -1
+
+    # Length per component: skan gives spacing-aware branch lengths; map each
+    # branch to its connected component via a pixel lying on the branch.
+    comp_len = {}
+    try:
+        S = Skeleton(skel, spacing=tuple(voxel))
+        df = summarize(S, separator="_")
+        for pi in df.index:
+            coords = np.round(S.path_coordinates(pi)).astype(int)
+            for a in range(coords.shape[1]):
+                coords[:, a] = np.clip(coords[:, a], 0, skel.shape[a] - 1)
+            c = 0
+            for p in (coords[len(coords) // 2], coords[0], coords[-1]):
+                c = int(lbl[tuple(p)])
+                if c:
+                    break
+            if c:
+                comp_len[c] = comp_len.get(c, 0.0) + float(df.loc[pi, "branch_distance"])
+    except Exception:
+        step = float(np.mean(voxel))
+        for c in range(1, n + 1):
+            comp_len[c] = float((lbl == c).sum()) * step
+
+    spacing = np.asarray(voxel, float)
+    slices = ndi.find_objects(lbl)
+    kept = np.zeros(n + 1, bool)
+    n_radial = n_offaxis = 0
+    do_dir = radial_axis is not None
+    for c in range(1, n + 1):
+        sl = slices[c - 1]
+        if sl is None:
+            continue
+        # A component with no measured length is a degenerate fragment (e.g. an
+        # isolated voxel skan emits no branch for) — never count it as a fiber,
+        # even when min_fiber_um is 0.  This also guarantees kept[c] => c in
+        # comp_len, so the total_length sum below can index comp_len safely.
+        if c not in comp_len or comp_len[c] < min_fiber_um:
+            continue
+        sub = lbl[sl] == c
+        dsub = 2.0 * distance_um[sl][sub]
+        mean_d = float(dsub.mean()) if dsub.size else 0.0
+        if mean_d < min_diam_um or (max_diam_um > 0 and mean_d > max_diam_um):
+            continue
+        kept[c] = True
+        if do_dir:
+            offset = np.array([s.start for s in sl], float)
+            pts = (np.argwhere(sub).astype(float) + offset) * spacing
+            if len(pts) >= 2:
+                pc = pts - pts.mean(0)
+                _, _, vt = np.linalg.svd(pc, full_matrices=False)
+                principal = vt[0]
+                cos_r = abs(principal[radial_axis]) / (np.linalg.norm(principal) + 1e-9)
+                n_radial += int(cos_r >= 0.5)
+                n_offaxis += int(cos_r < 0.5)
+            else:
+                n_offaxis += 1
+
+    kept_skel = kept[lbl]
+    n_fibers = int(kept[1:].sum())
+    total_length = float(sum(comp_len.get(c, 0.0)
+                             for c in range(1, n + 1) if kept[c]))
+    if not do_dir:
+        n_radial = n_offaxis = -1
+    return kept_skel, n_fibers, total_length, n_radial, n_offaxis
+
+
 def compute_metrics(trace: TraceResult, region_name: str,
                     roi_yx: Optional[np.ndarray] = None,
-                    min_fiber_um: float = 5.0) -> RegionMetrics:
-    """Quantify the skeleton, optionally restricted to a 2D ROI (broadcast in Z)."""
+                    min_fiber_um: float = 5.0,
+                    hair_cell_centroids: Optional[np.ndarray] = None,
+                    min_diameter_um: float = 0.0,
+                    max_diameter_um: float = 0.0,
+                    manual_hair_cells: Optional[int] = None,
+                    radial_axis_name: str = "Y") -> RegionMetrics:
+    """Quantify the skeleton, optionally restricted to an ROI (2D or 3D).
+
+    Fibers are filtered by length (``min_fiber_um``) and, when set, by mean
+    diameter (``min_diameter_um`` / ``max_diameter_um``).  ``manual_hair_cells``
+    overrides the detected hair-cell count for normalisation; otherwise
+    ``hair_cell_centroids`` (if given) are counted inside the region.
+    ``radial_axis_name`` ("Y"/"X"/"Z") is the IHC->OHC axis used to classify
+    fiber direction.
+    """
     dz, dy, dx = trace.voxel
     skel = trace.skeleton
     mask = trace.mask
     if roi_yx is not None:
-        roi3d = np.broadcast_to(roi_yx, skel.shape)
+        roi3d = np.broadcast_to(roi_yx, skel.shape) if roi_yx.ndim == 2 else roi_yx
         skel = skel & roi3d
         mask = mask & roi3d
 
     m = RegionMetrics(region=region_name)
-    m.fov_area_um2 = float(roi_yx.sum() if roi_yx is not None
-                           else mask.shape[1] * mask.shape[2]) * dx * dy
+
+    # 2D footprint of the ROI drives area / FOV (handles 2D and 3D ROIs).
+    if roi_yx is not None:
+        roi_foot = roi_yx if roi_yx.ndim == 2 else roi_yx.max(0)
+        m.fov_area_um2 = float(roi_foot.sum()) * dx * dy
+    else:
+        roi_foot = None
+        m.fov_area_um2 = float(mask.shape[1] * mask.shape[2]) * dx * dy
+
+    # Hair-cell count: manual override wins, else count detected centroids.
+    if manual_hair_cells is not None and manual_hair_cells >= 0:
+        m.n_hair_cells = int(manual_hair_cells)
+    elif hair_cell_centroids is not None:
+        if roi_foot is not None:
+            m.n_hair_cells = count_cells_in_roi(hair_cell_centroids, roi_foot)
+        else:
+            m.n_hair_cells = int(len(hair_cell_centroids))
 
     foot = mask.max(0)
     m.area_covered_um2 = float(foot.sum()) * dx * dy
     m.pct_area_covered = (100.0 * m.area_covered_um2 / m.fov_area_um2
                           if m.fov_area_um2 else 0.0)
 
-    if skel.sum() < 2:
-        return m
-
-    m.n_branch_points = _branch_point_count(skel)
-
-    diam = 2.0 * trace.distance_um[skel]
-    if diam.size:
-        m.mean_diameter_um = float(diam.mean())
-        m.median_diameter_um = float(np.median(diam))
+    if skel.sum() >= 2:
+        name = (radial_axis_name or "Y").upper()
+        if skel.ndim == 3:
+            radial_axis = {"Z": 0, "Y": 1, "X": 2}.get(name, 1)
+        else:
+            radial_axis = {"Y": 0, "X": 1}.get(name, 0)
+        kept_skel, n_fibers, total_len, n_radial, n_offaxis = _fiber_stats(
+            skel, trace.distance_um, (dz, dy, dx), float(min_fiber_um),
+            float(min_diameter_um), float(max_diameter_um), radial_axis)
+        m.n_fibers = n_fibers
+        m.total_length_um = total_len
+        m.n_branch_points = _branch_point_count(kept_skel)
+        diam = 2.0 * trace.distance_um[kept_skel]
+        if diam.size:
+            m.mean_diameter_um = float(diam.mean())
+            m.median_diameter_um = float(np.median(diam))
+        if n_radial >= 0:
+            m.n_radial = n_radial
+            m.n_offaxis = n_offaxis
+            tot = n_radial + n_offaxis
+            m.pct_radial = (100.0 * n_radial / tot) if tot else 0.0
+
+    # Normalisation by hair-cell count (when known and positive).
+    if m.n_hair_cells > 0:
+        m.fibers_per_hc = m.n_fibers / m.n_hair_cells
+        m.length_per_hc_um = m.total_length_um / m.n_hair_cells
 
-    # Length and fiber count via skan (per connected skeleton component).
-    try:
-        S = Skeleton(skel, spacing=(dz, dy, dx))
-        df = summarize(S, separator="_")
-        comp_len = df.groupby("skeleton_id")["branch_distance"].sum()
-        kept = comp_len[comp_len >= min_fiber_um]
-        m.n_fibers = int(kept.size)
-        m.total_length_um = float(kept.sum())
-    except Exception:
-        # Fallback: label components, approximate length by voxel count.
-        lbl, n = ndi.label(skel, structure=np.ones((3, 3, 3)))
-        m.n_fibers = int(n)
-        m.total_length_um = float(skel.sum()) * np.mean([dz, dy, dx])
     return m
 
 
@@ -404,26 +1135,44 @@ def skeleton_image(skel: np.ndarray, dilate: int = 1) -> np.ndarray:
 
 def region_overlay(skel: np.ndarray, ihc_roi: np.ndarray, ohc_roi: np.ndarray,
                    boundary_frac: float, axis: str = "Y",
+                   angle_deg: float = 0.0, curvature: float = 0.0,
                    dilate: int = 1) -> np.ndarray:
-    """Colour-coded RGB preview: IHC fibers cyan, OHC fibers magenta,
-    with the boundary line drawn in yellow."""
-    flat = skel.max(0) if skel.ndim == 3 else skel
+    """Colour-coded RGB preview: IHC fibers cyan, OHC fibers magenta, with the
+    boundary line in yellow.  Works with 2D (broadcast) or 3D (Z-split) ROIs."""
+    if skel.ndim == 3:
+        ihc3 = (np.broadcast_to(ihc_roi, skel.shape) if ihc_roi.ndim == 2
+                else ihc_roi)
+        ohc3 = (np.broadcast_to(ohc_roi, skel.shape) if ohc_roi.ndim == 2
+                else ohc_roi)
+        ihc_pix = (skel & ihc3).max(0)
+        ohc_pix = (skel & ohc3).max(0)
+        ny, nx = skel.shape[1], skel.shape[2]
+    else:
+        ihc_pix = skel & ihc_roi
+        ohc_pix = skel & ohc_roi
+        ny, nx = skel.shape
     if dilate:
-        flat = ndi.binary_dilation(flat, iterations=dilate)
-    ny, nx = flat.shape
+        ihc_pix = ndi.binary_dilation(ihc_pix, iterations=dilate)
+        ohc_pix = ndi.binary_dilation(ohc_pix, iterations=dilate)
     rgb = np.zeros((ny, nx, 3), np.uint8)
-    ihc_pix = flat & ihc_roi
-    ohc_pix = flat & ohc_roi
     rgb[ihc_pix] = (0, 220, 255)      # cyan  = IHC
     rgb[ohc_pix] = (255, 60, 200)     # magenta = OHC
-    if axis.upper() == "Y":
-        b = int(round(np.clip(boundary_frac, 0, 1) * ny))
-        b = min(max(b, 0), ny - 1)
-        rgb[b, :] = (255, 255, 0)
-    else:
-        b = int(round(np.clip(boundary_frac, 0, 1) * nx))
-        b = min(max(b, 0), nx - 1)
-        rgb[:, b] = (255, 255, 0)
+    _draw_boundary_rgb(rgb, boundary_frac, axis, angle_deg, curvature)
+    return rgb
+
+
+def overlay_on_original(base_gray_u8: np.ndarray, skel: np.ndarray,
+                        color=(255, 70, 70), dilate: int = 1) -> np.ndarray:
+    """Traced skeleton drawn in colour over the original (grayscale) MIP.
+
+    Lets the trace be checked against the raw signal — the reviewers' request
+    to view the tracing overlaid on the original image.
+    """
+    flat = skel.max(0) if skel.ndim == 3 else skel
+    if dilate:
+        flat = ndi.binary_dilation(flat, iterations=dilate)
+    rgb = np.stack([base_gray_u8] * 3, axis=-1).copy()
+    rgb[flat] = color
     return rgb
 
 

reference_medclipseg/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

reference_medclipseg/README.md

@@ -0,0 +1,141 @@
+<div align="center">
+  
+# MedCLIPSeg: Probabilistic Vision–Language Adaptation for Data-Efficient and Generalizable Medical Image Segmentation
+
+<h3>CVPR 2026</h3>
+
+**[Health-X Lab](http://www.healthx-lab.ca/)** | **[IMPACT Lab](https://users.encs.concordia.ca/~impact/)** 
+
+[Taha Koleilat](https://tahakoleilat.github.io/), 
+[Hojat Asgariandehkordi](https://scholar.google.com/citations?user=ndXNye4AAAAJ&hl=en), 
+[Omid Nejati Manzari](https://omid-nejati.github.io/), 
+[Berardino Barile](https://scholar.google.com/citations?user=odmpMGcAAAAJ&hl=en), 
+[Yiming Xiao](https://yimingxiao.weebly.com/curriculum-vitae.html)<sup>†</sup>, 
+[Hassan Rivaz](https://users.encs.concordia.ca/~hrivaz/)<sup>†</sup>
+
+<a href="https://arxiv.org/abs/2602.20423" target="_blank"><img alt="arXiv" src="https://img.shields.io/badge/arXiv-2602.20423-B31B1B?logo=arxiv&logoColor=white" height="25"/></a>
+<a href="https://tahakoleilat.github.io/MedCLIPSeg" target="_blank"><img alt="Project Website" src="https://img.shields.io/badge/%F0%9F%94%97%20Project-Website-blue" height="25"/></a>
+<a href="https://huggingface.co/datasets/TahaKoleilat/MedCLIPSeg" target="_blank"><img alt="HuggingFace Dataset" src="https://img.shields.io/badge/%F0%9F%A4%97%20_Datasets-Access-ffc107?color=ffc107&logoColor=white" height="25"/></a>
+<a href="https://huggingface.co/TahaKoleilat/MedCLIPSeg" target="_blank"><img alt="HuggingFace Models" src="https://img.shields.io/badge/Models-Reproduce-2ea44f?logo=huggingface&logoColor=white" height="25"/></a>
+<a href="#citation"><img alt="Citation" src="https://img.shields.io/badge/Citation-BibTeX-6C63FF?logo=bookstack&logoColor=white" height="25"/></a>
+
+† *Co-senior authors*
+</div>
+
+## Overview
+
+![main figure](assets/overview.png)
+> **<p align="justify"> Abstract:** *Medical image segmentation remains challenging due to limited annotations for training, ambiguous anatomical features, and domain shifts. While vision–language models such as **CLIP** offer strong cross-modal representations, their potential for dense, text-guided medical image segmentation remains underexplored. We present **MedCLIPSeg**, a novel framework that adapts CLIP for **robust, data-efficient, and uncertainty-aware** medical image segmentation. Our approach leverages **patch-level CLIP embeddings** through **probabilistic cross-modal attention**, enabling bidirectional interaction between image and text tokens and explicit modeling of predictive uncertainty. Together with a **soft patch-level contrastive loss** that encourages nuanced semantic learning across diverse textual prompts, **MedCLIPSeg** improves data efficiency and domain generalizability. Extensive experiments across **16 datasets**, spanning **five imaging modalities** and **six organs**, demonstrate that **MedCLIPSeg** outperforms prior methods in **accuracy, efficiency, and robustness**, while providing **interpretable uncertainty maps** that highlight the local reliability of segmentation results. This work demonstrates the potential of **probabilistic vision–language modeling** for text-driven medical image segmentation.* </p>
+
+## Method
+
+<p align="center">
+  <img src="assets/MedCLIPSeg.png" width="100%" />
+  <br>
+  <em>Overall architecture of MedCLIPSeg. The framework integrates probabilistic vision–language fusion into a CLIP-based segmentation pipeline.</em>
+</p>
+
+<br>
+
+<p align="center">
+  <img src="assets/MedCLIPSeg_PVL.png" width="85%" />
+  <br>
+  <em>Schematic illustration of the proposed Probabilistic Vision–Language (PVL) adapters used for bidirectional cross-modal interaction.</em>
+</p>
+
+
+1) **Bidirectional Vision–Language Fusion**: Introduce representation-level fusion modules that enable efficient bidirectional interaction between image and text features while keeping CLIP encoders frozen, improving data efficiency and robustness.
+
+2) **Probabilistic Cross-Modal Attention**: Model vision–language attention using variational Key–Value formulations to capture uncertainty, leading to improved segmentation accuracy and cross-domain generalization.
+
+3) **Pixel-Level Uncertainty Estimation**: Generate dense uncertainty maps by sampling attention Values from learned probability distributions, providing intuitive reliability estimates for clinical interpretation.
+
+4) **Extensive Multi-Modal Segmentation Evaluation**: Conduct comprehensive evaluation against state-of-the-art methods across 5 imaging modalities and 6 organs and 16 datasets, assessing data efficiency, domain generalization, and the contribution of individual model components.
+
+## Results
+Results reported below show DSC scores (%) for data efficiency and domain generalization evaluation benchmarks across 16 biomedical image segmentation datasets averaged.
+
+### Data-Efficiency Evaluation
+
+| **Method** | **10% Data** | **25% Data** | **50% Data** | **100% Data** |
+|-----------|:------------:|:------------:|:------------:|:-------------:|
+| [UNet](https://arxiv.org/abs/1505.04597) | 60.95 | 62.74 | 71.61 | 78.49 |
+| [UNet++](https://arxiv.org/abs/1807.10165) | 63.72 | 65.86 | 73.15 | 78.44 |
+| [DeepLabv3](https://arxiv.org/abs/1706.05587) | 61.32 | 65.39 | 68.58 | 73.28 |
+| [Attention U-Net](https://arxiv.org/abs/1804.03999) | 62.78 | 64.97 | 71.34 | 76.30 |
+| [nnU-Net](https://arxiv.org/abs/1809.10486) | 73.45 | 76.73 | 78.86 | 81.40 |
+| [Swin-UNet](https://arxiv.org/abs/2105.05537) | 53.04 | 54.69 | 55.89 | 65.03 |
+| [TransUNet](https://arxiv.org/abs/2102.04306) | 52.69 | 55.25 | 55.22 | 67.22 |
+| [LViT](https://arxiv.org/abs/2206.14718) | 66.51 | 75.66 | 78.88 | 83.35 |
+| [Ariadne’s Thread](https://arxiv.org/abs/2307.03942) | 61.34 | 63.09 | 65.65 | 70.07 |
+| [EoMT-CLIP](https://arxiv.org/abs/2503.19108) | 74.07 | 76.29 | 79.19 | 82.93 |
+| [CLIPSeg](https://arxiv.org/abs/2112.10003) | 74.66 | 78.31 | 79.63 | 84.87 |
+| [DenseCLIP](https://arxiv.org/abs/2112.01518) | 67.84 | 70.23 | 72.09 | 74.19 |
+| [ZegCLIP](https://arxiv.org/abs/2212.03588) | 61.25 | 72.46 | 76.21 | 78.98 |
+| [SAN](https://arxiv.org/abs/2302.12242) | 74.13 | 76.13 | 78.80 | 81.62 |
+| [MaPLe](https://arxiv.org/abs/2210.03117) | 66.27 | 71.53 | 74.60 | 74.60 |
+| [MaPLe + Decoder](https://arxiv.org/abs/2210.03117) | 74.81 | 79.64 | 82.81 | 84.94 |
+| [VLSM-Adapter](https://arxiv.org/abs/2405.06196) | 74.47 | 77.63 | 80.83 | 83.85 |
+| [CausalCLIPSeg](https://arxiv.org/abs/2503.15949) | 71.19 | 75.42 | 78.60 | 81.34 |
+| [CAT-Seg](https://arxiv.org/abs/2303.11797) | *78.76* | *81.12* | *83.32* | *85.90* |
+| **[MedCLIPSeg (Ours)](https://arxiv.org/abs/XXXX.XXXXX)** | **81.10** | **85.08** | **87.18** | **88.66** |
+
+### Domain Generalization
+
+| **Method** | **ID** | **OOD** | **HM** |
+|-----------|:--------------------:|:---------------------:|:-----------------:|
+| [LViT](https://arxiv.org/abs/2206.14718) | 83.31 | 64.99 | 73.02 |
+| [Ariadne’s Thread](https://arxiv.org/abs/2307.03942) | 68.25 | 27.23 | 38.93 |
+| [CLIPSeg](https://arxiv.org/abs/2112.10003) | 84.95 | 69.22 | 76.28 |
+| [DenseCLIP](https://arxiv.org/abs/2112.01518) | 77.69 | 58.11 | 66.49 |
+| [ZegCLIP](https://arxiv.org/abs/2212.03588) | 77.16 | 61.33 | 68.34 |
+| [SAN](https://arxiv.org/abs/2302.12242) | 84.45 | 69.87 | 76.47 |
+| [MaPLe](https://arxiv.org/abs/2210.03117) | 76.55 | 59.30 | 66.83 |
+| [MaPLe + Decoder](https://arxiv.org/abs/2210.03117) | 84.78 | 66.85 | 74.76 |
+| [VLSM-Adapter](https://arxiv.org/abs/2405.06196) | 85.78 | 73.28 | 79.04 |
+| [CausalCLIPSeg](https://arxiv.org/abs/2503.15949) | 81.52 | 53.86 | 64.86 |
+| [CAT-Seg](https://arxiv.org/abs/2303.11797) | 86.10 | 74.57 | 79.92 |
+| **[MedCLIPSeg (Ours)](https://arxiv.org/abs/XXXX.XXXXX)** | **89.11** | **79.02** | **83.76** |
+
+### Segmentation and Uncertainty Visualization
+
+<p align="center">
+  <img src="assets/MedCLIPSeg_Seg.png" width="100%" />
+</p>
+
+<p align="center">
+  <em>
+    Uncertainty peaks along lesion boundaries and remains consistent across diverse datasets, indicating reliable calibration and generalization.
+    In-distribution (ID) data are shown in <b style="color:#0000FF;">blue</b>, while out-of-distribution (OOD) data are shown in <b style="color:#FF0000;">red</b>.
+  </em>
+</p>
+
+## Model Checkpoints
+All the checkpoints can be found on the official [Hugging Face repo](https://huggingface.co/TahaKoleilat/MedCLIPSeg) for the Data Efficiency and Domain Generalization evaluation benchmarks. Take a look [here](https://github.com/HealthX-Lab/MedCLIPSeg/blob/main/assets/RUN.md#2-running-evaluation-from-given-checkpoints) to see how to run and reproduce all the results.
+
+## Installation 
+For installation and other package requirements, please follow the instructions detailed in [INSTALL.md](assets/INSTALL.md). 
+
+## Data preparation
+Please follow the instructions at [DATASETS.md](assets/DATASETS.md) to prepare all datasets.
+
+## Training and Evaluation
+Please refer to the [RUN.md](assets/RUN.md) for detailed instructions on training, evaluating and reproducing the results using our pre-trained models.
+
+<hr />
+
+## Citation
+If you use our work, please consider citing:
+```bibtex
+@inproceedings{koleilat2026medclipseg,
+    author    = {Koleilat, Taha and Asgariandehkordi, Hojat and Nejatimanzari, Omid and Barile, Berardino and Xiao, Yiming and Rivaz, Hassan},
+    title     = {MedCLIPSeg: Probabilistic Vision-Language Adaptation for Data-Efficient and Generalizable Medical Image Segmentation},
+    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+    month     = {June},
+    year      = {2026},
+    pages     = {1406-1417}
+}
+```
+
+## Acknowledgements
+We are grateful to the authors of [CLIP](https://github.com/openai/CLIP), [MaPLe](https://github.com/muzairkhattak/multimodal-prompt-learning), and [LViT](https://github.com/HUANGLIZI/LViT) for making their code publicly available. If you use our model or code, we kindly request that you also consider citing these foundational works.

reference_medclipseg/biomedclip/biomedclip.py

@@ -0,0 +1,80 @@
+import torch
+import torch.nn as nn
+from typing import Optional
+from open_clip.hf_model import ClsPooler
+
+class BiomedCLIP(nn.Module):
+    def __init__(self, model, prompt_depth = None, prompt_length = None, output_hidden_states=False):
+        super(BiomedCLIP, self).__init__()
+        self.vision_model = model.visual
+        self.text_model = model.text
+        self.logit_scale = model.logit_scale
+        self.prompt_depth = prompt_depth
+        self.output_hidden_states = output_hidden_states
+        self.prompt_tokens = nn.Parameter(torch.empty(prompt_depth, prompt_length, 768))
+        self.text_dtype = self.text_model.transformer.dtype
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        nn.init.normal_(self.prompt_tokens, mean=0, std=0.02)
+    def encode_image(self, x):
+        trunk = self.vision_model.trunk
+        x = trunk.patch_embed(x)
+        x = trunk._pos_embed(x)
+        x = trunk.norm_pre(x)
+
+        hidden_states = []
+
+        for i, block in enumerate(trunk.blocks):
+            if(i < self.prompt_depth - 1):
+                x = block(x, self.prompt_tokens[i])
+            else:
+                x = block(x)
+
+            hidden_states.append(x)
+
+        x = trunk.norm(x)
+
+        # if trunk.global_pool:
+        #     x = (
+        #         x[:, trunk.num_prefix_tokens :].mean(dim=1)
+        #         if trunk.global_pool == "avg"
+        #         else x[:, 0]
+        #     )
+
+        x = trunk.fc_norm(x)
+        x = trunk.head(x)
+
+        # Linear Projection: 768 -> 512
+        x = self.vision_model.head(x)
+
+        if self.output_hidden_states:
+            return x, hidden_states
+        else:
+            return x
+    def encode_text(self, x, attention_mask: Optional[torch.LongTensor] = None, output_hidden_states: bool = False,):
+
+        if attention_mask is None:
+            attention_mask = (x != self.text_model.config.pad_token_id).long()
+
+        inputs_embeds = self.text_model.transformer.embeddings.word_embeddings(x).type(self.text_dtype).to(self.device)
+
+        x = self.text_model.transformer.embeddings(
+            inputs_embeds=inputs_embeds
+        )
+
+        extended_attention_mask = attention_mask[:, None, None, :]
+        extended_attention_mask = extended_attention_mask.to(dtype=self.text_dtype)  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(self.text_dtype).min
+
+        for i, layer in enumerate(self.text_model.transformer.encoder.layer):
+            if(i < self.prompt_depth - 1):
+                x = layer(x, attention_mask=extended_attention_mask, prompt_tokens=self.prompt_tokens[i])
+            else:
+                x = layer(x, attention_mask=extended_attention_mask)
+            x = x[0]
+
+        pooled_out = x[:, 0, :]
+        projected = self.text_model.proj(pooled_out)
+
+
+        return projected
+    

reference_medclipseg/biomedclip/layers/__init__.py

@@ -0,0 +1,62 @@
+from .activations import *
+from .adaptive_avgmax_pool import \
+    adaptive_avgmax_pool2d, select_adaptive_pool2d, AdaptiveAvgMaxPool2d, SelectAdaptivePool2d
+from .attention2d import MultiQueryAttention2d, Attention2d, MultiQueryAttentionV2
+from .attention_pool import AttentionPoolLatent
+from .attention_pool2d import AttentionPool2d, RotAttentionPool2d, RotaryEmbedding
+from .blur_pool import BlurPool2d, create_aa
+from .classifier import create_classifier, ClassifierHead, NormMlpClassifierHead, ClNormMlpClassifierHead
+from .cond_conv2d import CondConv2d, get_condconv_initializer
+from .config import is_exportable, is_scriptable, is_no_jit, use_fused_attn, \
+    set_exportable, set_scriptable, set_no_jit, set_layer_config, set_fused_attn, \
+    set_reentrant_ckpt, use_reentrant_ckpt
+from .conv2d_same import Conv2dSame, conv2d_same
+from .conv_bn_act import ConvNormAct, ConvNormActAa, ConvBnAct
+from .create_act import create_act_layer, get_act_layer, get_act_fn
+from .create_attn import get_attn, create_attn
+from .create_conv2d import create_conv2d
+from .create_norm import get_norm_layer, create_norm_layer
+from .create_norm_act import get_norm_act_layer, create_norm_act_layer, get_norm_act_layer
+from .drop import DropBlock2d, DropPath, drop_block_2d, drop_path
+from .eca import EcaModule, CecaModule, EfficientChannelAttn, CircularEfficientChannelAttn
+from .evo_norm import EvoNorm2dB0, EvoNorm2dB1, EvoNorm2dB2,\
+    EvoNorm2dS0, EvoNorm2dS0a, EvoNorm2dS1, EvoNorm2dS1a, EvoNorm2dS2, EvoNorm2dS2a
+from .fast_norm import is_fast_norm, set_fast_norm, fast_group_norm, fast_layer_norm
+from .filter_response_norm import FilterResponseNormTlu2d, FilterResponseNormAct2d
+from .format import Format, get_channel_dim, get_spatial_dim, nchw_to, nhwc_to
+from .gather_excite import GatherExcite
+from .global_context import GlobalContext
+from .grid import ndgrid, meshgrid
+from .helpers import to_ntuple, to_2tuple, to_3tuple, to_4tuple, make_divisible, extend_tuple
+from .hybrid_embed import HybridEmbed, HybridEmbedWithSize
+from .inplace_abn import InplaceAbn
+from .layer_scale import LayerScale, LayerScale2d
+from .linear import Linear
+from .mixed_conv2d import MixedConv2d
+from .mlp import Mlp, GluMlp, GatedMlp, SwiGLU, SwiGLUPacked, ConvMlp, GlobalResponseNormMlp
+from .non_local_attn import NonLocalAttn, BatNonLocalAttn
+from .norm import GroupNorm, GroupNorm1, LayerNorm, LayerNorm2d, RmsNorm, RmsNorm2d, SimpleNorm, SimpleNorm2d
+from .norm_act import BatchNormAct2d, GroupNormAct, GroupNorm1Act, LayerNormAct, LayerNormAct2d,\
+    SyncBatchNormAct, convert_sync_batchnorm, FrozenBatchNormAct2d, freeze_batch_norm_2d, unfreeze_batch_norm_2d
+from .padding import get_padding, get_same_padding, pad_same
+from .patch_dropout import PatchDropout
+from .patch_embed import PatchEmbed, PatchEmbedWithSize, resample_patch_embed
+from .pool2d_same import AvgPool2dSame, create_pool2d
+from .pos_embed import resample_abs_pos_embed, resample_abs_pos_embed_nhwc
+from .pos_embed_rel import RelPosMlp, RelPosBias, RelPosBiasTf, gen_relative_position_index, gen_relative_log_coords, \
+    resize_rel_pos_bias_table, resize_rel_pos_bias_table_simple, resize_rel_pos_bias_table_levit
+from .pos_embed_sincos import pixel_freq_bands, freq_bands, build_sincos2d_pos_embed, build_fourier_pos_embed, \
+    build_rotary_pos_embed, apply_rot_embed, apply_rot_embed_cat, apply_rot_embed_list, apply_keep_indices_nlc, \
+    FourierEmbed, RotaryEmbedding, RotaryEmbeddingCat
+from .squeeze_excite import SEModule, SqueezeExcite, EffectiveSEModule, EffectiveSqueezeExcite
+from .selective_kernel import SelectiveKernel
+from .separable_conv import SeparableConv2d, SeparableConvNormAct
+from .space_to_depth import SpaceToDepth, DepthToSpace
+from .split_attn import SplitAttn
+from .split_batchnorm import SplitBatchNorm2d, convert_splitbn_model
+from .std_conv import StdConv2d, StdConv2dSame, ScaledStdConv2d, ScaledStdConv2dSame
+from .test_time_pool import TestTimePoolHead, apply_test_time_pool
+from .trace_utils import _assert, _float_to_int
+from .typing import LayerType, PadType
+from .weight_init import trunc_normal_, trunc_normal_tf_, variance_scaling_, lecun_normal_, \
+    init_weight_jax, init_weight_vit

reference_medclipseg/biomedclip/layers/activations.py

@@ -0,0 +1,173 @@
+""" Activations
+
+A collection of activations fn and modules with a common interface so that they can
+easily be swapped. All have an `inplace` arg even if not used.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+
+def swish(x, inplace: bool = False):
+    """Swish - Described in: https://arxiv.org/abs/1710.05941
+    """
+    return x.mul_(x.sigmoid()) if inplace else x.mul(x.sigmoid())
+
+
+class Swish(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(Swish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return swish(x, self.inplace)
+
+
+def mish(x, inplace: bool = False):
+    """Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    NOTE: I don't have a working inplace variant
+    """
+    return x.mul(F.softplus(x).tanh())
+
+
+class Mish(nn.Module):
+    """Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    """
+    def __init__(self, inplace: bool = False):
+        super(Mish, self).__init__()
+
+    def forward(self, x):
+        return mish(x)
+
+
+def sigmoid(x, inplace: bool = False):
+    return x.sigmoid_() if inplace else x.sigmoid()
+
+
+# PyTorch has this, but not with a consistent inplace argument interface
+class Sigmoid(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(Sigmoid, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return x.sigmoid_() if self.inplace else x.sigmoid()
+
+
+def tanh(x, inplace: bool = False):
+    return x.tanh_() if inplace else x.tanh()
+
+
+# PyTorch has this, but not with a consistent inplace argument interface
+class Tanh(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(Tanh, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return x.tanh_() if self.inplace else x.tanh()
+
+
+def hard_swish(x, inplace: bool = False):
+    inner = F.relu6(x + 3.).div_(6.)
+    return x.mul_(inner) if inplace else x.mul(inner)
+
+
+class HardSwish(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSwish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return hard_swish(x, self.inplace)
+
+
+def hard_sigmoid(x, inplace: bool = False):
+    if inplace:
+        return x.add_(3.).clamp_(0., 6.).div_(6.)
+    else:
+        return F.relu6(x + 3.) / 6.
+
+
+class HardSigmoid(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSigmoid, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return hard_sigmoid(x, self.inplace)
+
+
+def hard_mish(x, inplace: bool = False):
+    """ Hard Mish
+    Experimental, based on notes by Mish author Diganta Misra at
+      https://github.com/digantamisra98/H-Mish/blob/0da20d4bc58e696b6803f2523c58d3c8a82782d0/README.md
+    """
+    if inplace:
+        return x.mul_(0.5 * (x + 2).clamp(min=0, max=2))
+    else:
+        return 0.5 * x * (x + 2).clamp(min=0, max=2)
+
+
+class HardMish(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardMish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return hard_mish(x, self.inplace)
+
+
+class PReLU(nn.PReLU):
+    """Applies PReLU (w/ dummy inplace arg)
+    """
+    def __init__(self, num_parameters: int = 1, init: float = 0.25, inplace: bool = False) -> None:
+        super(PReLU, self).__init__(num_parameters=num_parameters, init=init)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.prelu(input, self.weight)
+
+
+def gelu(x: torch.Tensor, inplace: bool = False) -> torch.Tensor:
+    return F.gelu(x)
+
+
+class GELU(nn.Module):
+    """Applies the Gaussian Error Linear Units function (w/ dummy inplace arg)
+    """
+    def __init__(self, inplace: bool = False):
+        super(GELU, self).__init__()
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.gelu(input)
+
+
+def gelu_tanh(x: torch.Tensor, inplace: bool = False) -> torch.Tensor:
+    return F.gelu(x, approximate='tanh')
+
+
+class GELUTanh(nn.Module):
+    """Applies the Gaussian Error Linear Units function (w/ dummy inplace arg)
+    """
+    def __init__(self, inplace: bool = False):
+        super(GELUTanh, self).__init__()
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.gelu(input, approximate='tanh')
+
+
+def quick_gelu(x: torch.Tensor, inplace: bool = False) -> torch.Tensor:
+    return x * torch.sigmoid(1.702 * x)
+
+
+class QuickGELU(nn.Module):
+    """Applies the Gaussian Error Linear Units function (w/ dummy inplace arg)
+    """
+    def __init__(self, inplace: bool = False):
+        super(QuickGELU, self).__init__()
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return quick_gelu(input)

reference_medclipseg/biomedclip/layers/activations_me.py

@@ -0,0 +1,208 @@
+""" Activations (memory-efficient w/ custom autograd)
+
+A collection of activations fn and modules with a common interface so that they can
+easily be swapped. All have an `inplace` arg even if not used.
+
+These activations are not compatible with jit scripting or ONNX export of the model, please use
+basic versions of the activations.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+
+def swish_fwd(x):
+    return x.mul(torch.sigmoid(x))
+
+
+def swish_bwd(x, grad_output):
+    x_sigmoid = torch.sigmoid(x)
+    return grad_output * (x_sigmoid * (1 + x * (1 - x_sigmoid)))
+
+
+class SwishAutoFn(torch.autograd.Function):
+    """ optimised Swish w/ memory-efficient checkpoint
+    Inspired by conversation btw Jeremy Howard & Adam Pazske
+    https://twitter.com/jeremyphoward/status/1188251041835315200
+    """
+    @staticmethod
+    def symbolic(g, x):
+        return g.op("Mul", x, g.op("Sigmoid", x))
+
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return swish_fwd(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x = ctx.saved_tensors[0]
+        return swish_bwd(x, grad_output)
+
+
+def swish_me(x, inplace=False):
+    return SwishAutoFn.apply(x)
+
+
+class SwishMe(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(SwishMe, self).__init__()
+
+    def forward(self, x):
+        return SwishAutoFn.apply(x)
+
+
+def mish_fwd(x):
+    return x.mul(torch.tanh(F.softplus(x)))
+
+
+def mish_bwd(x, grad_output):
+    x_sigmoid = torch.sigmoid(x)
+    x_tanh_sp = F.softplus(x).tanh()
+    return grad_output.mul(x_tanh_sp + x * x_sigmoid * (1 - x_tanh_sp * x_tanh_sp))
+
+
+class MishAutoFn(torch.autograd.Function):
+    """ Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    A memory efficient variant of Mish
+    """
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return mish_fwd(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x = ctx.saved_tensors[0]
+        return mish_bwd(x, grad_output)
+
+
+def mish_me(x, inplace=False):
+    return MishAutoFn.apply(x)
+
+
+class MishMe(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(MishMe, self).__init__()
+
+    def forward(self, x):
+        return MishAutoFn.apply(x)
+
+
+def hard_sigmoid_fwd(x, inplace: bool = False):
+    return (x + 3).clamp(min=0, max=6).div(6.)
+
+
+def hard_sigmoid_bwd(x, grad_output):
+    m = torch.ones_like(x) * ((x >= -3.) & (x <= 3.)) / 6.
+    return grad_output * m
+
+
+class HardSigmoidAutoFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return hard_sigmoid_fwd(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x = ctx.saved_tensors[0]
+        return hard_sigmoid_bwd(x, grad_output)
+
+
+def hard_sigmoid_me(x, inplace: bool = False):
+    return HardSigmoidAutoFn.apply(x)
+
+
+class HardSigmoidMe(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSigmoidMe, self).__init__()
+
+    def forward(self, x):
+        return HardSigmoidAutoFn.apply(x)
+
+
+def hard_swish_fwd(x):
+    return x * (x + 3).clamp(min=0, max=6).div(6.)
+
+
+def hard_swish_bwd(x, grad_output):
+    m = torch.ones_like(x) * (x >= 3.)
+    m = torch.where((x >= -3.) & (x <= 3.),  x / 3. + .5, m)
+    return grad_output * m
+
+
+class HardSwishAutoFn(torch.autograd.Function):
+    """A memory efficient HardSwish activation"""
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return hard_swish_fwd(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x = ctx.saved_tensors[0]
+        return hard_swish_bwd(x, grad_output)
+
+    @staticmethod
+    def symbolic(g, self):
+        input = g.op("Add", self, g.op('Constant', value_t=torch.tensor(3, dtype=torch.float)))
+        hardtanh_ = g.op("Clip", input, g.op('Constant', value_t=torch.tensor(0, dtype=torch.float)), g.op('Constant', value_t=torch.tensor(6, dtype=torch.float)))
+        hardtanh_ = g.op("Div", hardtanh_, g.op('Constant', value_t=torch.tensor(6, dtype=torch.float)))
+        return g.op("Mul", self, hardtanh_)
+
+
+def hard_swish_me(x, inplace=False):
+    return HardSwishAutoFn.apply(x)
+
+
+class HardSwishMe(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardSwishMe, self).__init__()
+
+    def forward(self, x):
+        return HardSwishAutoFn.apply(x)
+
+
+def hard_mish_fwd(x):
+    return 0.5 * x * (x + 2).clamp(min=0, max=2)
+
+
+def hard_mish_bwd(x, grad_output):
+    m = torch.ones_like(x) * (x >= -2.)
+    m = torch.where((x >= -2.) & (x <= 0.), x + 1., m)
+    return grad_output * m
+
+
+class HardMishAutoFn(torch.autograd.Function):
+    """ A memory efficient variant of Hard Mish
+    Experimental, based on notes by Mish author Diganta Misra at
+      https://github.com/digantamisra98/H-Mish/blob/0da20d4bc58e696b6803f2523c58d3c8a82782d0/README.md
+    """
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return hard_mish_fwd(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x = ctx.saved_tensors[0]
+        return hard_mish_bwd(x, grad_output)
+
+
+def hard_mish_me(x, inplace: bool = False):
+    return HardMishAutoFn.apply(x)
+
+
+class HardMishMe(nn.Module):
+    def __init__(self, inplace: bool = False):
+        super(HardMishMe, self).__init__()
+
+    def forward(self, x):
+        return HardMishAutoFn.apply(x)
+
+
+

reference_medclipseg/biomedclip/layers/adaptive_avgmax_pool.py

@@ -0,0 +1,183 @@
+""" PyTorch selectable adaptive pooling
+Adaptive pooling with the ability to select the type of pooling from:
+    * 'avg' - Average pooling
+    * 'max' - Max pooling
+    * 'avgmax' - Sum of average and max pooling re-scaled by 0.5
+    * 'avgmaxc' - Concatenation of average and max pooling along feature dim, doubles feature dim
+
+Both a functional and a nn.Module version of the pooling is provided.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .format import get_spatial_dim, get_channel_dim
+
+_int_tuple_2_t = Union[int, Tuple[int, int]]
+
+
+def adaptive_pool_feat_mult(pool_type='avg'):
+    if pool_type.endswith('catavgmax'):
+        return 2
+    else:
+        return 1
+
+
+def adaptive_avgmax_pool2d(x, output_size: _int_tuple_2_t = 1):
+    x_avg = F.adaptive_avg_pool2d(x, output_size)
+    x_max = F.adaptive_max_pool2d(x, output_size)
+    return 0.5 * (x_avg + x_max)
+
+
+def adaptive_catavgmax_pool2d(x, output_size: _int_tuple_2_t = 1):
+    x_avg = F.adaptive_avg_pool2d(x, output_size)
+    x_max = F.adaptive_max_pool2d(x, output_size)
+    return torch.cat((x_avg, x_max), 1)
+
+
+def select_adaptive_pool2d(x, pool_type='avg', output_size: _int_tuple_2_t = 1):
+    """Selectable global pooling function with dynamic input kernel size
+    """
+    if pool_type == 'avg':
+        x = F.adaptive_avg_pool2d(x, output_size)
+    elif pool_type == 'avgmax':
+        x = adaptive_avgmax_pool2d(x, output_size)
+    elif pool_type == 'catavgmax':
+        x = adaptive_catavgmax_pool2d(x, output_size)
+    elif pool_type == 'max':
+        x = F.adaptive_max_pool2d(x, output_size)
+    else:
+        assert False, 'Invalid pool type: %s' % pool_type
+    return x
+
+
+class FastAdaptiveAvgPool(nn.Module):
+    def __init__(self, flatten: bool = False, input_fmt: F = 'NCHW'):
+        super(FastAdaptiveAvgPool, self).__init__()
+        self.flatten = flatten
+        self.dim = get_spatial_dim(input_fmt)
+
+    def forward(self, x):
+        return x.mean(self.dim, keepdim=not self.flatten)
+
+
+class FastAdaptiveMaxPool(nn.Module):
+    def __init__(self, flatten: bool = False, input_fmt: str = 'NCHW'):
+        super(FastAdaptiveMaxPool, self).__init__()
+        self.flatten = flatten
+        self.dim = get_spatial_dim(input_fmt)
+
+    def forward(self, x):
+        return x.amax(self.dim, keepdim=not self.flatten)
+
+
+class FastAdaptiveAvgMaxPool(nn.Module):
+    def __init__(self, flatten: bool = False, input_fmt: str = 'NCHW'):
+        super(FastAdaptiveAvgMaxPool, self).__init__()
+        self.flatten = flatten
+        self.dim = get_spatial_dim(input_fmt)
+
+    def forward(self, x):
+        x_avg = x.mean(self.dim, keepdim=not self.flatten)
+        x_max = x.amax(self.dim, keepdim=not self.flatten)
+        return 0.5 * x_avg + 0.5 * x_max
+
+
+class FastAdaptiveCatAvgMaxPool(nn.Module):
+    def __init__(self, flatten: bool = False, input_fmt: str = 'NCHW'):
+        super(FastAdaptiveCatAvgMaxPool, self).__init__()
+        self.flatten = flatten
+        self.dim_reduce = get_spatial_dim(input_fmt)
+        if flatten:
+            self.dim_cat = 1
+        else:
+            self.dim_cat = get_channel_dim(input_fmt)
+
+    def forward(self, x):
+        x_avg = x.mean(self.dim_reduce, keepdim=not self.flatten)
+        x_max = x.amax(self.dim_reduce, keepdim=not self.flatten)
+        return torch.cat((x_avg, x_max), self.dim_cat)
+
+
+class AdaptiveAvgMaxPool2d(nn.Module):
+    def __init__(self, output_size: _int_tuple_2_t = 1):
+        super(AdaptiveAvgMaxPool2d, self).__init__()
+        self.output_size = output_size
+
+    def forward(self, x):
+        return adaptive_avgmax_pool2d(x, self.output_size)
+
+
+class AdaptiveCatAvgMaxPool2d(nn.Module):
+    def __init__(self, output_size: _int_tuple_2_t = 1):
+        super(AdaptiveCatAvgMaxPool2d, self).__init__()
+        self.output_size = output_size
+
+    def forward(self, x):
+        return adaptive_catavgmax_pool2d(x, self.output_size)
+
+
+class SelectAdaptivePool2d(nn.Module):
+    """Selectable global pooling layer with dynamic input kernel size
+    """
+    def __init__(
+            self,
+            output_size: _int_tuple_2_t = 1,
+            pool_type: str = 'fast',
+            flatten: bool = False,
+            input_fmt: str = 'NCHW',
+    ):
+        super(SelectAdaptivePool2d, self).__init__()
+        assert input_fmt in ('NCHW', 'NHWC')
+        self.pool_type = pool_type or ''  # convert other falsy values to empty string for consistent TS typing
+        pool_type = pool_type.lower()
+        if not pool_type:
+            self.pool = nn.Identity()  # pass through
+            self.flatten = nn.Flatten(1) if flatten else nn.Identity()
+        elif pool_type.startswith('fast') or input_fmt != 'NCHW':
+            assert output_size == 1, 'Fast pooling and non NCHW input formats require output_size == 1.'
+            if pool_type.endswith('catavgmax'):
+                self.pool = FastAdaptiveCatAvgMaxPool(flatten, input_fmt=input_fmt)
+            elif pool_type.endswith('avgmax'):
+                self.pool = FastAdaptiveAvgMaxPool(flatten, input_fmt=input_fmt)
+            elif pool_type.endswith('max'):
+                self.pool = FastAdaptiveMaxPool(flatten, input_fmt=input_fmt)
+            elif pool_type == 'fast' or pool_type.endswith('avg'):
+                self.pool = FastAdaptiveAvgPool(flatten, input_fmt=input_fmt)
+            else:
+                assert False, 'Invalid pool type: %s' % pool_type
+            self.flatten = nn.Identity()
+        else:
+            assert input_fmt == 'NCHW'
+            if pool_type == 'avgmax':
+                self.pool = AdaptiveAvgMaxPool2d(output_size)
+            elif pool_type == 'catavgmax':
+                self.pool = AdaptiveCatAvgMaxPool2d(output_size)
+            elif pool_type == 'max':
+                self.pool = nn.AdaptiveMaxPool2d(output_size)
+            elif pool_type == 'avg':
+                self.pool = nn.AdaptiveAvgPool2d(output_size)
+            else:
+                assert False, 'Invalid pool type: %s' % pool_type
+            self.flatten = nn.Flatten(1) if flatten else nn.Identity()
+
+    def is_identity(self):
+        return not self.pool_type
+
+    def forward(self, x):
+        x = self.pool(x)
+        x = self.flatten(x)
+        return x
+
+    def feat_mult(self):
+        return adaptive_pool_feat_mult(self.pool_type)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(' \
+               + 'pool_type=' + self.pool_type \
+               + ', flatten=' + str(self.flatten) + ')'
+

reference_medclipseg/biomedclip/layers/attention2d.py

@@ -0,0 +1,351 @@
+from typing import List, Optional, Type, Union
+
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+from .config import use_fused_attn
+from .create_conv2d import create_conv2d
+from .helpers import to_2tuple
+from .pool2d_same import create_pool2d
+
+
+class MultiQueryAttentionV2(nn.Module):
+    """Multi Query Attention.
+
+    Fast Transformer Decoding: One Write-Head is All You Need
+    https://arxiv.org/pdf/1911.02150.pdf
+
+    This is an acceletor optimized version - removing multiple unnecessary
+    tensor transpose by re-arranging indices according to the following rules: 1)
+    contracted indices are at the end, 2) other indices have the same order in the
+    input and output tensores.
+
+    Compared to V1, this gives 3x speed up.
+    """
+
+    def __init__(
+            self,
+            dim: int,
+            dim_out: Optional[int] = None,
+            num_heads: int = 8,
+            key_dim: int = 64,
+            value_dim: int = 64,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+    ):
+        """Initializer."""
+        super().__init__()
+        dim_out = dim_out or dim
+        self.num_heads = num_heads
+        self.key_dim = key_dim
+        self.value_dim = value_dim
+        self.scale = key_dim ** -0.5
+
+        self.query_proj = nn.Parameter(torch.randn([self.num_heads, self.key_dim, dim]))
+        self.key_proj = nn.Parameter(torch.randn([dim, self.key_dim]))
+        self.value_proj = nn.Parameter(torch.randn([dim, self.value_dim]))
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.out_proj = nn.Parameter(torch.randn([dim_out, self.num_heads, self.value_dim]))
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def _reshape_input(self, t):
+        """Reshapes a tensor to three dimensions, keeping the first and last."""
+        s = t.shape
+        # Propagate the shape statically where possible.
+        #num = t.shape[1:-1].numel()
+        #return t.reshape(s[0], num, s[-1])
+        return t.reshape(s[0], s[1], -1).transpose(1, 2)
+
+    def forward(self, x, m: Optional[torch.Tensor] = None):
+        """Run layer computation."""
+        b, _, h, w = x.shape
+        m = m if m is not None else x
+
+        reshaped_x = self._reshape_input(x)
+        reshaped_m = self._reshape_input(m)
+
+        q = torch.einsum('bnd,hkd->bnhk', reshaped_x, self.query_proj)
+        k = torch.einsum('bmd,dk->bmk', reshaped_m, self.key_proj)
+
+        attn = torch.einsum('bnhk,bmk->bnhm', q, k) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        v = torch.einsum('bmd,dv->bmv', reshaped_m, self.value_proj)
+        o = torch.einsum('bnhm,bmv->bnhv', attn, v)
+        result = torch.einsum('bnhv,dhv->bdn', o, self.out_proj)
+        result = self.proj_drop(result)
+        return result.reshape(b, -1, h, w)
+
+
+class MultiQueryAttention2d(nn.Module):
+    """Multi Query Attention with spatial downsampling.
+
+     3 parameters are introduced for the spatial downsampling:
+     1. kv_stride: downsampling factor on Key and Values only.
+     2. query_strides: horizontal & vertical strides on Query only.
+
+    This is an optimized version.
+    1. Projections in Attention is explicit written out as 1x1 Conv2D.
+    2. Additional reshapes are introduced to bring a up to 3x speed up.
+    """
+    fused_attn: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            dim: int,
+            dim_out: Optional[int] = None,
+            num_heads: int = 8,
+            key_dim: Optional[int] = None,
+            value_dim: Optional[int] = None,
+            query_strides: int = 1,
+            kv_stride: int = 1,
+            dw_kernel_size: int = 3,
+            dilation: int = 1,
+            padding: Union[str, int, List[int]] = '',
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: Type[nn.Module] = nn.BatchNorm2d,
+            use_bias: bool = False,
+    ):
+        """Initializer.
+
+        Args:
+          num_heads: Number of attention heads.
+          key_dim: Size of the attention key dimension.
+          value_dim: Size of the attention value dimension.
+          query_strides: Vertical stride size for query only.
+          kv_stride: Key and value stride size.
+          dw_kernel_size: Spatial dimension of the depthwise kernel.
+        """
+        super().__init__()
+        dim_out = dim_out or dim
+        self.num_heads = num_heads
+        self.key_dim = key_dim or dim // num_heads
+        self.value_dim = value_dim or dim // num_heads
+        self.query_strides = to_2tuple(query_strides)
+        self.kv_stride = kv_stride
+        self.has_query_strides = any([s > 1 for s in self.query_strides])
+        self.scale = self.key_dim ** -0.5
+        self.fused_attn = use_fused_attn()
+        self.drop = attn_drop
+
+        self.query = nn.Sequential()
+        if self.has_query_strides:
+            # FIXME dilation
+            if padding == 'same':
+                self.query.add_module('down_pool', create_pool2d(
+                        'avg',
+                        kernel_size=self.query_strides,
+                        padding='same',
+                ))
+            else:
+                # no pad if not 'same' as kern=stride=even
+                self.query.add_module('down_pool', nn.AvgPool2d(kernel_size=query_strides))
+            self.query.add_module('norm', norm_layer(dim))
+        self.query.add_module('proj', create_conv2d(
+            dim,
+            self.num_heads * self.key_dim,
+            kernel_size=1,
+            bias=use_bias,
+        ))
+
+        self.key = nn.Sequential()
+        if kv_stride > 1:
+            self.key.add_module('down_conv', create_conv2d(
+                dim,
+                dim,
+                kernel_size=dw_kernel_size,
+                stride=kv_stride,
+                dilation=dilation,
+                padding=padding,
+                depthwise=True,
+            ))
+            self.key.add_module('norm', norm_layer(dim))
+        self.key.add_module('proj', create_conv2d(
+            dim,
+            self.key_dim,
+            kernel_size=1,
+            padding=padding,
+            bias=use_bias,
+        ))
+
+        self.value = nn.Sequential()
+        if kv_stride > 1:
+            self.value.add_module('down_conv', create_conv2d(
+                dim,
+                dim,
+                kernel_size=dw_kernel_size,
+                stride=kv_stride,
+                dilation=dilation,
+                padding=padding,
+                depthwise=True,
+            ))
+            self.value.add_module('norm', norm_layer(dim))
+        self.value.add_module('proj', create_conv2d(
+            dim,
+            self.value_dim,
+            kernel_size=1,
+            bias=use_bias,
+        ))
+
+        self.attn_drop = nn.Dropout(attn_drop)
+
+        self.output = nn.Sequential()
+        if self.has_query_strides:
+            self.output.add_module('upsample', nn.Upsample(scale_factor=self.query_strides, mode='bilinear', align_corners=False))
+        self.output.add_module('proj', create_conv2d(
+            self.value_dim * self.num_heads,
+            dim_out,
+            kernel_size=1,
+            bias=use_bias,
+        ))
+        self.output.add_module('drop',  nn.Dropout(proj_drop))
+
+        self.einsum = False
+
+    def init_weights(self):
+        # using xavier appeared to improve stability for mobilenetv4 hybrid w/ this layer
+        nn.init.xavier_uniform_(self.query.proj.weight)
+        nn.init.xavier_uniform_(self.key.proj.weight)
+        nn.init.xavier_uniform_(self.value.proj.weight)
+        if self.kv_stride > 1:
+            nn.init.xavier_uniform_(self.key.down_conv.weight)
+            nn.init.xavier_uniform_(self.value.down_conv.weight)
+        nn.init.xavier_uniform_(self.output.proj.weight)
+
+    def _reshape_input(self, t: torch.Tensor):
+        """Reshapes a tensor to three dimensions, keeping the batch and channels."""
+        s = t.shape
+        t = t.reshape(s[0], s[1], -1).transpose(1, 2)
+        if self.einsum:
+            return t
+        else:
+            return t.unsqueeze(1).contiguous()
+
+    def _reshape_projected_query(self, t: torch.Tensor, num_heads: int, key_dim: int):
+        """Reshapes projected query: [b, n, n, h x k] -> [b, n x n, h, k]."""
+        s = t.shape
+        t = t.reshape(s[0], num_heads, key_dim, -1)
+        if self.einsum:
+            return t.permute(0, 3, 1, 2).contiguous()
+        else:
+            return t.transpose(-1, -2).contiguous()
+
+    def _reshape_output(self, t: torch.Tensor, num_heads: int, h_px: int, w_px: int):
+        """Reshape output:[b, n x n x h, k] -> [b, n, n, hk]."""
+        s = t.shape
+        feat_dim = s[-1] * num_heads
+        if not self.einsum:
+            t = t.transpose(1, 2)
+        return t.reshape(s[0], h_px, w_px, feat_dim).permute(0, 3, 1, 2).contiguous()
+
+    def forward(self, x, attn_mask: Optional[torch.Tensor] = None):
+        """Run layer computation."""
+        B, C, H, W = s = x.shape
+
+        q = self.query(x)
+        # desired q shape: [b, h, k, n x n] - [b, l, h, k]
+        q = self._reshape_projected_query(q, self.num_heads, self.key_dim)
+
+        k = self.key(x)
+        # output shape of k: [b, k, p], p = m x m
+        k = self._reshape_input(k)
+
+        v = self.value(x)
+        # output shape of v: [ b, p, k], p = m x m
+        v = self._reshape_input(v)
+
+        # desired q shape: [b, n x n, h, k]
+        # desired k shape: [b, m x m, k]
+        # desired logits shape: [b, n x n, h, m x m]
+        if self.einsum:
+            attn = torch.einsum('blhk,bpk->blhp', q, k) * self.scale
+            if attn_mask is not None:
+                # NOTE: assumes mask is float and in correct shape
+                attn = attn + attn_mask
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            o = torch.einsum('blhp,bpk->blhk', attn, v)
+        else:
+            if self.fused_attn:
+                o = F.scaled_dot_product_attention(
+                    q, k, v,
+                    attn_mask=attn_mask,
+                    dropout_p=self.attn_drop.p if self.training else 0.
+                )
+            else:
+                q = q * self.scale
+                attn = q @ k.transpose(-1, -2)
+                if attn_mask is not None:
+                    # NOTE: assumes mask is float and in correct shape
+                    attn = attn + attn_mask
+                attn = attn.softmax(dim=-1)
+                attn = self.attn_drop(attn)
+                o = attn @ v
+
+        # reshape o into [b, hk, n, n,]
+        o = self._reshape_output(o, self.num_heads, H // self.query_strides[0], W // self.query_strides[1])
+        x = self.output(o)
+        return x
+
+
+class Attention2d(nn.Module):
+    fused_attn: torch.jit.Final[bool]
+
+    """ multi-head attention for 2D NCHW tensors"""
+    def __init__(
+            self,
+            dim: int,
+            dim_out: Optional[int] = None,
+            num_heads: int = 32,
+            bias: bool = True,
+            expand_first: bool = False,
+            head_first: bool = False,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.
+    ):
+        super().__init__()
+        dim_out = dim_out or dim
+        dim_attn = dim_out if expand_first else dim
+        self.num_heads = num_heads
+        self.dim_head = dim_attn // num_heads
+        self.head_first = head_first
+        self.fused_attn = use_fused_attn()
+
+        self.qkv = nn.Conv2d(dim, dim_attn * 3, 1, bias=bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Conv2d(dim_attn, dim_out, 1, bias=bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, attn_mask: Optional[torch.Tensor] = None):
+        B, C, H, W = x.shape
+
+        if self.head_first:
+            q, k, v = self.qkv(x).view(B, self.num_heads, self.dim_head * 3, -1).chunk(3, dim=2)
+        else:
+            q, k, v = self.qkv(x).reshape(B, 3, self.num_heads, self.dim_head, -1).unbind(1)
+
+        if self.fused_attn:
+            x = torch.nn.functional.scaled_dot_product_attention(
+                q.transpose(-1, -2).contiguous(),
+                k.transpose(-1, -2).contiguous(),
+                v.transpose(-1, -2).contiguous(),
+                attn_mask=attn_mask,
+                dropout_p=self.attn_drop.p if self.training else 0.,
+            ).transpose(-1, -2).reshape(B, -1, H, W)
+        else:
+            q = q.transpose(-1, -2)
+            v = v.transpose(-1, -2)
+            attn = q @ k * q.size(-1) ** -0.5
+            if attn_mask is not None:
+                # NOTE: assumes mask is float and in correct shape
+                attn = attn + attn_mask
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = (attn @ v).transpose(-1, -2).reshape(B, -1, H, W)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

reference_medclipseg/biomedclip/layers/attention_pool.py

@@ -0,0 +1,105 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import use_fused_attn
+from .mlp import Mlp
+from .weight_init import trunc_normal_tf_
+
+
+class AttentionPoolLatent(nn.Module):
+    """ Attention pooling w/ latent query
+    """
+    fused_attn: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            in_features: int,
+            out_features: int = None,
+            embed_dim: int = None,
+            num_heads: int = 8,
+            feat_size: Optional[int] = None,
+            mlp_ratio: float = 4.0,
+            qkv_bias: bool = True,
+            qk_norm: bool = False,
+            latent_len: int = 1,
+            latent_dim: int = None,
+            pos_embed: str = '',
+            pool_type: str = 'token',
+            norm_layer: Optional[nn.Module] = None,
+            drop: float = 0.0,
+    ):
+        super().__init__()
+        embed_dim = embed_dim or in_features
+        out_features = out_features or in_features
+        assert embed_dim % num_heads == 0
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.feat_size = feat_size
+        self.scale = self.head_dim ** -0.5
+        self.pool = pool_type
+        self.fused_attn = use_fused_attn()
+
+        if pos_embed == 'abs':
+            assert feat_size is not None
+            self.pos_embed = nn.Parameter(torch.zeros(feat_size, in_features))
+        else:
+            self.pos_embed = None
+
+        self.latent_dim = latent_dim or embed_dim
+        self.latent_len = latent_len
+        self.latent = nn.Parameter(torch.zeros(1, self.latent_len, embed_dim))
+
+        self.q = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)
+        self.kv = nn.Linear(embed_dim, embed_dim * 2, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.proj = nn.Linear(embed_dim, embed_dim)
+        self.proj_drop = nn.Dropout(drop)
+
+        self.norm = norm_layer(out_features) if norm_layer is not None else nn.Identity()
+        self.mlp = Mlp(embed_dim, int(embed_dim * mlp_ratio))
+
+        self.init_weights()
+
+    def init_weights(self):
+        if self.pos_embed is not None:
+            trunc_normal_tf_(self.pos_embed, std=self.pos_embed.shape[1] ** -0.5)
+        trunc_normal_tf_(self.latent, std=self.latent_dim ** -0.5)
+
+    def forward(self, x):
+        B, N, C = x.shape
+
+        if self.pos_embed is not None:
+            # FIXME interpolate
+            x = x + self.pos_embed.unsqueeze(0).to(x.dtype)
+
+        q_latent = self.latent.expand(B, -1, -1)
+        q = self.q(q_latent).reshape(B, self.latent_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        kv = self.kv(x).reshape(B, N, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        k, v = kv.unbind(0)
+
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        if self.fused_attn:
+            x = F.scaled_dot_product_attention(q, k, v)
+        else:
+            q = q * self.scale
+            attn = q @ k.transpose(-2, -1)
+            attn = attn.softmax(dim=-1)
+            x = attn @ v
+        x = x.transpose(1, 2).reshape(B, self.latent_len, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        x = x + self.mlp(self.norm(x))
+
+        # optional pool if latent seq_len > 1 and pooled output is desired
+        if self.pool == 'token':
+            x = x[:, 0]
+        elif self.pool == 'avg':
+            x = x.mean(1)
+        return x

reference_medclipseg/biomedclip/layers/attention_pool2d.py

@@ -0,0 +1,278 @@
+""" Attention Pool 2D
+
+Implementations of 2D spatial feature pooling using multi-head attention instead of average pool.
+
+Based on idea in CLIP by OpenAI, licensed Apache 2.0
+https://github.com/openai/CLIP/blob/3b473b0e682c091a9e53623eebc1ca1657385717/clip/model.py
+
+Hacked together by / Copyright 2021 Ross Wightman
+"""
+from typing import Optional, Union, Tuple
+
+import torch
+import torch.nn as nn
+
+from. config import use_fused_attn
+from .helpers import to_2tuple
+from .pos_embed import resample_abs_pos_embed
+from .pos_embed_sincos import apply_rot_embed, RotaryEmbedding
+from .weight_init import trunc_normal_
+
+
+class RotAttentionPool2d(nn.Module):
+    """ Attention based 2D feature pooling w/ rotary (relative) pos embedding.
+    This is a multi-head attention based replacement for (spatial) average pooling in NN architectures.
+
+    Adapted from the AttentionPool2d in CLIP w/ rotary embedding instead of learned embed.
+    https://github.com/openai/CLIP/blob/3b473b0e682c091a9e53623eebc1ca1657385717/clip/model.py
+
+    NOTE: While this impl does not require a fixed feature size, performance at differeing resolutions from
+    train varies widely and falls off dramatically. I'm not sure if there is a way around this... -RW
+    """
+    fused_attn: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            in_features: int,
+            out_features: Optional[int] = None,
+            ref_feat_size: Union[int, Tuple[int, int]] = 7,
+            embed_dim: Optional[int] = None,
+            head_dim: Optional[int] = 64,
+            num_heads: Optional[int] = None,
+            qkv_bias: bool = True,
+            qkv_separate: bool = False,
+            pool_type: str = 'token',
+            class_token: bool = False,
+            drop_rate: float = 0.,
+    ):
+        super().__init__()
+        assert pool_type in ('', 'token')
+        self.embed_dim = embed_dim = embed_dim or in_features
+        self.in_features = in_features
+        self.out_features = out_features or in_features
+        ref_feat_size = to_2tuple(ref_feat_size)
+        if num_heads is not None:
+            assert embed_dim % num_heads == 0
+            head_dim = embed_dim // num_heads
+        else:
+            assert embed_dim % head_dim == 0
+            num_heads = embed_dim // head_dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.pool_type = pool_type.lower()
+        self.scale = self.head_dim ** -0.5
+        self.fused_attn = use_fused_attn()
+
+        if class_token:
+            self.cls_token = nn.Parameter(torch.zeros(1, embed_dim))
+        else:
+            self.cls_token = None
+
+        if qkv_separate:
+            self.q = nn.Linear(in_features, embed_dim, bias=qkv_bias)
+            self.k = nn.Linear(in_features, embed_dim, bias=qkv_bias)
+            self.v = nn.Linear(in_features, embed_dim, bias=qkv_bias)
+            self.qkv = None
+        else:
+            self.qkv = nn.Linear(in_features, embed_dim * 3, bias=qkv_bias)
+        self.drop = nn.Dropout(drop_rate)
+        self.proj = nn.Linear(embed_dim, self.out_features)
+        self.pos_embed = RotaryEmbedding(self.head_dim, in_pixels=False, ref_feat_shape=ref_feat_size)
+
+    def init_weights(self, zero_init_last: bool = False):
+        if self.qkv is None:
+            in_features = self.q.in_features
+            trunc_normal_(self.q.weight, std=in_features ** -0.5)
+            nn.init.zeros_(self.q.bias)
+            trunc_normal_(self.k.weight, std=in_features ** -0.5)
+            nn.init.zeros_(self.k.bias)
+            trunc_normal_(self.v.weight, std=in_features ** -0.5)
+            nn.init.zeros_(self.v.bias)
+        else:
+            in_features = self.qkv.in_features
+            trunc_normal_(self.qkv.weight, std=in_features ** -0.5)
+            nn.init.zeros_(self.qkv.bias)
+
+    def reset(self, num_classes: Optional[int] = None, pool_type: Optional[str] = None):
+        # NOTE: this module is being used as a head, so need compatible reset()
+        if pool_type is not None:
+            assert pool_type in ('', 'token')
+            self.pool_type = pool_type
+        if num_classes is not None:
+            self.proj = nn.Linear(self.in_features, num_classes) if num_classes > 0 else nn.Identity()
+            self.out_features = num_classes if num_classes > 0 else self.embed_dim
+
+    def _pool(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
+        if self.pool_type == 'token':
+            x = x[:, 0]
+        else:
+            # if not pooled, return spatial output without token
+            x = x[:, 1:].reshape(x.shape[0], H, W, -1).permute(0, 3, 1, 2)
+        return x
+
+    def forward(self, x, pre_logits: bool = False):
+        B, _, H, W = x.shape
+        N = H * W
+        x = x.flatten(2).transpose(1, 2)
+        if self.cls_token is None:
+            x = torch.cat([x.mean(1, keepdim=True), x], dim=1)
+        else:
+            x = torch.cat([self.cls_token.expand(x.shape[0], -1, -1), x], dim=1)
+        if self.qkv is None:
+            q = self.q(x).reshape(B, N + 1, self.num_heads, self.head_dim).transpose(1, 2)
+            k = self.k(x).reshape(B, N + 1, self.num_heads, self.head_dim).transpose(1, 2)
+            v = self.v(x).reshape(B, N + 1, self.num_heads, self.head_dim).transpose(1, 2)
+        else:
+            x = self.qkv(x).reshape(B, N + 1, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+            q, k, v = x.unbind(0)
+
+        rse, rce = self.pos_embed.get_embed((H, W))
+        q = torch.cat([q[:, :, :1, :], apply_rot_embed(q[:, :, 1:, :], rse, rce)], dim=2).type_as(v)
+        k = torch.cat([k[:, :, :1, :], apply_rot_embed(k[:, :, 1:, :], rse, rce)], dim=2).type_as(v)
+
+        if self.fused_attn:
+            x = nn.functional.scaled_dot_product_attention(q, k, v)
+        else:
+            q = q * self.scale
+            attn = q @ k.transpose(-2, -1)
+            attn = attn.softmax(dim=-1)
+            x = attn @ v
+        x = x.transpose(1, 2).reshape(B, N + 1, -1)
+        x = self.drop(x)
+        if pre_logits:
+            x = self._pool(x, H, W)
+            return x
+        x = self.proj(x)
+        x = self._pool(x, H, W)
+        return x
+
+
+class AttentionPool2d(nn.Module):
+    """ Attention based 2D feature pooling w/ learned (absolute) pos embedding.
+    This is a multi-head attention based replacement for (spatial) average pooling in NN architectures.
+
+    It was based on impl in CLIP by OpenAI
+    https://github.com/openai/CLIP/blob/3b473b0e682c091a9e53623eebc1ca1657385717/clip/model.py
+
+    NOTE: This requires feature size upon construction and well prevent adaptive sizing of the network.
+    """
+    fused_attn: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            in_features: int,
+            feat_size: Union[int, Tuple[int, int]] = 7,
+            out_features: Optional[int] = None,
+            embed_dim: Optional[int] = None,
+            head_dim: Optional[int] = 64,
+            num_heads: Optional[int] = None,
+            qkv_bias: bool = True,
+            qkv_separate: bool = False,
+            pool_type: str = 'token',
+            class_token: bool = False,
+            drop_rate: float = 0.,
+    ):
+        super().__init__()
+        assert pool_type in ('', 'token')
+        self.embed_dim = embed_dim = embed_dim or in_features
+        self.in_features = in_features
+        self.out_features = out_features or in_features
+        if num_heads is not None:
+            assert embed_dim % num_heads == 0
+            head_dim = embed_dim // num_heads
+        else:
+            assert embed_dim % head_dim == 0
+            num_heads = embed_dim // head_dim
+        self.feat_size = to_2tuple(feat_size)
+        self.seq_len = self.feat_size[0] * self.feat_size[1]
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.pool_type = pool_type
+        self.scale = self.head_dim ** -0.5
+        self.fused_attn = use_fused_attn()
+
+        if class_token:
+            self.cls_token = nn.Parameter(torch.zeros(1, embed_dim))
+        else:
+            self.cls_token = None
+
+        if qkv_separate:
+            self.q = nn.Linear(in_features, embed_dim, bias=qkv_bias)
+            self.k = nn.Linear(in_features, embed_dim, bias=qkv_bias)
+            self.v = nn.Linear(in_features, embed_dim, bias=qkv_bias)
+            self.qkv = None
+        else:
+            self.q = self.k = self.v = None
+            self.qkv = nn.Linear(in_features, embed_dim * 3, bias=qkv_bias)
+        self.drop = nn.Dropout(drop_rate)
+        self.proj = nn.Linear(embed_dim, self.out_features)
+        self.pos_embed = nn.Parameter(torch.zeros(self.seq_len + 1, in_features))
+
+        self.init_weights()
+
+    def init_weights(self, zero_init_last: bool = False):
+        if self.qkv is None:
+            in_features = self.q.in_features
+            trunc_normal_(self.q.weight, std=in_features ** -0.5)
+            nn.init.zeros_(self.q.bias)
+            trunc_normal_(self.k.weight, std=in_features ** -0.5)
+            nn.init.zeros_(self.k.bias)
+            trunc_normal_(self.v.weight, std=in_features ** -0.5)
+            nn.init.zeros_(self.v.bias)
+        else:
+            in_features = self.qkv.in_features
+            trunc_normal_(self.qkv.weight, std=in_features ** -0.5)
+            nn.init.zeros_(self.qkv.bias)
+        trunc_normal_(self.pos_embed, std=in_features ** -0.5)
+
+    def reset(self, num_classes: Optional[int] = None, pool_type: Optional[str] = None):
+        # NOTE: this module is being used as a head, so need compatible reset()
+        if pool_type is not None:
+            assert pool_type in ('', 'token')
+            self.pool_type = pool_type
+        if num_classes is not None:
+            self.proj = nn.Linear(self.in_features, num_classes) if num_classes > 0 else nn.Identity()
+            self.out_features = num_classes if num_classes > 0 else self.embed_dim
+
+    def _pool(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
+        if self.pool_type == 'token':
+            x = x[:, 0]
+        else:
+            # if not pooled, return spatial output without token
+            x = x[:, 1:].reshape(x.shape[0], H, W, -1).permute(0, 3, 1, 2)
+        return x
+
+    def forward(self, x, pre_logits: bool = False):
+        B, _, H, W = x.shape
+        N = H * W
+        x = x.flatten(2).transpose(1, 2)
+        if self.cls_token is None:
+            x = torch.cat([x.mean(1, keepdim=True), x], dim=1)
+        else:
+            x = torch.cat([self.cls_token.expand(x.shape[0], -1, -1), x], dim=1)
+        pos_embed = resample_abs_pos_embed(self.pos_embed.unsqueeze(0), (H, W), num_prefix_tokens=1)
+        x = x + pos_embed
+
+        if self.qkv is None:
+            q = self.q(x).reshape(B, N + 1, self.num_heads, self.head_dim).transpose(1, 2)
+            k = self.k(x).reshape(B, N + 1, self.num_heads, self.head_dim).transpose(1, 2)
+            v = self.v(x).reshape(B, N + 1, self.num_heads, self.head_dim).transpose(1, 2)
+        else:
+            x = self.qkv(x).reshape(B, -1, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+            q, k, v = x.unbind(0)
+
+        if self.fused_attn:
+            x = nn.functional.scaled_dot_product_attention(q, k, v)
+        else:
+            q = q * self.scale
+            attn = q @ k.transpose(-2, -1)
+            attn = attn.softmax(dim=-1)
+            x = attn @ v
+        x = x.transpose(1, 2).reshape(B, N + 1, -1)
+        x = self.drop(x)
+        if pre_logits:
+            x = self._pool(x, H, W)
+            return x
+        x = self.proj(x)
+        x = self._pool(x, H, W)
+        return x

reference_medclipseg/biomedclip/layers/blur_pool.py

@@ -0,0 +1,91 @@
+"""
+BlurPool layer inspired by
+ - Kornia's Max_BlurPool2d
+ - Making Convolutional Networks Shift-Invariant Again :cite:`zhang2019shiftinvar`
+
+Hacked together by Chris Ha and Ross Wightman
+"""
+from functools import partial
+from typing import Optional, Type
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from .padding import get_padding
+from .typing import LayerType
+
+
+class BlurPool2d(nn.Module):
+    r"""Creates a module that computes blurs and downsample a given feature map.
+    See :cite:`zhang2019shiftinvar` for more details.
+    Corresponds to the Downsample class, which does blurring and subsampling
+
+    Args:
+        channels = Number of input channels
+        filt_size (int): binomial filter size for blurring. currently supports 3 (default) and 5.
+        stride (int): downsampling filter stride
+
+    Returns:
+        torch.Tensor: the transformed tensor.
+    """
+    def __init__(
+            self,
+            channels: Optional[int] = None,
+            filt_size: int = 3,
+            stride: int = 2,
+            pad_mode: str = 'reflect',
+    ) -> None:
+        super(BlurPool2d, self).__init__()
+        assert filt_size > 1
+        self.channels = channels
+        self.filt_size = filt_size
+        self.stride = stride
+        self.pad_mode = pad_mode
+        self.padding = [get_padding(filt_size, stride, dilation=1)] * 4
+
+        coeffs = torch.tensor((np.poly1d((0.5, 0.5)) ** (self.filt_size - 1)).coeffs.astype(np.float32))
+        blur_filter = (coeffs[:, None] * coeffs[None, :])[None, None, :, :]
+        if channels is not None:
+            blur_filter = blur_filter.repeat(self.channels, 1, 1, 1)
+        self.register_buffer('filt', blur_filter, persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.pad(x, self.padding, mode=self.pad_mode)
+        if self.channels is None:
+            channels = x.shape[1]
+            weight = self.filt.expand(channels, 1, self.filt_size, self.filt_size)
+        else:
+            channels = self.channels
+            weight = self.filt
+        return F.conv2d(x, weight, stride=self.stride, groups=channels)
+
+
+def create_aa(
+        aa_layer: LayerType,
+        channels: Optional[int] = None,
+        stride: int = 2,
+        enable: bool = True,
+        noop: Optional[Type[nn.Module]] = nn.Identity
+) -> nn.Module:
+    """ Anti-aliasing """
+    if not aa_layer or not enable:
+        return noop() if noop is not None else None
+
+    if isinstance(aa_layer, str):
+        aa_layer = aa_layer.lower().replace('_', '').replace('-', '')
+        if aa_layer == 'avg' or aa_layer == 'avgpool':
+            aa_layer = nn.AvgPool2d
+        elif aa_layer == 'blur' or aa_layer == 'blurpool':
+            aa_layer = BlurPool2d
+        elif aa_layer == 'blurpc':
+            aa_layer = partial(BlurPool2d, pad_mode='constant')
+
+        else:
+            assert False, f"Unknown anti-aliasing layer ({aa_layer})."
+
+    try:
+        return aa_layer(channels=channels, stride=stride)
+    except TypeError as e:
+        return aa_layer(stride)

reference_medclipseg/biomedclip/layers/bottleneck_attn.py

@@ -0,0 +1,157 @@
+""" Bottleneck Self Attention (Bottleneck Transformers)
+
+Paper: `Bottleneck Transformers for Visual Recognition` - https://arxiv.org/abs/2101.11605
+
+@misc{2101.11605,
+Author = {Aravind Srinivas and Tsung-Yi Lin and Niki Parmar and Jonathon Shlens and Pieter Abbeel and Ashish Vaswani},
+Title = {Bottleneck Transformers for Visual Recognition},
+Year = {2021},
+}
+
+Based on ref gist at: https://gist.github.com/aravindsrinivas/56359b79f0ce4449bcb04ab4b56a57a2
+
+This impl is a WIP but given that it is based on the ref gist likely not too far off.
+
+Hacked together by / Copyright 2021 Ross Wightman
+"""
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .helpers import to_2tuple, make_divisible
+from .weight_init import trunc_normal_
+from .trace_utils import _assert
+
+
+def rel_logits_1d(q, rel_k, permute_mask: List[int]):
+    """ Compute relative logits along one dimension
+
+    As per: https://gist.github.com/aravindsrinivas/56359b79f0ce4449bcb04ab4b56a57a2
+    Originally from: `Attention Augmented Convolutional Networks` - https://arxiv.org/abs/1904.09925
+
+    Args:
+        q: (batch, heads, height, width, dim)
+        rel_k: (2 * width - 1, dim)
+        permute_mask: permute output dim according to this
+    """
+    B, H, W, dim = q.shape
+    x = (q @ rel_k.transpose(-1, -2))
+    x = x.reshape(-1, W, 2 * W -1)
+
+    # pad to shift from relative to absolute indexing
+    x_pad = F.pad(x, [0, 1]).flatten(1)
+    x_pad = F.pad(x_pad, [0, W - 1])
+
+    # reshape and slice out the padded elements
+    x_pad = x_pad.reshape(-1, W + 1, 2 * W - 1)
+    x = x_pad[:, :W, W - 1:]
+
+    # reshape and tile
+    x = x.reshape(B, H, 1, W, W).expand(-1, -1, H, -1, -1)
+    return x.permute(permute_mask)
+
+
+class PosEmbedRel(nn.Module):
+    """ Relative Position Embedding
+    As per: https://gist.github.com/aravindsrinivas/56359b79f0ce4449bcb04ab4b56a57a2
+    Originally from: `Attention Augmented Convolutional Networks` - https://arxiv.org/abs/1904.09925
+    """
+    def __init__(self, feat_size, dim_head, scale):
+        super().__init__()
+        self.height, self.width = to_2tuple(feat_size)
+        self.dim_head = dim_head
+        self.height_rel = nn.Parameter(torch.randn(self.height * 2 - 1, dim_head) * scale)
+        self.width_rel = nn.Parameter(torch.randn(self.width * 2 - 1, dim_head) * scale)
+
+    def forward(self, q):
+        B, HW, _ = q.shape
+
+        # relative logits in width dimension.
+        q = q.reshape(B, self.height, self.width, -1)
+        rel_logits_w = rel_logits_1d(q, self.width_rel, permute_mask=(0, 1, 3, 2, 4))
+
+        # relative logits in height dimension.
+        q = q.transpose(1, 2)
+        rel_logits_h = rel_logits_1d(q, self.height_rel, permute_mask=(0, 3, 1, 4, 2))
+
+        rel_logits = rel_logits_h + rel_logits_w
+        rel_logits = rel_logits.reshape(B, HW, HW)
+        return rel_logits
+
+
+class BottleneckAttn(nn.Module):
+    """ Bottleneck Attention
+    Paper: `Bottleneck Transformers for Visual Recognition` - https://arxiv.org/abs/2101.11605
+
+    The internal dimensions of the attention module are controlled by the interaction of several arguments.
+      * the output dimension of the module is specified by dim_out, which falls back to input dim if not set
+      * the value (v) dimension is set to dim_out // num_heads, the v projection determines the output dim
+      * the query and key (qk) dimensions are determined by
+        * num_heads * dim_head if dim_head is not None
+        * num_heads * (dim_out * attn_ratio // num_heads) if dim_head is None
+      * as seen above, attn_ratio determines the ratio of q and k relative to the output if dim_head not used
+
+    Args:
+        dim (int): input dimension to the module
+        dim_out (int): output dimension of the module, same as dim if not set
+        stride (int): output stride of the module, avg pool used if stride == 2 (default: 1).
+        num_heads (int): parallel attention heads (default: 4)
+        dim_head (int): dimension of query and key heads, calculated from dim_out * attn_ratio // num_heads if not set
+        qk_ratio (float): ratio of q and k dimensions to output dimension when dim_head not set. (default: 1.0)
+        qkv_bias (bool): add bias to q, k, and v projections
+        scale_pos_embed (bool): scale the position embedding as well as Q @ K
+    """
+    def __init__(
+            self, dim, dim_out=None, feat_size=None, stride=1, num_heads=4, dim_head=None,
+            qk_ratio=1.0, qkv_bias=False, scale_pos_embed=False):
+        super().__init__()
+        assert feat_size is not None, 'A concrete feature size matching expected input (H, W) is required'
+        dim_out = dim_out or dim
+        assert dim_out % num_heads == 0
+        self.num_heads = num_heads
+        self.dim_head_qk = dim_head or make_divisible(dim_out * qk_ratio, divisor=8) // num_heads
+        self.dim_head_v = dim_out // self.num_heads
+        self.dim_out_qk = num_heads * self.dim_head_qk
+        self.dim_out_v = num_heads * self.dim_head_v
+        self.scale = self.dim_head_qk ** -0.5
+        self.scale_pos_embed = scale_pos_embed
+
+        self.qkv = nn.Conv2d(dim, self.dim_out_qk * 2 + self.dim_out_v, 1, bias=qkv_bias)
+
+        # NOTE I'm only supporting relative pos embedding for now
+        self.pos_embed = PosEmbedRel(feat_size, dim_head=self.dim_head_qk, scale=self.scale)
+
+        self.pool = nn.AvgPool2d(2, 2) if stride == 2 else nn.Identity()
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        trunc_normal_(self.qkv.weight, std=self.qkv.weight.shape[1] ** -0.5)  # fan-in
+        trunc_normal_(self.pos_embed.height_rel, std=self.scale)
+        trunc_normal_(self.pos_embed.width_rel, std=self.scale)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        _assert(H == self.pos_embed.height, '')
+        _assert(W == self.pos_embed.width, '')
+
+        x = self.qkv(x)  # B, (2 * dim_head_qk + dim_head_v) * num_heads, H, W
+
+        # NOTE head vs channel split ordering in qkv projection was decided before I allowed qk to differ from v
+        # So, this is more verbose than if heads were before qkv splits, but throughput is not impacted.
+        q, k, v = torch.split(x, [self.dim_out_qk, self.dim_out_qk, self.dim_out_v], dim=1)
+        q = q.reshape(B * self.num_heads, self.dim_head_qk, -1).transpose(-1, -2)
+        k = k.reshape(B * self.num_heads, self.dim_head_qk, -1)  # no transpose, for q @ k
+        v = v.reshape(B * self.num_heads, self.dim_head_v, -1).transpose(-1, -2)
+
+        if self.scale_pos_embed:
+            attn = (q @ k + self.pos_embed(q)) * self.scale  # B * num_heads, H * W, H * W
+        else:
+            attn = (q @ k) * self.scale + self.pos_embed(q)
+        attn = attn.softmax(dim=-1)
+
+        out = (attn @ v).transpose(-1, -2).reshape(B, self.dim_out_v, H, W)  # B, dim_out, H, W
+        out = self.pool(out)
+        return out

reference_medclipseg/biomedclip/layers/cbam.py

@@ -0,0 +1,112 @@
+""" CBAM (sort-of) Attention
+
+Experimental impl of CBAM: Convolutional Block Attention Module: https://arxiv.org/abs/1807.06521
+
+WARNING: Results with these attention layers have been mixed. They can significantly reduce performance on
+some tasks, especially fine-grained it seems. I may end up removing this impl.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch
+from torch import nn as nn
+import torch.nn.functional as F
+
+from .conv_bn_act import ConvNormAct
+from .create_act import create_act_layer, get_act_layer
+from .helpers import make_divisible
+
+
+class ChannelAttn(nn.Module):
+    """ Original CBAM channel attention module, currently avg + max pool variant only.
+    """
+    def __init__(
+            self, channels, rd_ratio=1./16, rd_channels=None, rd_divisor=1,
+            act_layer=nn.ReLU, gate_layer='sigmoid', mlp_bias=False):
+        super(ChannelAttn, self).__init__()
+        if not rd_channels:
+            rd_channels = make_divisible(channels * rd_ratio, rd_divisor, round_limit=0.)
+        self.fc1 = nn.Conv2d(channels, rd_channels, 1, bias=mlp_bias)
+        self.act = act_layer(inplace=True)
+        self.fc2 = nn.Conv2d(rd_channels, channels, 1, bias=mlp_bias)
+        self.gate = create_act_layer(gate_layer)
+
+    def forward(self, x):
+        x_avg = self.fc2(self.act(self.fc1(x.mean((2, 3), keepdim=True))))
+        x_max = self.fc2(self.act(self.fc1(x.amax((2, 3), keepdim=True))))
+        return x * self.gate(x_avg + x_max)
+
+
+class LightChannelAttn(ChannelAttn):
+    """An experimental 'lightweight' that sums avg + max pool first
+    """
+    def __init__(
+            self, channels, rd_ratio=1./16, rd_channels=None, rd_divisor=1,
+            act_layer=nn.ReLU, gate_layer='sigmoid', mlp_bias=False):
+        super(LightChannelAttn, self).__init__(
+            channels, rd_ratio, rd_channels, rd_divisor, act_layer, gate_layer, mlp_bias)
+
+    def forward(self, x):
+        x_pool = 0.5 * x.mean((2, 3), keepdim=True) + 0.5 * x.amax((2, 3), keepdim=True)
+        x_attn = self.fc2(self.act(self.fc1(x_pool)))
+        return x * F.sigmoid(x_attn)
+
+
+class SpatialAttn(nn.Module):
+    """ Original CBAM spatial attention module
+    """
+    def __init__(self, kernel_size=7, gate_layer='sigmoid'):
+        super(SpatialAttn, self).__init__()
+        self.conv = ConvNormAct(2, 1, kernel_size, apply_act=False)
+        self.gate = create_act_layer(gate_layer)
+
+    def forward(self, x):
+        x_attn = torch.cat([x.mean(dim=1, keepdim=True), x.amax(dim=1, keepdim=True)], dim=1)
+        x_attn = self.conv(x_attn)
+        return x * self.gate(x_attn)
+
+
+class LightSpatialAttn(nn.Module):
+    """An experimental 'lightweight' variant that sums avg_pool and max_pool results.
+    """
+    def __init__(self, kernel_size=7, gate_layer='sigmoid'):
+        super(LightSpatialAttn, self).__init__()
+        self.conv = ConvNormAct(1, 1, kernel_size, apply_act=False)
+        self.gate = create_act_layer(gate_layer)
+
+    def forward(self, x):
+        x_attn = 0.5 * x.mean(dim=1, keepdim=True) + 0.5 * x.amax(dim=1, keepdim=True)
+        x_attn = self.conv(x_attn)
+        return x * self.gate(x_attn)
+
+
+class CbamModule(nn.Module):
+    def __init__(
+            self, channels, rd_ratio=1./16, rd_channels=None, rd_divisor=1,
+            spatial_kernel_size=7, act_layer=nn.ReLU, gate_layer='sigmoid', mlp_bias=False):
+        super(CbamModule, self).__init__()
+        self.channel = ChannelAttn(
+            channels, rd_ratio=rd_ratio, rd_channels=rd_channels,
+            rd_divisor=rd_divisor, act_layer=act_layer, gate_layer=gate_layer, mlp_bias=mlp_bias)
+        self.spatial = SpatialAttn(spatial_kernel_size, gate_layer=gate_layer)
+
+    def forward(self, x):
+        x = self.channel(x)
+        x = self.spatial(x)
+        return x
+
+
+class LightCbamModule(nn.Module):
+    def __init__(
+            self, channels, rd_ratio=1./16, rd_channels=None, rd_divisor=1,
+            spatial_kernel_size=7, act_layer=nn.ReLU, gate_layer='sigmoid', mlp_bias=False):
+        super(LightCbamModule, self).__init__()
+        self.channel = LightChannelAttn(
+            channels, rd_ratio=rd_ratio, rd_channels=rd_channels,
+            rd_divisor=rd_divisor, act_layer=act_layer, gate_layer=gate_layer, mlp_bias=mlp_bias)
+        self.spatial = LightSpatialAttn(spatial_kernel_size)
+
+    def forward(self, x):
+        x = self.channel(x)
+        x = self.spatial(x)
+        return x
+

reference_medclipseg/biomedclip/layers/classifier.py

@@ -0,0 +1,283 @@
+""" Classifier head and layer factory
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from collections import OrderedDict
+from functools import partial
+from typing import Optional, Union, Callable
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from .adaptive_avgmax_pool import SelectAdaptivePool2d
+from .create_act import get_act_layer
+from .create_norm import get_norm_layer
+
+
+def _create_pool(
+        num_features: int,
+        num_classes: int,
+        pool_type: str = 'avg',
+        use_conv: bool = False,
+        input_fmt: Optional[str] = None,
+):
+    flatten_in_pool = not use_conv  # flatten when we use a Linear layer after pooling
+    if not pool_type:
+        flatten_in_pool = False  # disable flattening if pooling is pass-through (no pooling)
+    global_pool = SelectAdaptivePool2d(
+        pool_type=pool_type,
+        flatten=flatten_in_pool,
+        input_fmt=input_fmt,
+    )
+    num_pooled_features = num_features * global_pool.feat_mult()
+    return global_pool, num_pooled_features
+
+
+def _create_fc(num_features, num_classes, use_conv=False):
+    if num_classes <= 0:
+        fc = nn.Identity()  # pass-through (no classifier)
+    elif use_conv:
+        fc = nn.Conv2d(num_features, num_classes, 1, bias=True)
+    else:
+        fc = nn.Linear(num_features, num_classes, bias=True)
+    return fc
+
+
+def create_classifier(
+        num_features: int,
+        num_classes: int,
+        pool_type: str = 'avg',
+        use_conv: bool = False,
+        input_fmt: str = 'NCHW',
+        drop_rate: Optional[float] = None,
+):
+    global_pool, num_pooled_features = _create_pool(
+        num_features,
+        num_classes,
+        pool_type,
+        use_conv=use_conv,
+        input_fmt=input_fmt,
+    )
+    fc = _create_fc(
+        num_pooled_features,
+        num_classes,
+        use_conv=use_conv,
+    )
+    if drop_rate is not None:
+        dropout = nn.Dropout(drop_rate)
+        return global_pool, dropout, fc
+    return global_pool, fc
+
+
+class ClassifierHead(nn.Module):
+    """Classifier head w/ configurable global pooling and dropout."""
+
+    def __init__(
+            self,
+            in_features: int,
+            num_classes: int,
+            pool_type: str = 'avg',
+            drop_rate: float = 0.,
+            use_conv: bool = False,
+            input_fmt: str = 'NCHW',
+    ):
+        """
+        Args:
+            in_features: The number of input features.
+            num_classes:  The number of classes for the final classifier layer (output).
+            pool_type: Global pooling type, pooling disabled if empty string ('').
+            drop_rate: Pre-classifier dropout rate.
+        """
+        super(ClassifierHead, self).__init__()
+        self.in_features = in_features
+        self.use_conv = use_conv
+        self.input_fmt = input_fmt
+
+        global_pool, fc = create_classifier(
+            in_features,
+            num_classes,
+            pool_type,
+            use_conv=use_conv,
+            input_fmt=input_fmt,
+        )
+        self.global_pool = global_pool
+        self.drop = nn.Dropout(drop_rate)
+        self.fc = fc
+        self.flatten = nn.Flatten(1) if use_conv and pool_type else nn.Identity()
+
+    def reset(self, num_classes: int, pool_type: Optional[str] = None):
+        if pool_type is not None and pool_type != self.global_pool.pool_type:
+            self.global_pool, self.fc = create_classifier(
+                self.in_features,
+                num_classes,
+                pool_type=pool_type,
+                use_conv=self.use_conv,
+                input_fmt=self.input_fmt,
+            )
+            self.flatten = nn.Flatten(1) if self.use_conv and pool_type else nn.Identity()
+        else:
+            num_pooled_features = self.in_features * self.global_pool.feat_mult()
+            self.fc = _create_fc(
+                num_pooled_features,
+                num_classes,
+                use_conv=self.use_conv,
+            )
+
+    def forward(self, x, pre_logits: bool = False):
+        x = self.global_pool(x)
+        x = self.drop(x)
+        if pre_logits:
+            return self.flatten(x)
+        x = self.fc(x)
+        return self.flatten(x)
+
+
+class NormMlpClassifierHead(nn.Module):
+    """ A Pool -> Norm -> Mlp Classifier Head for '2D' NCHW tensors
+    """
+    def __init__(
+            self,
+            in_features: int,
+            num_classes: int,
+            hidden_size: Optional[int] = None,
+            pool_type: str = 'avg',
+            drop_rate: float = 0.,
+            norm_layer: Union[str, Callable] = 'layernorm2d',
+            act_layer: Union[str, Callable] = 'tanh',
+    ):
+        """
+        Args:
+            in_features: The number of input features.
+            num_classes:  The number of classes for the final classifier layer (output).
+            hidden_size: The hidden size of the MLP (pre-logits FC layer) if not None.
+            pool_type: Global pooling type, pooling disabled if empty string ('').
+            drop_rate: Pre-classifier dropout rate.
+            norm_layer: Normalization layer type.
+            act_layer: MLP activation layer type (only used if hidden_size is not None).
+        """
+        super().__init__()
+        self.in_features = in_features
+        self.hidden_size = hidden_size
+        self.num_features = in_features
+        self.use_conv = not pool_type
+        norm_layer = get_norm_layer(norm_layer)
+        act_layer = get_act_layer(act_layer)
+        linear_layer = partial(nn.Conv2d, kernel_size=1) if self.use_conv else nn.Linear
+
+        self.global_pool = SelectAdaptivePool2d(pool_type=pool_type)
+        self.norm = norm_layer(in_features)
+        self.flatten = nn.Flatten(1) if pool_type else nn.Identity()
+        if hidden_size:
+            self.pre_logits = nn.Sequential(OrderedDict([
+                ('fc', linear_layer(in_features, hidden_size)),
+                ('act', act_layer()),
+            ]))
+            self.num_features = hidden_size
+        else:
+            self.pre_logits = nn.Identity()
+        self.drop = nn.Dropout(drop_rate)
+        self.fc = linear_layer(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+    def reset(self, num_classes: int, pool_type: Optional[str] = None):
+        if pool_type is not None:
+            self.global_pool = SelectAdaptivePool2d(pool_type=pool_type)
+            self.flatten = nn.Flatten(1) if pool_type else nn.Identity()
+        self.use_conv = self.global_pool.is_identity()
+        linear_layer = partial(nn.Conv2d, kernel_size=1) if self.use_conv else nn.Linear
+        if self.hidden_size:
+            if ((isinstance(self.pre_logits.fc, nn.Conv2d) and not self.use_conv) or
+                    (isinstance(self.pre_logits.fc, nn.Linear) and self.use_conv)):
+                with torch.no_grad():
+                    new_fc = linear_layer(self.in_features, self.hidden_size)
+                    new_fc.weight.copy_(self.pre_logits.fc.weight.reshape(new_fc.weight.shape))
+                    new_fc.bias.copy_(self.pre_logits.fc.bias)
+                    self.pre_logits.fc = new_fc
+        self.fc = linear_layer(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward(self, x, pre_logits: bool = False):
+        x = self.global_pool(x)
+        x = self.norm(x)
+        x = self.flatten(x)
+        x = self.pre_logits(x)
+        x = self.drop(x)
+        if pre_logits:
+            return x
+        x = self.fc(x)
+        return x
+
+
+class ClNormMlpClassifierHead(nn.Module):
+    """ A Pool -> Norm -> Mlp Classifier Head for n-D NxxC tensors
+    """
+    def __init__(
+            self,
+            in_features: int,
+            num_classes: int,
+            hidden_size: Optional[int] = None,
+            pool_type: str = 'avg',
+            drop_rate: float = 0.,
+            norm_layer: Union[str, Callable] = 'layernorm',
+            act_layer: Union[str, Callable] = 'gelu',
+            input_fmt: str = 'NHWC',
+    ):
+        """
+        Args:
+            in_features: The number of input features.
+            num_classes:  The number of classes for the final classifier layer (output).
+            hidden_size: The hidden size of the MLP (pre-logits FC layer) if not None.
+            pool_type: Global pooling type, pooling disabled if empty string ('').
+            drop_rate: Pre-classifier dropout rate.
+            norm_layer: Normalization layer type.
+            act_layer: MLP activation layer type (only used if hidden_size is not None).
+        """
+        super().__init__()
+        self.in_features = in_features
+        self.hidden_size = hidden_size
+        self.num_features = in_features
+        assert pool_type in ('', 'avg', 'max', 'avgmax')
+        self.pool_type = pool_type
+        assert input_fmt in ('NHWC', 'NLC')
+        self.pool_dim = 1 if input_fmt == 'NLC' else (1, 2)
+        norm_layer = get_norm_layer(norm_layer)
+        act_layer = get_act_layer(act_layer)
+
+        self.norm = norm_layer(in_features)
+        if hidden_size:
+            self.pre_logits = nn.Sequential(OrderedDict([
+                ('fc', nn.Linear(in_features, hidden_size)),
+                ('act', act_layer()),
+            ]))
+            self.num_features = hidden_size
+        else:
+            self.pre_logits = nn.Identity()
+        self.drop = nn.Dropout(drop_rate)
+        self.fc = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+    def reset(self, num_classes: int, pool_type: Optional[str] = None, reset_other: bool = False):
+        if pool_type is not None:
+            self.pool_type = pool_type
+        if reset_other:
+            self.pre_logits = nn.Identity()
+            self.norm = nn.Identity()
+        self.fc = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+    def _global_pool(self, x):
+        if self.pool_type:
+            if self.pool_type == 'avg':
+                x = x.mean(dim=self.pool_dim)
+            elif self.pool_type == 'max':
+                x = x.amax(dim=self.pool_dim)
+            elif self.pool_type == 'avgmax':
+                x = 0.5 * (x.amax(dim=self.pool_dim) + x.mean(dim=self.pool_dim))
+        return x
+
+    def forward(self, x, pre_logits: bool = False):
+        x = self._global_pool(x)
+        x = self.norm(x)
+        x = self.pre_logits(x)
+        x = self.drop(x)
+        if pre_logits:
+            return x
+        x = self.fc(x)
+        return x

reference_medclipseg/biomedclip/layers/cond_conv2d.py

@@ -0,0 +1,123 @@
+""" PyTorch Conditionally Parameterized Convolution (CondConv)
+
+Paper: CondConv: Conditionally Parameterized Convolutions for Efficient Inference
+(https://arxiv.org/abs/1904.04971)
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+
+import math
+from functools import partial
+import numpy as np
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+from .helpers import to_2tuple
+from .conv2d_same import conv2d_same
+from .padding import get_padding_value
+
+
+def get_condconv_initializer(initializer, num_experts, expert_shape):
+    def condconv_initializer(weight):
+        """CondConv initializer function."""
+        num_params = np.prod(expert_shape)
+        if (len(weight.shape) != 2 or weight.shape[0] != num_experts or
+                weight.shape[1] != num_params):
+            raise (ValueError(
+                'CondConv variables must have shape [num_experts, num_params]'))
+        for i in range(num_experts):
+            initializer(weight[i].view(expert_shape))
+    return condconv_initializer
+
+
+class CondConv2d(nn.Module):
+    """ Conditionally Parameterized Convolution
+    Inspired by: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/condconv/condconv_layers.py
+
+    Grouped convolution hackery for parallel execution of the per-sample kernel filters inspired by this discussion:
+    https://github.com/pytorch/pytorch/issues/17983
+    """
+    __constants__ = ['in_channels', 'out_channels', 'dynamic_padding']
+
+    def __init__(self, in_channels, out_channels, kernel_size=3,
+                 stride=1, padding='', dilation=1, groups=1, bias=False, num_experts=4):
+        super(CondConv2d, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = to_2tuple(kernel_size)
+        self.stride = to_2tuple(stride)
+        padding_val, is_padding_dynamic = get_padding_value(
+            padding, kernel_size, stride=stride, dilation=dilation)
+        self.dynamic_padding = is_padding_dynamic  # if in forward to work with torchscript
+        self.padding = to_2tuple(padding_val)
+        self.dilation = to_2tuple(dilation)
+        self.groups = groups
+        self.num_experts = num_experts
+
+        self.weight_shape = (self.out_channels, self.in_channels // self.groups) + self.kernel_size
+        weight_num_param = 1
+        for wd in self.weight_shape:
+            weight_num_param *= wd
+        self.weight = torch.nn.Parameter(torch.Tensor(self.num_experts, weight_num_param))
+
+        if bias:
+            self.bias_shape = (self.out_channels,)
+            self.bias = torch.nn.Parameter(torch.Tensor(self.num_experts, self.out_channels))
+        else:
+            self.register_parameter('bias', None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init_weight = get_condconv_initializer(
+            partial(nn.init.kaiming_uniform_, a=math.sqrt(5)), self.num_experts, self.weight_shape)
+        init_weight(self.weight)
+        if self.bias is not None:
+            fan_in = np.prod(self.weight_shape[1:])
+            bound = 1 / math.sqrt(fan_in)
+            init_bias = get_condconv_initializer(
+                partial(nn.init.uniform_, a=-bound, b=bound), self.num_experts, self.bias_shape)
+            init_bias(self.bias)
+
+    def forward(self, x, routing_weights):
+        B, C, H, W = x.shape
+        weight = torch.matmul(routing_weights, self.weight)
+        new_weight_shape = (B * self.out_channels, self.in_channels // self.groups) + self.kernel_size
+        weight = weight.view(new_weight_shape)
+        bias = None
+        if self.bias is not None:
+            bias = torch.matmul(routing_weights, self.bias)
+            bias = bias.view(B * self.out_channels)
+        # move batch elements with channels so each batch element can be efficiently convolved with separate kernel
+        # reshape instead of view to work with channels_last input
+        x = x.reshape(1, B * C, H, W)
+        if self.dynamic_padding:
+            out = conv2d_same(
+                x, weight, bias, stride=self.stride, padding=self.padding,
+                dilation=self.dilation, groups=self.groups * B)
+        else:
+            out = F.conv2d(
+                x, weight, bias, stride=self.stride, padding=self.padding,
+                dilation=self.dilation, groups=self.groups * B)
+        out = out.permute([1, 0, 2, 3]).view(B, self.out_channels, out.shape[-2], out.shape[-1])
+
+        # Literal port (from TF definition)
+        # x = torch.split(x, 1, 0)
+        # weight = torch.split(weight, 1, 0)
+        # if self.bias is not None:
+        #     bias = torch.matmul(routing_weights, self.bias)
+        #     bias = torch.split(bias, 1, 0)
+        # else:
+        #     bias = [None] * B
+        # out = []
+        # for xi, wi, bi in zip(x, weight, bias):
+        #     wi = wi.view(*self.weight_shape)
+        #     if bi is not None:
+        #         bi = bi.view(*self.bias_shape)
+        #     out.append(self.conv_fn(
+        #         xi, wi, bi, stride=self.stride, padding=self.padding,
+        #         dilation=self.dilation, groups=self.groups))
+        # out = torch.cat(out, 0)
+        return out

reference_medclipseg/biomedclip/layers/config.py

@@ -0,0 +1,165 @@
+""" Model / Layer Config singleton state
+"""
+import os
+import warnings
+from typing import Any, Optional
+
+import torch
+
+__all__ = [
+    'is_exportable', 'is_scriptable', 'is_no_jit', 'use_fused_attn',
+    'set_exportable', 'set_scriptable', 'set_no_jit', 'set_layer_config', 'set_fused_attn',
+    'set_reentrant_ckpt', 'use_reentrant_ckpt'
+]
+
+# Set to True if prefer to have layers with no jit optimization (includes activations)
+_NO_JIT = False
+
+# Set to True if prefer to have activation layers with no jit optimization
+# NOTE not currently used as no difference between no_jit and no_activation jit as only layers obeying
+# the jit flags so far are activations. This will change as more layers are updated and/or added.
+_NO_ACTIVATION_JIT = False
+
+# Set to True if exporting a model with Same padding via ONNX
+_EXPORTABLE = False
+
+# Set to True if wanting to use torch.jit.script on a model
+_SCRIPTABLE = False
+
+
+# use torch.scaled_dot_product_attention where possible
+_HAS_FUSED_ATTN = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+if 'TIMM_FUSED_ATTN' in os.environ:
+    _USE_FUSED_ATTN = int(os.environ['TIMM_FUSED_ATTN'])
+else:
+    _USE_FUSED_ATTN = 1  # 0 == off, 1 == on (for tested use), 2 == on (for experimental use)
+
+
+if 'TIMM_REENTRANT_CKPT' in os.environ:
+    _USE_REENTRANT_CKPT = bool(os.environ['TIMM_REENTRANT_CKPT'])
+else:
+    _USE_REENTRANT_CKPT = False  # defaults to disabled (off)
+
+
+def is_no_jit():
+    return _NO_JIT
+
+
+class set_no_jit:
+    def __init__(self, mode: bool) -> None:
+        global _NO_JIT
+        self.prev = _NO_JIT
+        _NO_JIT = mode
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(self, *args: Any) -> bool:
+        global _NO_JIT
+        _NO_JIT = self.prev
+        return False
+
+
+def is_exportable():
+    return _EXPORTABLE
+
+
+class set_exportable:
+    def __init__(self, mode: bool) -> None:
+        global _EXPORTABLE
+        self.prev = _EXPORTABLE
+        _EXPORTABLE = mode
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(self, *args: Any) -> bool:
+        global _EXPORTABLE
+        _EXPORTABLE = self.prev
+        return False
+
+
+def is_scriptable():
+    return _SCRIPTABLE
+
+
+class set_scriptable:
+    def __init__(self, mode: bool) -> None:
+        global _SCRIPTABLE
+        self.prev = _SCRIPTABLE
+        _SCRIPTABLE = mode
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(self, *args: Any) -> bool:
+        global _SCRIPTABLE
+        _SCRIPTABLE = self.prev
+        return False
+
+
+class set_layer_config:
+    """ Layer config context manager that allows setting all layer config flags at once.
+    If a flag arg is None, it will not change the current value.
+    """
+    def __init__(
+            self,
+            scriptable: Optional[bool] = None,
+            exportable: Optional[bool] = None,
+            no_jit: Optional[bool] = None,
+            no_activation_jit: Optional[bool] = None):
+        global _SCRIPTABLE
+        global _EXPORTABLE
+        global _NO_JIT
+        global _NO_ACTIVATION_JIT
+        self.prev = _SCRIPTABLE, _EXPORTABLE, _NO_JIT, _NO_ACTIVATION_JIT
+        if scriptable is not None:
+            _SCRIPTABLE = scriptable
+        if exportable is not None:
+            _EXPORTABLE = exportable
+        if no_jit is not None:
+            _NO_JIT = no_jit
+        if no_activation_jit is not None:
+            _NO_ACTIVATION_JIT = no_activation_jit
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(self, *args: Any) -> bool:
+        global _SCRIPTABLE
+        global _EXPORTABLE
+        global _NO_JIT
+        global _NO_ACTIVATION_JIT
+        _SCRIPTABLE, _EXPORTABLE, _NO_JIT, _NO_ACTIVATION_JIT = self.prev
+        return False
+
+
+def use_fused_attn(experimental: bool = False) -> bool:
+    # NOTE: ONNX export cannot handle F.scaled_dot_product_attention as of pytorch 2.0
+    if not _HAS_FUSED_ATTN or _EXPORTABLE:
+        return False
+    if experimental:
+        return _USE_FUSED_ATTN > 1
+    return _USE_FUSED_ATTN > 0
+
+
+def set_fused_attn(enable: bool = True, experimental: bool = False):
+    global _USE_FUSED_ATTN
+    if not _HAS_FUSED_ATTN:
+        warnings.warn('This version of pytorch does not have F.scaled_dot_product_attention, fused_attn flag ignored.')
+        return
+    if experimental and enable:
+        _USE_FUSED_ATTN = 2
+    elif enable:
+        _USE_FUSED_ATTN = 1
+    else:
+        _USE_FUSED_ATTN = 0
+
+
+def use_reentrant_ckpt() -> bool:
+    return _USE_REENTRANT_CKPT
+
+
+def set_reentrant_ckpt(enable: bool = True):
+    global _USE_REENTRANT_CKPT
+    _USE_REENTRANT_CKPT = enable

reference_medclipseg/biomedclip/layers/conv2d_same.py

@@ -0,0 +1,110 @@
+""" Conv2d w/ Same Padding
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple, Optional
+
+from .config import is_exportable, is_scriptable
+from .padding import pad_same, pad_same_arg, get_padding_value
+
+
+_USE_EXPORT_CONV = False
+
+
+def conv2d_same(
+        x,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        stride: Tuple[int, int] = (1, 1),
+        padding: Tuple[int, int] = (0, 0),
+        dilation: Tuple[int, int] = (1, 1),
+        groups: int = 1,
+):
+    x = pad_same(x, weight.shape[-2:], stride, dilation)
+    return F.conv2d(x, weight, bias, stride, (0, 0), dilation, groups)
+
+
+class Conv2dSame(nn.Conv2d):
+    """ Tensorflow like 'SAME' convolution wrapper for 2D convolutions
+    """
+
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=1,
+            padding=0,
+            dilation=1,
+            groups=1,
+            bias=True,
+    ):
+        super(Conv2dSame, self).__init__(
+            in_channels, out_channels, kernel_size,
+            stride, 0, dilation, groups, bias,
+        )
+
+    def forward(self, x):
+        return conv2d_same(
+            x, self.weight, self.bias,
+            self.stride, self.padding, self.dilation, self.groups,
+        )
+
+
+class Conv2dSameExport(nn.Conv2d):
+    """ ONNX export friendly Tensorflow like 'SAME' convolution wrapper for 2D convolutions
+
+    NOTE: This does not currently work with torch.jit.script
+    """
+
+    # pylint: disable=unused-argument
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=1,
+            padding=0,
+            dilation=1,
+            groups=1,
+            bias=True,
+    ):
+        super(Conv2dSameExport, self).__init__(
+            in_channels, out_channels, kernel_size,
+            stride, 0, dilation, groups, bias,
+        )
+        self.pad = None
+        self.pad_input_size = (0, 0)
+
+    def forward(self, x):
+        input_size = x.size()[-2:]
+        if self.pad is None:
+            pad_arg = pad_same_arg(input_size, self.weight.size()[-2:], self.stride, self.dilation)
+            self.pad = nn.ZeroPad2d(pad_arg)
+            self.pad_input_size = input_size
+
+        x = self.pad(x)
+        return F.conv2d(
+            x, self.weight, self.bias,
+            self.stride, self.padding, self.dilation, self.groups,
+        )
+
+
+def create_conv2d_pad(in_chs, out_chs, kernel_size, **kwargs):
+    padding = kwargs.pop('padding', '')
+    kwargs.setdefault('bias', False)
+    padding, is_dynamic = get_padding_value(padding, kernel_size, **kwargs)
+    if is_dynamic:
+        if _USE_EXPORT_CONV and is_exportable():
+            # older PyTorch ver needed this to export same padding reasonably
+            assert not is_scriptable()  # Conv2DSameExport does not work with jit
+            return Conv2dSameExport(in_chs, out_chs, kernel_size, **kwargs)
+        else:
+            return Conv2dSame(in_chs, out_chs, kernel_size, **kwargs)
+    else:
+        return nn.Conv2d(in_chs, out_chs, kernel_size, padding=padding, **kwargs)
+
+

reference_medclipseg/biomedclip/layers/conv_bn_act.py

@@ -0,0 +1,92 @@
+""" Conv2d + BN + Act
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from typing import Any, Dict, Optional, Type
+
+from torch import nn as nn
+
+from .typing import LayerType, PadType
+from .blur_pool import create_aa
+from .create_conv2d import create_conv2d
+from .create_norm_act import get_norm_act_layer
+
+
+class ConvNormAct(nn.Module):
+    def __init__(
+            self,
+            in_channels: int,
+            out_channels: int,
+            kernel_size: int = 1,
+            stride: int = 1,
+            padding: PadType = '',
+            dilation: int = 1,
+            groups: int = 1,
+            bias: bool = False,
+            apply_norm: bool = True,
+            apply_act: bool = True,
+            norm_layer: LayerType = nn.BatchNorm2d,
+            act_layer: Optional[LayerType] = nn.ReLU,
+            aa_layer: Optional[LayerType] = None,
+            drop_layer: Optional[Type[nn.Module]] = None,
+            conv_kwargs: Optional[Dict[str, Any]] = None,
+            norm_kwargs: Optional[Dict[str, Any]] = None,
+            act_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        super(ConvNormAct, self).__init__()
+        conv_kwargs = conv_kwargs or {}
+        norm_kwargs = norm_kwargs or {}
+        act_kwargs = act_kwargs or {}
+        use_aa = aa_layer is not None and stride > 1
+
+        self.conv = create_conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=1 if use_aa else stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            **conv_kwargs,
+        )
+
+        if apply_norm:
+            # NOTE for backwards compatibility with models that use separate norm and act layer definitions
+            norm_act_layer = get_norm_act_layer(norm_layer, act_layer)
+            # NOTE for backwards (weight) compatibility, norm layer name remains `.bn`
+            if drop_layer:
+                norm_kwargs['drop_layer'] = drop_layer
+            self.bn = norm_act_layer(
+                out_channels,
+                apply_act=apply_act,
+                act_kwargs=act_kwargs,
+                **norm_kwargs,
+            )
+        else:
+            self.bn = nn.Sequential()
+            if drop_layer:
+                norm_kwargs['drop_layer'] = drop_layer
+                self.bn.add_module('drop', drop_layer())
+
+        self.aa = create_aa(aa_layer, out_channels, stride=stride, enable=use_aa, noop=None)
+
+    @property
+    def in_channels(self):
+        return self.conv.in_channels
+
+    @property
+    def out_channels(self):
+        return self.conv.out_channels
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        aa = getattr(self, 'aa', None)
+        if aa is not None:
+            x = self.aa(x)
+        return x
+
+
+ConvBnAct = ConvNormAct
+ConvNormActAa = ConvNormAct   # backwards compat, when they were separate

reference_medclipseg/biomedclip/layers/create_act.py

@@ -0,0 +1,138 @@
+""" Activation Factory
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from typing import Union, Callable, Type
+
+from .activations import *
+from .activations_me import *
+from .config import is_exportable, is_scriptable
+
+# PyTorch has an optimized, native 'silu' (aka 'swish') operator as of PyTorch 1.7.
+# Also hardsigmoid, hardswish, and soon mish. This code will use native version if present.
+# Eventually, the custom SiLU, Mish, Hard*, layers will be removed and only native variants will be used.
+_has_silu = 'silu' in dir(torch.nn.functional)
+_has_hardswish = 'hardswish' in dir(torch.nn.functional)
+_has_hardsigmoid = 'hardsigmoid' in dir(torch.nn.functional)
+_has_mish = 'mish' in dir(torch.nn.functional)
+
+
+_ACT_FN_DEFAULT = dict(
+    silu=F.silu if _has_silu else swish,
+    swish=F.silu if _has_silu else swish,
+    mish=F.mish if _has_mish else mish,
+    relu=F.relu,
+    relu6=F.relu6,
+    leaky_relu=F.leaky_relu,
+    elu=F.elu,
+    celu=F.celu,
+    selu=F.selu,
+    gelu=gelu,
+    gelu_tanh=gelu_tanh,
+    quick_gelu=quick_gelu,
+    sigmoid=sigmoid,
+    tanh=tanh,
+    hard_sigmoid=F.hardsigmoid if _has_hardsigmoid else hard_sigmoid,
+    hard_swish=F.hardswish if _has_hardswish else hard_swish,
+    hard_mish=hard_mish,
+)
+
+_ACT_FN_ME = dict(
+    silu=F.silu if _has_silu else swish_me,
+    swish=F.silu if _has_silu else swish_me,
+    mish=F.mish if _has_mish else mish_me,
+    hard_sigmoid=F.hardsigmoid if _has_hardsigmoid else hard_sigmoid_me,
+    hard_swish=F.hardswish if _has_hardswish else hard_swish_me,
+    hard_mish=hard_mish_me,
+)
+
+_ACT_FNS = (_ACT_FN_ME, _ACT_FN_DEFAULT)
+for a in _ACT_FNS:
+    a.setdefault('hardsigmoid', a.get('hard_sigmoid'))
+    a.setdefault('hardswish', a.get('hard_swish'))
+
+
+_ACT_LAYER_DEFAULT = dict(
+    silu=nn.SiLU if _has_silu else Swish,
+    swish=nn.SiLU if _has_silu else Swish,
+    mish=nn.Mish if _has_mish else Mish,
+    relu=nn.ReLU,
+    relu6=nn.ReLU6,
+    leaky_relu=nn.LeakyReLU,
+    elu=nn.ELU,
+    prelu=PReLU,
+    celu=nn.CELU,
+    selu=nn.SELU,
+    gelu=GELU,
+    gelu_tanh=GELUTanh,
+    quick_gelu=QuickGELU,
+    sigmoid=Sigmoid,
+    tanh=Tanh,
+    hard_sigmoid=nn.Hardsigmoid if _has_hardsigmoid else HardSigmoid,
+    hard_swish=nn.Hardswish if _has_hardswish else HardSwish,
+    hard_mish=HardMish,
+    identity=nn.Identity,
+)
+
+_ACT_LAYER_ME = dict(
+    silu=nn.SiLU if _has_silu else SwishMe,
+    swish=nn.SiLU if _has_silu else SwishMe,
+    mish=nn.Mish if _has_mish else MishMe,
+    hard_sigmoid=nn.Hardsigmoid if _has_hardsigmoid else HardSigmoidMe,
+    hard_swish=nn.Hardswish if _has_hardswish else HardSwishMe,
+    hard_mish=HardMishMe,
+)
+
+_ACT_LAYERS = (_ACT_LAYER_ME, _ACT_LAYER_DEFAULT)
+for a in _ACT_LAYERS:
+    a.setdefault('hardsigmoid', a.get('hard_sigmoid'))
+    a.setdefault('hardswish', a.get('hard_swish'))
+
+
+def get_act_fn(name: Union[Callable, str] = 'relu'):
+    """ Activation Function Factory
+    Fetching activation fns by name with this function allows export or torch script friendly
+    functions to be returned dynamically based on current config.
+    """
+    if not name:
+        return None
+    if isinstance(name, Callable):
+        return name
+    name = name.lower()
+    if not (is_exportable() or is_scriptable()):
+        # If not exporting or scripting the model, first look for a memory-efficient version with
+        # custom autograd, then fallback
+        if name in _ACT_FN_ME:
+            return _ACT_FN_ME[name]
+    return _ACT_FN_DEFAULT[name]
+
+
+def get_act_layer(name: Union[Type[nn.Module], str] = 'relu'):
+    """ Activation Layer Factory
+    Fetching activation layers by name with this function allows export or torch script friendly
+    functions to be returned dynamically based on current config.
+    """
+    if name is None:
+        return None
+    if not isinstance(name, str):
+        # callable, module, etc
+        return name
+    if not name:
+        return None
+    name = name.lower()
+    if not (is_exportable() or is_scriptable()):
+        if name in _ACT_LAYER_ME:
+            return _ACT_LAYER_ME[name]
+    return _ACT_LAYER_DEFAULT[name]
+
+
+def create_act_layer(name: Union[Type[nn.Module], str], inplace=None, **kwargs):
+    act_layer = get_act_layer(name)
+    if act_layer is None:
+        return None
+    if inplace is None:
+        return act_layer(**kwargs)
+    try:
+        return act_layer(inplace=inplace, **kwargs)
+    except TypeError:
+        # recover if act layer doesn't have inplace arg
+        return act_layer(**kwargs)

reference_medclipseg/biomedclip/layers/create_attn.py

@@ -0,0 +1,89 @@
+""" Attention Factory
+
+Hacked together by / Copyright 2021 Ross Wightman
+"""
+import torch
+from functools import partial
+
+from .bottleneck_attn import BottleneckAttn
+from .cbam import CbamModule, LightCbamModule
+from .eca import EcaModule, CecaModule
+from .gather_excite import GatherExcite
+from .global_context import GlobalContext
+from .halo_attn import HaloAttn
+from .lambda_layer import LambdaLayer
+from .non_local_attn import NonLocalAttn, BatNonLocalAttn
+from .selective_kernel import SelectiveKernel
+from .split_attn import SplitAttn
+from .squeeze_excite import SEModule, EffectiveSEModule
+
+
+def get_attn(attn_type):
+    if isinstance(attn_type, torch.nn.Module):
+        return attn_type
+    module_cls = None
+    if attn_type:
+        if isinstance(attn_type, str):
+            attn_type = attn_type.lower()
+            # Lightweight attention modules (channel and/or coarse spatial).
+            # Typically added to existing network architecture blocks in addition to existing convolutions.
+            if attn_type == 'se':
+                module_cls = SEModule
+            elif attn_type == 'ese':
+                module_cls = EffectiveSEModule
+            elif attn_type == 'eca':
+                module_cls = EcaModule
+            elif attn_type == 'ecam':
+                module_cls = partial(EcaModule, use_mlp=True)
+            elif attn_type == 'ceca':
+                module_cls = CecaModule
+            elif attn_type == 'ge':
+                module_cls = GatherExcite
+            elif attn_type == 'gc':
+                module_cls = GlobalContext
+            elif attn_type == 'gca':
+                module_cls = partial(GlobalContext, fuse_add=True, fuse_scale=False)
+            elif attn_type == 'cbam':
+                module_cls = CbamModule
+            elif attn_type == 'lcbam':
+                module_cls = LightCbamModule
+
+            # Attention / attention-like modules w/ significant params
+            # Typically replace some of the existing workhorse convs in a network architecture.
+            # All of these accept a stride argument and can spatially downsample the input.
+            elif attn_type == 'sk':
+                module_cls = SelectiveKernel
+            elif attn_type == 'splat':
+                module_cls = SplitAttn
+
+            # Self-attention / attention-like modules w/ significant compute and/or params
+            # Typically replace some of the existing workhorse convs in a network architecture.
+            # All of these accept a stride argument and can spatially downsample the input.
+            elif attn_type == 'lambda':
+                return LambdaLayer
+            elif attn_type == 'bottleneck':
+                return BottleneckAttn
+            elif attn_type == 'halo':
+                return HaloAttn
+            elif attn_type == 'nl':
+                module_cls = NonLocalAttn
+            elif attn_type == 'bat':
+                module_cls = BatNonLocalAttn
+
+            # Woops!
+            else:
+                assert False, "Invalid attn module (%s)" % attn_type
+        elif isinstance(attn_type, bool):
+            if attn_type:
+                module_cls = SEModule
+        else:
+            module_cls = attn_type
+    return module_cls
+
+
+def create_attn(attn_type, channels, **kwargs):
+    module_cls = get_attn(attn_type)
+    if module_cls is not None:
+        # NOTE: it's expected the first (positional) argument of all attention layers is the # input channels
+        return module_cls(channels, **kwargs)
+    return None

reference_medclipseg/biomedclip/layers/create_conv2d.py

@@ -0,0 +1,36 @@
+""" Create Conv2d Factory Method
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+
+from .mixed_conv2d import MixedConv2d
+from .cond_conv2d import CondConv2d
+from .conv2d_same import create_conv2d_pad
+
+
+def create_conv2d(in_channels, out_channels, kernel_size, **kwargs):
+    """ Select a 2d convolution implementation based on arguments
+    Creates and returns one of torch.nn.Conv2d, Conv2dSame, MixedConv2d, or CondConv2d.
+
+    Used extensively by EfficientNet, MobileNetv3 and related networks.
+    """
+    if isinstance(kernel_size, list):
+        assert 'num_experts' not in kwargs  # MixNet + CondConv combo not supported currently
+        if 'groups' in kwargs:
+            groups = kwargs.pop('groups')
+            if groups == in_channels:
+                kwargs['depthwise'] = True
+            else:
+                assert groups == 1
+        # We're going to use only lists for defining the MixedConv2d kernel groups,
+        # ints, tuples, other iterables will continue to pass to normal conv and specify h, w.
+        m = MixedConv2d(in_channels, out_channels, kernel_size, **kwargs)
+    else:
+        depthwise = kwargs.pop('depthwise', False)
+        # for DW out_channels must be multiple of in_channels as must have out_channels % groups == 0
+        groups = in_channels if depthwise else kwargs.pop('groups', 1)
+        if 'num_experts' in kwargs and kwargs['num_experts'] > 0:
+            m = CondConv2d(in_channels, out_channels, kernel_size, groups=groups, **kwargs)
+        else:
+            m = create_conv2d_pad(in_channels, out_channels, kernel_size, groups=groups, **kwargs)
+    return m

reference_medclipseg/biomedclip/layers/create_norm.py

@@ -0,0 +1,60 @@
+""" Norm Layer Factory
+
+Create norm modules by string (to mirror create_act and creat_norm-act fns)
+
+Copyright 2022 Ross Wightman
+"""
+import functools
+import types
+from typing import Type
+
+import torch.nn as nn
+
+from .norm import GroupNorm, GroupNorm1, LayerNorm, LayerNorm2d, RmsNorm, RmsNorm2d, SimpleNorm, SimpleNorm2d
+from torchvision.ops.misc import FrozenBatchNorm2d
+
+_NORM_MAP = dict(
+    batchnorm=nn.BatchNorm2d,
+    batchnorm2d=nn.BatchNorm2d,
+    batchnorm1d=nn.BatchNorm1d,
+    groupnorm=GroupNorm,
+    groupnorm1=GroupNorm1,
+    layernorm=LayerNorm,
+    layernorm2d=LayerNorm2d,
+    rmsnorm=RmsNorm,
+    rmsnorm2d=RmsNorm2d,
+    simplenorm=SimpleNorm,
+    simplenorm2d=SimpleNorm2d,
+    frozenbatchnorm2d=FrozenBatchNorm2d,
+)
+_NORM_TYPES = {m for n, m in _NORM_MAP.items()}
+
+
+def create_norm_layer(layer_name, num_features, **kwargs):
+    layer = get_norm_layer(layer_name)
+    layer_instance = layer(num_features, **kwargs)
+    return layer_instance
+
+
+def get_norm_layer(norm_layer):
+    if norm_layer is None:
+        return None
+    assert isinstance(norm_layer, (type, str, types.FunctionType, functools.partial))
+    norm_kwargs = {}
+
+    # unbind partial fn, so args can be rebound later
+    if isinstance(norm_layer, functools.partial):
+        norm_kwargs.update(norm_layer.keywords)
+        norm_layer = norm_layer.func
+
+    if isinstance(norm_layer, str):
+        if not norm_layer:
+            return None
+        layer_name = norm_layer.replace('_', '').lower()
+        norm_layer = _NORM_MAP[layer_name]
+    else:
+        norm_layer = norm_layer
+
+    if norm_kwargs:
+        norm_layer = functools.partial(norm_layer, **norm_kwargs)  # bind/rebind args
+    return norm_layer

reference_medclipseg/biomedclip/layers/create_norm_act.py

@@ -0,0 +1,95 @@
+""" NormAct (Normalization + Activation Layer) Factory
+
+Create norm + act combo modules that attempt to be backwards compatible with separate norm + act
+instances in models. Where these are used it will be possible to swap separate BN + act layers with
+combined modules like IABN or EvoNorms.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import types
+import functools
+
+from .evo_norm import *
+from .filter_response_norm import FilterResponseNormAct2d, FilterResponseNormTlu2d
+from .norm_act import BatchNormAct2d, GroupNormAct, LayerNormAct, LayerNormAct2d
+from .inplace_abn import InplaceAbn
+
+_NORM_ACT_MAP = dict(
+    batchnorm=BatchNormAct2d,
+    batchnorm2d=BatchNormAct2d,
+    groupnorm=GroupNormAct,
+    groupnorm1=functools.partial(GroupNormAct, num_groups=1),
+    layernorm=LayerNormAct,
+    layernorm2d=LayerNormAct2d,
+    evonormb0=EvoNorm2dB0,
+    evonormb1=EvoNorm2dB1,
+    evonormb2=EvoNorm2dB2,
+    evonorms0=EvoNorm2dS0,
+    evonorms0a=EvoNorm2dS0a,
+    evonorms1=EvoNorm2dS1,
+    evonorms1a=EvoNorm2dS1a,
+    evonorms2=EvoNorm2dS2,
+    evonorms2a=EvoNorm2dS2a,
+    frn=FilterResponseNormAct2d,
+    frntlu=FilterResponseNormTlu2d,
+    inplaceabn=InplaceAbn,
+    iabn=InplaceAbn,
+)
+_NORM_ACT_TYPES = {m for n, m in _NORM_ACT_MAP.items()}
+# has act_layer arg to define act type
+_NORM_ACT_REQUIRES_ARG = {
+    BatchNormAct2d, GroupNormAct, LayerNormAct, LayerNormAct2d, FilterResponseNormAct2d, InplaceAbn}
+
+
+def create_norm_act_layer(layer_name, num_features, act_layer=None, apply_act=True, jit=False, **kwargs):
+    layer = get_norm_act_layer(layer_name, act_layer=act_layer)
+    layer_instance = layer(num_features, apply_act=apply_act, **kwargs)
+    if jit:
+        layer_instance = torch.jit.script(layer_instance)
+    return layer_instance
+
+
+def get_norm_act_layer(norm_layer, act_layer=None):
+    if norm_layer is None:
+        return None
+    assert isinstance(norm_layer, (type, str,  types.FunctionType, functools.partial))
+    assert act_layer is None or isinstance(act_layer, (type, str, types.FunctionType, functools.partial))
+    norm_act_kwargs = {}
+
+    # unbind partial fn, so args can be rebound later
+    if isinstance(norm_layer, functools.partial):
+        norm_act_kwargs.update(norm_layer.keywords)
+        norm_layer = norm_layer.func
+
+    if isinstance(norm_layer, str):
+        if not norm_layer:
+            return None
+        layer_name = norm_layer.replace('_', '').lower().split('-')[0]
+        norm_act_layer = _NORM_ACT_MAP[layer_name]
+    elif norm_layer in _NORM_ACT_TYPES:
+        norm_act_layer = norm_layer
+    elif isinstance(norm_layer,  types.FunctionType):
+        # if function type, must be a lambda/fn that creates a norm_act layer
+        norm_act_layer = norm_layer
+    else:
+        type_name = norm_layer.__name__.lower()
+        if type_name.startswith('batchnorm'):
+            norm_act_layer = BatchNormAct2d
+        elif type_name.startswith('groupnorm'):
+            norm_act_layer = GroupNormAct
+        elif type_name.startswith('groupnorm1'):
+            norm_act_layer = functools.partial(GroupNormAct, num_groups=1)
+        elif type_name.startswith('layernorm2d'):
+            norm_act_layer = LayerNormAct2d
+        elif type_name.startswith('layernorm'):
+            norm_act_layer = LayerNormAct
+        else:
+            assert False, f"No equivalent norm_act layer for {type_name}"
+
+    if norm_act_layer in _NORM_ACT_REQUIRES_ARG:
+        # pass `act_layer` through for backwards compat where `act_layer=None` implies no activation.
+        # In the future, may force use of `apply_act` with `act_layer` arg bound to relevant NormAct types
+        norm_act_kwargs.setdefault('act_layer', act_layer)
+    if norm_act_kwargs:
+        norm_act_layer = functools.partial(norm_act_layer, **norm_act_kwargs)  # bind/rebind args
+    return norm_act_layer

reference_medclipseg/biomedclip/layers/drop.py

@@ -0,0 +1,182 @@
+""" DropBlock, DropPath
+
+PyTorch implementations of DropBlock and DropPath (Stochastic Depth) regularization layers.
+
+Papers:
+DropBlock: A regularization method for convolutional networks (https://arxiv.org/abs/1810.12890)
+
+Deep Networks with Stochastic Depth (https://arxiv.org/abs/1603.09382)
+
+Code:
+DropBlock impl inspired by two Tensorflow impl that I liked:
+ - https://github.com/tensorflow/tpu/blob/master/models/official/resnet/resnet_model.py#L74
+ - https://github.com/clovaai/assembled-cnn/blob/master/nets/blocks.py
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .grid import ndgrid
+
+
+def drop_block_2d(
+        x,
+        drop_prob: float = 0.1,
+        block_size: int = 7,
+        gamma_scale: float = 1.0,
+        with_noise: bool = False,
+        inplace: bool = False,
+        batchwise: bool = False
+):
+    """ DropBlock. See https://arxiv.org/pdf/1810.12890.pdf
+
+    DropBlock with an experimental gaussian noise option. This layer has been tested on a few training
+    runs with success, but needs further validation and possibly optimization for lower runtime impact.
+    """
+    B, C, H, W = x.shape
+    total_size = W * H
+    clipped_block_size = min(block_size, min(W, H))
+    # seed_drop_rate, the gamma parameter
+    gamma = gamma_scale * drop_prob * total_size / clipped_block_size ** 2 / (
+            (W - block_size + 1) * (H - block_size + 1))
+
+    # Forces the block to be inside the feature map.
+    w_i, h_i = ndgrid(torch.arange(W, device=x.device), torch.arange(H, device=x.device))
+    valid_block = ((w_i >= clipped_block_size // 2) & (w_i < W - (clipped_block_size - 1) // 2)) & \
+                  ((h_i >= clipped_block_size // 2) & (h_i < H - (clipped_block_size - 1) // 2))
+    valid_block = torch.reshape(valid_block, (1, 1, H, W)).to(dtype=x.dtype)
+
+    if batchwise:
+        # one mask for whole batch, quite a bit faster
+        uniform_noise = torch.rand((1, C, H, W), dtype=x.dtype, device=x.device)
+    else:
+        uniform_noise = torch.rand_like(x)
+    block_mask = ((2 - gamma - valid_block + uniform_noise) >= 1).to(dtype=x.dtype)
+    block_mask = -F.max_pool2d(
+        -block_mask,
+        kernel_size=clipped_block_size,  # block_size,
+        stride=1,
+        padding=clipped_block_size // 2)
+
+    if with_noise:
+        normal_noise = torch.randn((1, C, H, W), dtype=x.dtype, device=x.device) if batchwise else torch.randn_like(x)
+        if inplace:
+            x.mul_(block_mask).add_(normal_noise * (1 - block_mask))
+        else:
+            x = x * block_mask + normal_noise * (1 - block_mask)
+    else:
+        normalize_scale = (block_mask.numel() / block_mask.to(dtype=torch.float32).sum().add(1e-7)).to(x.dtype)
+        if inplace:
+            x.mul_(block_mask * normalize_scale)
+        else:
+            x = x * block_mask * normalize_scale
+    return x
+
+
+def drop_block_fast_2d(
+        x: torch.Tensor,
+        drop_prob: float = 0.1,
+        block_size: int = 7,
+        gamma_scale: float = 1.0,
+        with_noise: bool = False,
+        inplace: bool = False,
+):
+    """ DropBlock. See https://arxiv.org/pdf/1810.12890.pdf
+
+    DropBlock with an experimental gaussian noise option. Simplied from above without concern for valid
+    block mask at edges.
+    """
+    B, C, H, W = x.shape
+    total_size = W * H
+    clipped_block_size = min(block_size, min(W, H))
+    gamma = gamma_scale * drop_prob * total_size / clipped_block_size ** 2 / (
+            (W - block_size + 1) * (H - block_size + 1))
+
+    block_mask = torch.empty_like(x).bernoulli_(gamma)
+    block_mask = F.max_pool2d(
+        block_mask.to(x.dtype), kernel_size=clipped_block_size, stride=1, padding=clipped_block_size // 2)
+
+    if with_noise:
+        normal_noise = torch.empty_like(x).normal_()
+        if inplace:
+            x.mul_(1. - block_mask).add_(normal_noise * block_mask)
+        else:
+            x = x * (1. - block_mask) + normal_noise * block_mask
+    else:
+        block_mask = 1 - block_mask
+        normalize_scale = (block_mask.numel() / block_mask.to(dtype=torch.float32).sum().add(1e-6)).to(dtype=x.dtype)
+        if inplace:
+            x.mul_(block_mask * normalize_scale)
+        else:
+            x = x * block_mask * normalize_scale
+    return x
+
+
+class DropBlock2d(nn.Module):
+    """ DropBlock. See https://arxiv.org/pdf/1810.12890.pdf
+    """
+
+    def __init__(
+            self,
+            drop_prob: float = 0.1,
+            block_size: int = 7,
+            gamma_scale: float = 1.0,
+            with_noise: bool = False,
+            inplace: bool = False,
+            batchwise: bool = False,
+            fast: bool = True):
+        super(DropBlock2d, self).__init__()
+        self.drop_prob = drop_prob
+        self.gamma_scale = gamma_scale
+        self.block_size = block_size
+        self.with_noise = with_noise
+        self.inplace = inplace
+        self.batchwise = batchwise
+        self.fast = fast  # FIXME finish comparisons of fast vs not
+
+    def forward(self, x):
+        if not self.training or not self.drop_prob:
+            return x
+        if self.fast:
+            return drop_block_fast_2d(
+                x, self.drop_prob, self.block_size, self.gamma_scale, self.with_noise, self.inplace)
+        else:
+            return drop_block_2d(
+                x, self.drop_prob, self.block_size, self.gamma_scale, self.with_noise, self.inplace, self.batchwise)
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+    def extra_repr(self):
+        return f'drop_prob={round(self.drop_prob,3):0.3f}'

reference_medclipseg/biomedclip/layers/eca.py

@@ -0,0 +1,145 @@
+"""
+ECA module from ECAnet
+
+paper: ECA-Net: Efficient Channel Attention for Deep Convolutional Neural Networks
+https://arxiv.org/abs/1910.03151
+
+Original ECA model borrowed from https://github.com/BangguWu/ECANet
+
+Modified circular ECA implementation and adaption for use in timm package
+by Chris Ha https://github.com/VRandme
+
+Original License:
+
+MIT License
+
+Copyright (c) 2019 BangguWu, Qilong Wang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+import math
+from torch import nn
+import torch.nn.functional as F
+
+
+from .create_act import create_act_layer
+from .helpers import make_divisible
+
+
+class EcaModule(nn.Module):
+    """Constructs an ECA module.
+
+    Args:
+        channels: Number of channels of the input feature map for use in adaptive kernel sizes
+            for actual calculations according to channel.
+            gamma, beta: when channel is given parameters of mapping function
+            refer to original paper https://arxiv.org/pdf/1910.03151.pdf
+            (default=None. if channel size not given, use k_size given for kernel size.)
+        kernel_size: Adaptive selection of kernel size (default=3)
+        gamm: used in kernel_size calc, see above
+        beta: used in kernel_size calc, see above
+        act_layer: optional non-linearity after conv, enables conv bias, this is an experiment
+        gate_layer: gating non-linearity to use
+    """
+    def __init__(
+            self, channels=None, kernel_size=3, gamma=2, beta=1, act_layer=None, gate_layer='sigmoid',
+            rd_ratio=1/8, rd_channels=None, rd_divisor=8, use_mlp=False):
+        super(EcaModule, self).__init__()
+        if channels is not None:
+            t = int(abs(math.log(channels, 2) + beta) / gamma)
+            kernel_size = max(t if t % 2 else t + 1, 3)
+        assert kernel_size % 2 == 1
+        padding = (kernel_size - 1) // 2
+        if use_mlp:
+            # NOTE 'mlp' mode is a timm experiment, not in paper
+            assert channels is not None
+            if rd_channels is None:
+                rd_channels = make_divisible(channels * rd_ratio, divisor=rd_divisor)
+            act_layer = act_layer or nn.ReLU
+            self.conv = nn.Conv1d(1, rd_channels, kernel_size=1, padding=0, bias=True)
+            self.act = create_act_layer(act_layer)
+            self.conv2 = nn.Conv1d(rd_channels, 1, kernel_size=kernel_size, padding=padding, bias=True)
+        else:
+            self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=padding, bias=False)
+            self.act = None
+            self.conv2 = None
+        self.gate = create_act_layer(gate_layer)
+
+    def forward(self, x):
+        y = x.mean((2, 3)).view(x.shape[0], 1, -1)  # view for 1d conv
+        y = self.conv(y)
+        if self.conv2 is not None:
+            y = self.act(y)
+            y = self.conv2(y)
+        y = self.gate(y).view(x.shape[0], -1, 1, 1)
+        return x * y.expand_as(x)
+
+
+EfficientChannelAttn = EcaModule  # alias
+
+
+class CecaModule(nn.Module):
+    """Constructs a circular ECA module.
+
+    ECA module where the conv uses circular padding rather than zero padding.
+    Unlike the spatial dimension, the channels do not have inherent ordering nor
+    locality. Although this module in essence, applies such an assumption, it is unnecessary
+    to limit the channels on either "edge" from being circularly adapted to each other.
+    This will fundamentally increase connectivity and possibly increase performance metrics
+    (accuracy, robustness), without significantly impacting resource metrics
+    (parameter size, throughput,latency, etc)
+
+    Args:
+        channels: Number of channels of the input feature map for use in adaptive kernel sizes
+            for actual calculations according to channel.
+            gamma, beta: when channel is given parameters of mapping function
+            refer to original paper https://arxiv.org/pdf/1910.03151.pdf
+            (default=None. if channel size not given, use k_size given for kernel size.)
+        kernel_size: Adaptive selection of kernel size (default=3)
+        gamm: used in kernel_size calc, see above
+        beta: used in kernel_size calc, see above
+        act_layer: optional non-linearity after conv, enables conv bias, this is an experiment
+        gate_layer: gating non-linearity to use
+    """
+
+    def __init__(self, channels=None, kernel_size=3, gamma=2, beta=1, act_layer=None, gate_layer='sigmoid'):
+        super(CecaModule, self).__init__()
+        if channels is not None:
+            t = int(abs(math.log(channels, 2) + beta) / gamma)
+            kernel_size = max(t if t % 2 else t + 1, 3)
+        has_act = act_layer is not None
+        assert kernel_size % 2 == 1
+
+        # PyTorch circular padding mode is buggy as of pytorch 1.4
+        # see https://github.com/pytorch/pytorch/pull/17240
+        # implement manual circular padding
+        self.padding = (kernel_size - 1) // 2
+        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=0, bias=has_act)
+        self.gate = create_act_layer(gate_layer)
+
+    def forward(self, x):
+        y = x.mean((2, 3)).view(x.shape[0], 1, -1)
+        # Manually implement circular padding, F.pad does not seemed to be bugged
+        y = F.pad(y, (self.padding, self.padding), mode='circular')
+        y = self.conv(y)
+        y = self.gate(y).view(x.shape[0], -1, 1, 1)
+        return x * y.expand_as(x)
+
+
+CircularEfficientChannelAttn = CecaModule

reference_medclipseg/biomedclip/layers/evo_norm.py

@@ -0,0 +1,352 @@
+""" EvoNorm in PyTorch
+
+Based on `Evolving Normalization-Activation Layers` - https://arxiv.org/abs/2004.02967
+@inproceedings{NEURIPS2020,
+ author = {Liu, Hanxiao and Brock, Andy and Simonyan, Karen and Le, Quoc},
+ booktitle = {Advances in Neural Information Processing Systems},
+ editor = {H. Larochelle and M. Ranzato and R. Hadsell and M. F. Balcan and H. Lin},
+ pages = {13539--13550},
+ publisher = {Curran Associates, Inc.},
+ title = {Evolving Normalization-Activation Layers},
+ url = {https://proceedings.neurips.cc/paper/2020/file/9d4c03631b8b0c85ae08bf05eda37d0f-Paper.pdf},
+ volume = {33},
+ year = {2020}
+}
+
+An attempt at getting decent performing EvoNorms running in PyTorch.
+While faster than other PyTorch impl, still quite a ways off the built-in BatchNorm
+in terms of memory usage and throughput on GPUs.
+
+I'm testing these modules on TPU w/ PyTorch XLA. Promising start but
+currently working around some issues with builtin torch/tensor.var/std. Unlike
+GPU, similar train speeds for EvoNormS variants and BatchNorm.
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from typing import Sequence, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .create_act import create_act_layer
+from .trace_utils import _assert
+
+
+def instance_std(x, eps: float = 1e-5):
+    std = x.float().var(dim=(2, 3), unbiased=False, keepdim=True).add(eps).sqrt().to(x.dtype)
+    return std.expand(x.shape)
+
+
+def instance_std_tpu(x, eps: float = 1e-5):
+    std = manual_var(x, dim=(2, 3)).add(eps).sqrt()
+    return std.expand(x.shape)
+# instance_std = instance_std_tpu
+
+
+def instance_rms(x, eps: float = 1e-5):
+    rms = x.float().square().mean(dim=(2, 3), keepdim=True).add(eps).sqrt().to(x.dtype)
+    return rms.expand(x.shape)
+
+
+def manual_var(x, dim: Union[int, Sequence[int]], diff_sqm: bool = False):
+    xm = x.mean(dim=dim, keepdim=True)
+    if diff_sqm:
+        # difference of squared mean and mean squared, faster on TPU can be less stable
+        var = ((x * x).mean(dim=dim, keepdim=True) - (xm * xm)).clamp(0)
+    else:
+        var = ((x - xm) * (x - xm)).mean(dim=dim, keepdim=True)
+    return var
+
+
+def group_std(x, groups: int = 32, eps: float = 1e-5, flatten: bool = False):
+    B, C, H, W = x.shape
+    x_dtype = x.dtype
+    _assert(C % groups == 0, '')
+    if flatten:
+        x = x.reshape(B, groups, -1)  # FIXME simpler shape causing TPU / XLA issues
+        std = x.float().var(dim=2, unbiased=False, keepdim=True).add(eps).sqrt().to(x_dtype)
+    else:
+        x = x.reshape(B, groups, C // groups, H, W)
+        std = x.float().var(dim=(2, 3, 4), unbiased=False, keepdim=True).add(eps).sqrt().to(x_dtype)
+    return std.expand(x.shape).reshape(B, C, H, W)
+
+
+def group_std_tpu(x, groups: int = 32, eps: float = 1e-5, diff_sqm: bool = False, flatten: bool = False):
+    # This is a workaround for some stability / odd behaviour of .var and .std
+    # running on PyTorch XLA w/ TPUs. These manual var impl are producing much better results
+    B, C, H, W = x.shape
+    _assert(C % groups == 0, '')
+    if flatten:
+        x = x.reshape(B, groups, -1)  # FIXME simpler shape causing TPU / XLA issues
+        var = manual_var(x, dim=-1, diff_sqm=diff_sqm)
+    else:
+        x = x.reshape(B, groups, C // groups, H, W)
+        var = manual_var(x, dim=(2, 3, 4), diff_sqm=diff_sqm)
+    return var.add(eps).sqrt().expand(x.shape).reshape(B, C, H, W)
+#group_std = group_std_tpu  # FIXME TPU temporary
+
+
+def group_rms(x, groups: int = 32, eps: float = 1e-5):
+    B, C, H, W = x.shape
+    _assert(C % groups == 0, '')
+    x_dtype = x.dtype
+    x = x.reshape(B, groups, C // groups, H, W)
+    rms = x.float().square().mean(dim=(2, 3, 4), keepdim=True).add(eps).sqrt_().to(x_dtype)
+    return rms.expand(x.shape).reshape(B, C, H, W)
+
+
+class EvoNorm2dB0(nn.Module):
+    def __init__(self, num_features, apply_act=True, momentum=0.1, eps=1e-3, **_):
+        super().__init__()
+        self.apply_act = apply_act  # apply activation (non-linearity)
+        self.momentum = momentum
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(num_features))
+        self.bias = nn.Parameter(torch.zeros(num_features))
+        self.v = nn.Parameter(torch.ones(num_features)) if apply_act else None
+        self.register_buffer('running_var', torch.ones(num_features))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.ones_(self.weight)
+        nn.init.zeros_(self.bias)
+        if self.v is not None:
+            nn.init.ones_(self.v)
+
+    def forward(self, x):
+        _assert(x.dim() == 4, 'expected 4D input')
+        x_dtype = x.dtype
+        v_shape = (1, -1, 1, 1)
+        if self.v is not None:
+            if self.training:
+                var = x.float().var(dim=(0, 2, 3), unbiased=False)
+                # var = manual_var(x, dim=(0, 2, 3)).squeeze()
+                n = x.numel() / x.shape[1]
+                self.running_var.copy_(
+                    self.running_var * (1 - self.momentum) +
+                    var.detach() * self.momentum * (n / (n - 1)))
+            else:
+                var = self.running_var
+            left = var.add(self.eps).sqrt_().to(x_dtype).view(v_shape).expand_as(x)
+            v = self.v.to(x_dtype).view(v_shape)
+            right = x * v + instance_std(x, self.eps)
+            x = x / left.max(right)
+        return x * self.weight.to(x_dtype).view(v_shape) + self.bias.to(x_dtype).view(v_shape)
+
+
+class EvoNorm2dB1(nn.Module):
+    def __init__(self, num_features, apply_act=True, momentum=0.1, eps=1e-5, **_):
+        super().__init__()
+        self.apply_act = apply_act  # apply activation (non-linearity)
+        self.momentum = momentum
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(num_features))
+        self.bias = nn.Parameter(torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.ones_(self.weight)
+        nn.init.zeros_(self.bias)
+
+    def forward(self, x):
+        _assert(x.dim() == 4, 'expected 4D input')
+        x_dtype = x.dtype
+        v_shape = (1, -1, 1, 1)
+        if self.apply_act:
+            if self.training:
+                var = x.float().var(dim=(0, 2, 3), unbiased=False)
+                n = x.numel() / x.shape[1]
+                self.running_var.copy_(
+                    self.running_var * (1 - self.momentum) +
+                    var.detach().to(self.running_var.dtype) * self.momentum * (n / (n - 1)))
+            else:
+                var = self.running_var
+            var = var.to(x_dtype).view(v_shape)
+            left = var.add(self.eps).sqrt_()
+            right = (x + 1) * instance_rms(x, self.eps)
+            x = x / left.max(right)
+        return x * self.weight.view(v_shape).to(x_dtype) + self.bias.view(v_shape).to(x_dtype)
+
+
+class EvoNorm2dB2(nn.Module):
+    def __init__(self, num_features, apply_act=True, momentum=0.1, eps=1e-5, **_):
+        super().__init__()
+        self.apply_act = apply_act  # apply activation (non-linearity)
+        self.momentum = momentum
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(num_features))
+        self.bias = nn.Parameter(torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.ones_(self.weight)
+        nn.init.zeros_(self.bias)
+
+    def forward(self, x):
+        _assert(x.dim() == 4, 'expected 4D input')
+        x_dtype = x.dtype
+        v_shape = (1, -1, 1, 1)
+        if self.apply_act:
+            if self.training:
+                var = x.float().var(dim=(0, 2, 3), unbiased=False)
+                n = x.numel() / x.shape[1]
+                self.running_var.copy_(
+                    self.running_var * (1 - self.momentum) +
+                    var.detach().to(self.running_var.dtype) * self.momentum * (n / (n - 1)))
+            else:
+                var = self.running_var
+            var = var.to(x_dtype).view(v_shape)
+            left = var.add(self.eps).sqrt_()
+            right = instance_rms(x, self.eps) - x
+            x = x / left.max(right)
+        return x * self.weight.view(v_shape).to(x_dtype) + self.bias.view(v_shape).to(x_dtype)
+
+
+class EvoNorm2dS0(nn.Module):
+    def __init__(self, num_features, groups=32, group_size=None, apply_act=True, eps=1e-5, **_):
+        super().__init__()
+        self.apply_act = apply_act  # apply activation (non-linearity)
+        if group_size:
+            assert num_features % group_size == 0
+            self.groups = num_features // group_size
+        else:
+            self.groups = groups
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(num_features))
+        self.bias = nn.Parameter(torch.zeros(num_features))
+        self.v = nn.Parameter(torch.ones(num_features)) if apply_act else None
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.ones_(self.weight)
+        nn.init.zeros_(self.bias)
+        if self.v is not None:
+            nn.init.ones_(self.v)
+
+    def forward(self, x):
+        _assert(x.dim() == 4, 'expected 4D input')
+        x_dtype = x.dtype
+        v_shape = (1, -1, 1, 1)
+        if self.v is not None:
+            v = self.v.view(v_shape).to(x_dtype)
+            x = x * (x * v).sigmoid() / group_std(x, self.groups, self.eps)
+        return x * self.weight.view(v_shape).to(x_dtype) + self.bias.view(v_shape).to(x_dtype)
+
+
+class EvoNorm2dS0a(EvoNorm2dS0):
+    def __init__(self, num_features, groups=32, group_size=None, apply_act=True, eps=1e-3, **_):
+        super().__init__(
+            num_features, groups=groups, group_size=group_size, apply_act=apply_act, eps=eps)
+
+    def forward(self, x):
+        _assert(x.dim() == 4, 'expected 4D input')
+        x_dtype = x.dtype
+        v_shape = (1, -1, 1, 1)
+        d = group_std(x, self.groups, self.eps)
+        if self.v is not None:
+            v = self.v.view(v_shape).to(x_dtype)
+            x = x * (x * v).sigmoid()
+        x = x / d
+        return x * self.weight.view(v_shape).to(x_dtype) + self.bias.view(v_shape).to(x_dtype)
+
+
+class EvoNorm2dS1(nn.Module):
+    def __init__(
+            self, num_features, groups=32, group_size=None,
+            apply_act=True, act_layer=None, eps=1e-5, **_):
+        super().__init__()
+        act_layer = act_layer or nn.SiLU
+        self.apply_act = apply_act  # apply activation (non-linearity)
+        if act_layer is not None and apply_act:
+            self.act = create_act_layer(act_layer)
+        else:
+            self.act = nn.Identity()
+        if group_size:
+            assert num_features % group_size == 0
+            self.groups = num_features // group_size
+        else:
+            self.groups = groups
+        self.eps = eps
+        self.pre_act_norm = False
+        self.weight = nn.Parameter(torch.ones(num_features))
+        self.bias = nn.Parameter(torch.zeros(num_features))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.ones_(self.weight)
+        nn.init.zeros_(self.bias)
+
+    def forward(self, x):
+        _assert(x.dim() == 4, 'expected 4D input')
+        x_dtype = x.dtype
+        v_shape = (1, -1, 1, 1)
+        if self.apply_act:
+            x = self.act(x) / group_std(x, self.groups, self.eps)
+        return x * self.weight.view(v_shape).to(x_dtype) + self.bias.view(v_shape).to(x_dtype)
+
+
+class EvoNorm2dS1a(EvoNorm2dS1):
+    def __init__(
+            self, num_features, groups=32, group_size=None,
+            apply_act=True, act_layer=None, eps=1e-3, **_):
+        super().__init__(
+            num_features, groups=groups, group_size=group_size, apply_act=apply_act, act_layer=act_layer, eps=eps)
+
+    def forward(self, x):
+        _assert(x.dim() == 4, 'expected 4D input')
+        x_dtype = x.dtype
+        v_shape = (1, -1, 1, 1)
+        x = self.act(x) / group_std(x, self.groups, self.eps)
+        return x * self.weight.view(v_shape).to(x_dtype) + self.bias.view(v_shape).to(x_dtype)
+
+
+class EvoNorm2dS2(nn.Module):
+    def __init__(
+            self, num_features, groups=32, group_size=None,
+            apply_act=True, act_layer=None, eps=1e-5, **_):
+        super().__init__()
+        act_layer = act_layer or nn.SiLU
+        self.apply_act = apply_act  # apply activation (non-linearity)
+        if act_layer is not None and apply_act:
+            self.act = create_act_layer(act_layer)
+        else:
+            self.act = nn.Identity()
+        if group_size:
+            assert num_features % group_size == 0
+            self.groups = num_features // group_size
+        else:
+            self.groups = groups
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(num_features))
+        self.bias = nn.Parameter(torch.zeros(num_features))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.ones_(self.weight)
+        nn.init.zeros_(self.bias)
+
+    def forward(self, x):
+        _assert(x.dim() == 4, 'expected 4D input')
+        x_dtype = x.dtype
+        v_shape = (1, -1, 1, 1)
+        if self.apply_act:
+            x = self.act(x) / group_rms(x, self.groups, self.eps)
+        return x * self.weight.view(v_shape).to(x_dtype) + self.bias.view(v_shape).to(x_dtype)
+
+
+class EvoNorm2dS2a(EvoNorm2dS2):
+    def __init__(
+            self, num_features, groups=32, group_size=None,
+            apply_act=True, act_layer=None, eps=1e-3, **_):
+        super().__init__(
+            num_features, groups=groups, group_size=group_size, apply_act=apply_act, act_layer=act_layer, eps=eps)
+
+    def forward(self, x):
+        _assert(x.dim() == 4, 'expected 4D input')
+        x_dtype = x.dtype
+        v_shape = (1, -1, 1, 1)
+        x = self.act(x) / group_rms(x, self.groups, self.eps)
+        return x * self.weight.view(v_shape).to(x_dtype) + self.bias.view(v_shape).to(x_dtype)

reference_medclipseg/biomedclip/layers/fast_norm.py

@@ -0,0 +1,206 @@
+""" 'Fast' Normalization Functions
+
+For GroupNorm and LayerNorm these functions bypass typical AMP upcast to float32.
+
+Additionally, for LayerNorm, the APEX fused LN is used if available (which also does not upcast)
+
+Hacked together by / Copyright 2022 Ross Wightman
+"""
+from typing import List, Optional
+
+import torch
+from torch.nn import functional as F
+
+try:
+    from apex.normalization.fused_layer_norm import fused_layer_norm_affine
+    has_apex = True
+except ImportError:
+    has_apex = False
+
+try:
+    from apex.normalization.fused_layer_norm import fused_rms_norm_affine, fused_rms_norm
+    has_apex_rmsnorm = True
+except ImportError:
+    has_apex_rmsnorm = False
+
+
+has_torch_rms_norm = hasattr(F, 'rms_norm')
+
+# fast (ie lower precision LN) can be disabled with this flag if issues crop up
+_USE_FAST_NORM = False  # defaulting to False for now
+
+
+def get_autocast_dtype(device: str = 'cuda'):
+    try:
+        return torch.get_autocast_dtype(device)
+    except (AttributeError, TypeError):
+        # dispatch to older device specific fns, only covering cuda/cpu devices here
+        if device == 'cpu':
+            return torch.get_autocast_cpu_dtype()
+        else:
+            assert device == 'cuda'
+            return torch.get_autocast_gpu_dtype()
+
+
+def is_autocast_enabled(device: str = 'cuda'):
+    try:
+        return torch.is_autocast_enabled(device)
+    except TypeError:
+        # dispatch to older device specific fns, only covering cuda/cpu devices here
+        if device == 'cpu':
+            return torch.is_autocast_cpu_enabled()
+        else:
+            assert device == 'cuda'
+            return torch.is_autocast_enabled()  # defaults cuda (only cuda on older pytorch)
+
+
+def is_fast_norm():
+    return _USE_FAST_NORM
+
+
+def set_fast_norm(enable=True):
+    global _USE_FAST_NORM
+    _USE_FAST_NORM = enable
+
+
+def fast_group_norm(
+    x: torch.Tensor,
+    num_groups: int,
+    weight: Optional[torch.Tensor] = None,
+    bias: Optional[torch.Tensor] = None,
+    eps: float = 1e-5
+) -> torch.Tensor:
+    if torch.jit.is_scripting():
+        # currently cannot use is_autocast_enabled within torchscript
+        return F.group_norm(x, num_groups, weight, bias, eps)
+
+    if is_autocast_enabled(x.device.type):
+        # normally native AMP casts GN inputs to float32
+        # here we use the low precision autocast dtype
+        dt = get_autocast_dtype(x.device.type)
+        x, weight, bias = x.to(dt), weight.to(dt), bias.to(dt) if bias is not None else None
+
+    with torch.amp.autocast(device_type=x.device.type, enabled=False):
+        return F.group_norm(x, num_groups, weight, bias, eps)
+
+
+def fast_layer_norm(
+    x: torch.Tensor,
+    normalized_shape: List[int],
+    weight: Optional[torch.Tensor] = None,
+    bias: Optional[torch.Tensor] = None,
+    eps: float = 1e-5
+) -> torch.Tensor:
+    if torch.jit.is_scripting():
+        # currently cannot use is_autocast_enabled within torchscript
+        return F.layer_norm(x, normalized_shape, weight, bias, eps)
+
+    if has_apex:
+        return fused_layer_norm_affine(x, weight, bias, normalized_shape, eps)
+
+    if is_autocast_enabled(x.device.type):
+        # normally native AMP casts LN inputs to float32
+        # apex LN does not, this is behaving like Apex
+        dt = get_autocast_dtype(x.device.type)
+        x, weight, bias = x.to(dt), weight.to(dt), bias.to(dt) if bias is not None else None
+
+    with torch.amp.autocast(device_type=x.device.type, enabled=False):
+        return F.layer_norm(x, normalized_shape, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    normalized_shape: List[int],
+    weight: Optional[torch.Tensor] = None,
+    eps: float = 1e-5,
+):
+    norm_ndim = len(normalized_shape)
+    v = x.pow(2)
+    if torch.jit.is_scripting():
+        # ndim = len(x.shape)
+        # dims = list(range(ndim - norm_ndim, ndim))  # this doesn't work on pytorch <= 1.13.x
+        # NOTE -ve dims cause torchscript to crash in some cases, out of options to work around
+        assert norm_ndim == 1
+        v = torch.mean(v, dim=-1).unsqueeze(-1)  # ts crashes with -ve dim + keepdim=True
+    else:
+        dims = tuple(range(-1, -norm_ndim - 1, -1))
+        v = torch.mean(v, dim=dims, keepdim=True)
+    x = x * torch.rsqrt(v + eps)
+    if weight is not None:
+        x = x * weight
+    return x
+
+
+def fast_rms_norm(
+    x: torch.Tensor,
+    normalized_shape: List[int],
+    weight: Optional[torch.Tensor] = None,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    if torch.jit.is_scripting():
+        # this must be by itself, cannot merge with has_apex_rmsnorm
+        return rms_norm(x, normalized_shape, weight, eps)
+
+    if has_apex_rmsnorm:
+        if weight is None:
+            return fused_rms_norm(x, normalized_shape, eps)
+        else:
+            return fused_rms_norm_affine(x, weight, normalized_shape, eps)
+
+    if is_autocast_enabled(x.device.type):
+        # normally native AMP casts LN inputs to float32
+        # apex LN does not, this is behaving like Apex
+        dt = get_autocast_dtype(x.device.type)
+        x, weight = x.to(dt), weight.to(dt)
+
+    with torch.amp.autocast(device_type=x.device.type, enabled=False):
+        if has_torch_rms_norm:
+            x = F.rms_norm(x, normalized_shape, weight, eps)
+        else:
+            x = rms_norm(x, normalized_shape, weight, eps)
+
+    return x
+
+
+def simple_norm(
+    x: torch.Tensor,
+    normalized_shape: List[int],
+    weight: Optional[torch.Tensor] = None,
+    eps: float = 1e-5,
+):
+    norm_ndim = len(normalized_shape)
+    if torch.jit.is_scripting():
+        # ndim = len(x.shape)
+        # dims = list(range(ndim - norm_ndim, ndim))  # this doesn't work on pytorch <= 1.13.x
+        # NOTE -ve dims cause torchscript to crash in some cases, out of options to work around
+        assert norm_ndim == 1
+        v = torch.var(x, dim=-1).unsqueeze(-1)  # ts crashes with -ve dim + keepdim=True
+    else:
+        dims = tuple(range(-1, -norm_ndim - 1, -1))
+        v = torch.var(x, dim=dims, keepdim=True)
+    x = x * torch.rsqrt(v + eps)
+    if weight is not None:
+        x = x * weight
+    return x
+
+
+def fast_simple_norm(
+    x: torch.Tensor,
+    normalized_shape: List[int],
+    weight: Optional[torch.Tensor] = None,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    if torch.jit.is_scripting():
+        # this must be by itself, cannot merge with has_apex_rmsnorm
+        return simple_norm(x, normalized_shape, weight, eps)
+
+    if is_autocast_enabled(x.device.type):
+        # normally native AMP casts LN inputs to float32
+        # apex LN does not, this is behaving like Apex
+        dt = get_autocast_dtype(x.device.type)
+        x, weight = x.to(dt), weight.to(dt)
+
+    with torch.amp.autocast(device_type=x.device.type, enabled=False):
+        x = simple_norm(x, normalized_shape, weight, eps)
+    return x
+

reference_medclipseg/biomedclip/layers/filter_response_norm.py

@@ -0,0 +1,68 @@
+""" Filter Response Norm in PyTorch
+
+Based on `Filter Response Normalization Layer` - https://arxiv.org/abs/1911.09737
+
+Hacked together by / Copyright 2021 Ross Wightman
+"""
+import torch
+import torch.nn as nn
+
+from .create_act import create_act_layer
+from .trace_utils import _assert
+
+
+def inv_instance_rms(x, eps: float = 1e-5):
+    rms = x.square().float().mean(dim=(2, 3), keepdim=True).add(eps).rsqrt().to(x.dtype)
+    return rms.expand(x.shape)
+
+
+class FilterResponseNormTlu2d(nn.Module):
+    def __init__(self, num_features, apply_act=True, eps=1e-5, rms=True, **_):
+        super(FilterResponseNormTlu2d, self).__init__()
+        self.apply_act = apply_act  # apply activation (non-linearity)
+        self.rms = rms
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(num_features))
+        self.bias = nn.Parameter(torch.zeros(num_features))
+        self.tau = nn.Parameter(torch.zeros(num_features)) if apply_act else None
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.ones_(self.weight)
+        nn.init.zeros_(self.bias)
+        if self.tau is not None:
+            nn.init.zeros_(self.tau)
+
+    def forward(self, x):
+        _assert(x.dim() == 4, 'expected 4D input')
+        x_dtype = x.dtype
+        v_shape = (1, -1, 1, 1)
+        x = x * inv_instance_rms(x, self.eps)
+        x = x * self.weight.view(v_shape).to(dtype=x_dtype) + self.bias.view(v_shape).to(dtype=x_dtype)
+        return torch.maximum(x, self.tau.reshape(v_shape).to(dtype=x_dtype)) if self.tau is not None else x
+
+
+class FilterResponseNormAct2d(nn.Module):
+    def __init__(self, num_features, apply_act=True, act_layer=nn.ReLU, inplace=None, rms=True, eps=1e-5, **_):
+        super(FilterResponseNormAct2d, self).__init__()
+        if act_layer is not None and apply_act:
+            self.act = create_act_layer(act_layer, inplace=inplace)
+        else:
+            self.act = nn.Identity()
+        self.rms = rms
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(num_features))
+        self.bias = nn.Parameter(torch.zeros(num_features))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.ones_(self.weight)
+        nn.init.zeros_(self.bias)
+
+    def forward(self, x):
+        _assert(x.dim() == 4, 'expected 4D input')
+        x_dtype = x.dtype
+        v_shape = (1, -1, 1, 1)
+        x = x * inv_instance_rms(x, self.eps)
+        x = x * self.weight.view(v_shape).to(dtype=x_dtype) + self.bias.view(v_shape).to(dtype=x_dtype)
+        return self.act(x)

reference_medclipseg/biomedclip/layers/format.py

@@ -0,0 +1,58 @@
+from enum import Enum
+from typing import Union
+
+import torch
+
+
+class Format(str, Enum):
+    NCHW = 'NCHW'
+    NHWC = 'NHWC'
+    NCL = 'NCL'
+    NLC = 'NLC'
+
+
+FormatT = Union[str, Format]
+
+
+def get_spatial_dim(fmt: FormatT):
+    fmt = Format(fmt)
+    if fmt is Format.NLC:
+        dim = (1,)
+    elif fmt is Format.NCL:
+        dim = (2,)
+    elif fmt is Format.NHWC:
+        dim = (1, 2)
+    else:
+        dim = (2, 3)
+    return dim
+
+
+def get_channel_dim(fmt: FormatT):
+    fmt = Format(fmt)
+    if fmt is Format.NHWC:
+        dim = 3
+    elif fmt is Format.NLC:
+        dim = 2
+    else:
+        dim = 1
+    return dim
+
+
+def nchw_to(x: torch.Tensor, fmt: Format):
+    if fmt == Format.NHWC:
+        x = x.permute(0, 2, 3, 1)
+    elif fmt == Format.NLC:
+        x = x.flatten(2).transpose(1, 2)
+    elif fmt == Format.NCL:
+        x = x.flatten(2)
+    return x
+
+
+def nhwc_to(x: torch.Tensor, fmt: Format):
+    if fmt == Format.NCHW:
+        x = x.permute(0, 3, 1, 2)
+    elif fmt == Format.NLC:
+        x = x.flatten(1, 2)
+    elif fmt == Format.NCL:
+        x = x.flatten(1, 2).transpose(1, 2)
+    return x

reference_medclipseg/biomedclip/layers/gather_excite.py

@@ -0,0 +1,90 @@
+""" Gather-Excite Attention Block
+
+Paper: `Gather-Excite: Exploiting Feature Context in CNNs` - https://arxiv.org/abs/1810.12348
+
+Official code here, but it's only partial impl in Caffe: https://github.com/hujie-frank/GENet
+
+I've tried to support all of the extent both w/ and w/o params. I don't believe I've seen another
+impl that covers all of the cases.
+
+NOTE: extent=0 + extra_params=False is equivalent to Squeeze-and-Excitation
+
+Hacked together by / Copyright 2021 Ross Wightman
+"""
+import math
+
+from torch import nn as nn
+import torch.nn.functional as F
+
+from .create_act import create_act_layer, get_act_layer
+from .create_conv2d import create_conv2d
+from .helpers import make_divisible
+from .mlp import ConvMlp
+
+
+class GatherExcite(nn.Module):
+    """ Gather-Excite Attention Module
+    """
+    def __init__(
+            self, channels, feat_size=None, extra_params=False, extent=0, use_mlp=True,
+            rd_ratio=1./16, rd_channels=None,  rd_divisor=1, add_maxpool=False,
+            act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d, gate_layer='sigmoid'):
+        super(GatherExcite, self).__init__()
+        self.add_maxpool = add_maxpool
+        act_layer = get_act_layer(act_layer)
+        self.extent = extent
+        if extra_params:
+            self.gather = nn.Sequential()
+            if extent == 0:
+                assert feat_size is not None, 'spatial feature size must be specified for global extent w/ params'
+                self.gather.add_module(
+                    'conv1', create_conv2d(channels, channels, kernel_size=feat_size, stride=1, depthwise=True))
+                if norm_layer:
+                    self.gather.add_module(f'norm1', nn.BatchNorm2d(channels))
+            else:
+                assert extent % 2 == 0
+                num_conv = int(math.log2(extent))
+                for i in range(num_conv):
+                    self.gather.add_module(
+                        f'conv{i + 1}',
+                        create_conv2d(channels, channels, kernel_size=3, stride=2, depthwise=True))
+                    if norm_layer:
+                        self.gather.add_module(f'norm{i + 1}', nn.BatchNorm2d(channels))
+                    if i != num_conv - 1:
+                        self.gather.add_module(f'act{i + 1}', act_layer(inplace=True))
+        else:
+            self.gather = None
+            if self.extent == 0:
+                self.gk = 0
+                self.gs = 0
+            else:
+                assert extent % 2 == 0
+                self.gk = self.extent * 2 - 1
+                self.gs = self.extent
+
+        if not rd_channels:
+            rd_channels = make_divisible(channels * rd_ratio, rd_divisor, round_limit=0.)
+        self.mlp = ConvMlp(channels, rd_channels, act_layer=act_layer) if use_mlp else nn.Identity()
+        self.gate = create_act_layer(gate_layer)
+
+    def forward(self, x):
+        size = x.shape[-2:]
+        if self.gather is not None:
+            x_ge = self.gather(x)
+        else:
+            if self.extent == 0:
+                # global extent
+                x_ge = x.mean(dim=(2, 3), keepdims=True)
+                if self.add_maxpool:
+                    # experimental codepath, may remove or change
+                    x_ge = 0.5 * x_ge + 0.5 * x.amax((2, 3), keepdim=True)
+            else:
+                x_ge = F.avg_pool2d(
+                    x, kernel_size=self.gk, stride=self.gs, padding=self.gk // 2, count_include_pad=False)
+                if self.add_maxpool:
+                    # experimental codepath, may remove or change
+                    x_ge = 0.5 * x_ge + 0.5 * F.max_pool2d(x, kernel_size=self.gk, stride=self.gs, padding=self.gk // 2)
+        x_ge = self.mlp(x_ge)
+        if x_ge.shape[-1] != 1 or x_ge.shape[-2] != 1:
+            x_ge = F.interpolate(x_ge, size=size)
+        return x * self.gate(x_ge)

reference_medclipseg/biomedclip/layers/global_context.py

@@ -0,0 +1,67 @@
+""" Global Context Attention Block
+
+Paper: `GCNet: Non-local Networks Meet Squeeze-Excitation Networks and Beyond`
+    - https://arxiv.org/abs/1904.11492
+
+Official code consulted as reference: https://github.com/xvjiarui/GCNet
+
+Hacked together by / Copyright 2021 Ross Wightman
+"""
+from torch import nn as nn
+import torch.nn.functional as F
+
+from .create_act import create_act_layer, get_act_layer
+from .helpers import make_divisible
+from .mlp import ConvMlp
+from .norm import LayerNorm2d
+
+
+class GlobalContext(nn.Module):
+
+    def __init__(self, channels, use_attn=True, fuse_add=False, fuse_scale=True, init_last_zero=False,
+                 rd_ratio=1./8, rd_channels=None, rd_divisor=1, act_layer=nn.ReLU, gate_layer='sigmoid'):
+        super(GlobalContext, self).__init__()
+        act_layer = get_act_layer(act_layer)
+
+        self.conv_attn = nn.Conv2d(channels, 1, kernel_size=1, bias=True) if use_attn else None
+
+        if rd_channels is None:
+            rd_channels = make_divisible(channels * rd_ratio, rd_divisor, round_limit=0.)
+        if fuse_add:
+            self.mlp_add = ConvMlp(channels, rd_channels, act_layer=act_layer, norm_layer=LayerNorm2d)
+        else:
+            self.mlp_add = None
+        if fuse_scale:
+            self.mlp_scale = ConvMlp(channels, rd_channels, act_layer=act_layer, norm_layer=LayerNorm2d)
+        else:
+            self.mlp_scale = None
+
+        self.gate = create_act_layer(gate_layer)
+        self.init_last_zero = init_last_zero
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        if self.conv_attn is not None:
+            nn.init.kaiming_normal_(self.conv_attn.weight, mode='fan_in', nonlinearity='relu')
+        if self.mlp_add is not None:
+            nn.init.zeros_(self.mlp_add.fc2.weight)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+
+        if self.conv_attn is not None:
+            attn = self.conv_attn(x).reshape(B, 1, H * W)  # (B, 1, H * W)
+            attn = F.softmax(attn, dim=-1).unsqueeze(3)  # (B, 1, H * W, 1)
+            context = x.reshape(B, C, H * W).unsqueeze(1) @ attn
+            context = context.view(B, C, 1, 1)
+        else:
+            context = x.mean(dim=(2, 3), keepdim=True)
+
+        if self.mlp_scale is not None:
+            mlp_x = self.mlp_scale(context)
+            x = x * self.gate(mlp_x)
+        if self.mlp_add is not None:
+            mlp_x = self.mlp_add(context)
+            x = x + mlp_x
+
+        return x

reference_medclipseg/biomedclip/layers/grid.py

@@ -0,0 +1,49 @@
+from typing import Tuple
+
+import torch
+
+
+def ndgrid(*tensors) -> Tuple[torch.Tensor, ...]:
+    """generate N-D grid in dimension order.
+
+    The ndgrid function is like meshgrid except that the order of the first two input arguments are switched.
+
+    That is, the statement
+    [X1,X2,X3] = ndgrid(x1,x2,x3)
+
+    produces the same result as
+
+    [X2,X1,X3] = meshgrid(x2,x1,x3)
+
+    This naming is based on MATLAB, the purpose is to avoid confusion due to torch's change to make
+    torch.meshgrid behaviour move from matching ndgrid ('ij') indexing to numpy meshgrid defaults of ('xy').
+
+    """
+    try:
+        return torch.meshgrid(*tensors, indexing='ij')
+    except TypeError:
+        # old PyTorch < 1.10 will follow this path as it does not have indexing arg,
+        # the old behaviour of meshgrid was 'ij'
+        return torch.meshgrid(*tensors)
+
+
+def meshgrid(*tensors) -> Tuple[torch.Tensor, ...]:
+    """generate N-D grid in spatial dim order.
+
+    The meshgrid function is similar to ndgrid except that the order of the
+    first two input and output arguments is switched.
+
+    That is, the statement
+
+    [X,Y,Z] = meshgrid(x,y,z)
+    produces the same result as
+
+    [Y,X,Z] = ndgrid(y,x,z)
+    Because of this, meshgrid is better suited to problems in two- or three-dimensional Cartesian space,
+    while ndgrid is better suited to multidimensional problems that aren't spatially based.
+    """
+
+    # NOTE: this will throw in PyTorch < 1.10 as meshgrid did not support indexing arg or have
+    # capability of generating grid in xy order before then.
+    return torch.meshgrid(*tensors, indexing='xy')
+

reference_medclipseg/biomedclip/layers/grn.py

@@ -0,0 +1,39 @@
+""" Global Response Normalization Module
+
+Based on the GRN layer presented in
+`ConvNeXt-V2 - Co-designing and Scaling ConvNets with Masked Autoencoders` - https://arxiv.org/abs/2301.00808
+
+This implementation
+* works for both NCHW and NHWC tensor layouts
+* uses affine param names matching existing torch norm layers
+* slightly improves eager mode performance via fused addcmul
+
+Hacked together by / Copyright 2023 Ross Wightman
+"""
+
+import torch
+from torch import nn as nn
+
+
+class GlobalResponseNorm(nn.Module):
+    """ Global Response Normalization layer
+    """
+    def __init__(self, dim, eps=1e-6, channels_last=True):
+        super().__init__()
+        self.eps = eps
+        if channels_last:
+            self.spatial_dim = (1, 2)
+            self.channel_dim = -1
+            self.wb_shape = (1, 1, 1, -1)
+        else:
+            self.spatial_dim = (2, 3)
+            self.channel_dim = 1
+            self.wb_shape = (1, -1, 1, 1)
+
+        self.weight = nn.Parameter(torch.zeros(dim))
+        self.bias = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x):
+        x_g = x.norm(p=2, dim=self.spatial_dim, keepdim=True)
+        x_n = x_g / (x_g.mean(dim=self.channel_dim, keepdim=True) + self.eps)
+        return x + torch.addcmul(self.bias.view(self.wb_shape), self.weight.view(self.wb_shape), x * x_n)

reference_medclipseg/biomedclip/layers/halo_attn.py

@@ -0,0 +1,233 @@
+""" Halo Self Attention
+
+Paper: `Scaling Local Self-Attention for Parameter Efficient Visual Backbones`
+    - https://arxiv.org/abs/2103.12731
+
+@misc{2103.12731,
+Author = {Ashish Vaswani and Prajit Ramachandran and Aravind Srinivas and Niki Parmar and Blake Hechtman and
+    Jonathon Shlens},
+Title = {Scaling Local Self-Attention for Parameter Efficient Visual Backbones},
+Year = {2021},
+}
+
+Status:
+This impl is a WIP, there is no official ref impl and some details in paper weren't clear to me.
+The attention mechanism works but it's slow as implemented.
+
+Hacked together by / Copyright 2021 Ross Wightman
+"""
+from typing import List
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .helpers import make_divisible
+from .weight_init import trunc_normal_
+from .trace_utils import _assert
+
+
+def rel_logits_1d(q, rel_k, permute_mask: List[int]):
+    """ Compute relative logits along one dimension
+
+    As per: https://gist.github.com/aravindsrinivas/56359b79f0ce4449bcb04ab4b56a57a2
+    Originally from: `Attention Augmented Convolutional Networks` - https://arxiv.org/abs/1904.09925
+
+    Args:
+        q: (batch, height, width, dim)
+        rel_k: (2 * window - 1, dim)
+        permute_mask: permute output dim according to this
+    """
+    B, H, W, dim = q.shape
+    rel_size = rel_k.shape[0]
+    win_size = (rel_size + 1) // 2
+
+    x = (q @ rel_k.transpose(-1, -2))
+    x = x.reshape(-1, W, rel_size)
+
+    # pad to shift from relative to absolute indexing
+    x_pad = F.pad(x, [0, 1]).flatten(1)
+    x_pad = F.pad(x_pad, [0, rel_size - W])
+
+    # reshape and slice out the padded elements
+    x_pad = x_pad.reshape(-1, W + 1, rel_size)
+    x = x_pad[:, :W, win_size - 1:]
+
+    # reshape and tile
+    x = x.reshape(B, H, 1, W, win_size).expand(-1, -1, win_size, -1, -1)
+    return x.permute(permute_mask)
+
+
+class PosEmbedRel(nn.Module):
+    """ Relative Position Embedding
+    As per: https://gist.github.com/aravindsrinivas/56359b79f0ce4449bcb04ab4b56a57a2
+    Originally from: `Attention Augmented Convolutional Networks` - https://arxiv.org/abs/1904.09925
+
+    """
+    def __init__(self, block_size, win_size, dim_head, scale):
+        """
+        Args:
+            block_size (int): block size
+            win_size (int): neighbourhood window size
+            dim_head (int): attention head dim
+            scale (float): scale factor (for init)
+        """
+        super().__init__()
+        self.block_size = block_size
+        self.dim_head = dim_head
+        self.height_rel = nn.Parameter(torch.randn(win_size * 2 - 1, dim_head) * scale)
+        self.width_rel = nn.Parameter(torch.randn(win_size * 2 - 1, dim_head) * scale)
+
+    def forward(self, q):
+        B, BB, HW, _ = q.shape
+
+        # relative logits in width dimension.
+        q = q.reshape(-1, self.block_size, self.block_size, self.dim_head)
+        rel_logits_w = rel_logits_1d(q, self.width_rel, permute_mask=(0, 1, 3, 2, 4))
+
+        # relative logits in height dimension.
+        q = q.transpose(1, 2)
+        rel_logits_h = rel_logits_1d(q, self.height_rel, permute_mask=(0, 3, 1, 4, 2))
+
+        rel_logits = rel_logits_h + rel_logits_w
+        rel_logits = rel_logits.reshape(B, BB, HW, -1)
+        return rel_logits
+
+
+class HaloAttn(nn.Module):
+    """ Halo Attention
+
+    Paper: `Scaling Local Self-Attention for Parameter Efficient Visual Backbones`
+        - https://arxiv.org/abs/2103.12731
+
+    The internal dimensions of the attention module are controlled by the interaction of several arguments.
+      * the output dimension of the module is specified by dim_out, which falls back to input dim if not set
+      * the value (v) dimension is set to dim_out // num_heads, the v projection determines the output dim
+      * the query and key (qk) dimensions are determined by
+        * num_heads * dim_head if dim_head is not None
+        * num_heads * (dim_out * attn_ratio // num_heads) if dim_head is None
+      * as seen above, attn_ratio determines the ratio of q and k relative to the output if dim_head not used
+
+    Args:
+        dim (int): input dimension to the module
+        dim_out (int): output dimension of the module, same as dim if not set
+        feat_size (Tuple[int, int]): size of input feature_map (not used, for arg compat with bottle/lambda)
+        stride: output stride of the module, query downscaled if > 1 (default: 1).
+        num_heads: parallel attention heads (default: 8).
+        dim_head: dimension of query and key heads, calculated from dim_out * attn_ratio // num_heads if not set
+        block_size (int): size of blocks. (default: 8)
+        halo_size (int): size of halo overlap. (default: 3)
+        qk_ratio (float): ratio of q and k dimensions to output dimension when dim_head not set. (default: 1.0)
+        qkv_bias (bool) : add bias to q, k, and v projections
+        avg_down (bool): use average pool downsample instead of strided query blocks
+        scale_pos_embed (bool): scale the position embedding as well as Q @ K
+    """
+    def __init__(
+            self, dim, dim_out=None, feat_size=None, stride=1, num_heads=8, dim_head=None, block_size=8, halo_size=3,
+            qk_ratio=1.0, qkv_bias=False, avg_down=False, scale_pos_embed=False):
+        super().__init__()
+        dim_out = dim_out or dim
+        assert dim_out % num_heads == 0
+        assert stride in (1, 2)
+        self.num_heads = num_heads
+        self.dim_head_qk = dim_head or make_divisible(dim_out * qk_ratio, divisor=8) // num_heads
+        self.dim_head_v = dim_out // self.num_heads
+        self.dim_out_qk = num_heads * self.dim_head_qk
+        self.dim_out_v = num_heads * self.dim_head_v
+        self.scale = self.dim_head_qk ** -0.5
+        self.scale_pos_embed = scale_pos_embed
+        self.block_size = self.block_size_ds = block_size
+        self.halo_size = halo_size
+        self.win_size = block_size + halo_size * 2  # neighbourhood window size
+        self.block_stride = 1
+        use_avg_pool = False
+        if stride > 1:
+            use_avg_pool = avg_down or block_size % stride != 0
+            self.block_stride = 1 if use_avg_pool else stride
+            self.block_size_ds = self.block_size // self.block_stride
+
+        # FIXME not clear if this stride behaviour is what the paper intended
+        # Also, the paper mentions using a 3D conv for dealing with the blocking/gather, and leaving
+        # data in unfolded block form. I haven't wrapped my head around how that'd look.
+        self.q = nn.Conv2d(dim, self.dim_out_qk, 1, stride=self.block_stride, bias=qkv_bias)
+        self.kv = nn.Conv2d(dim, self.dim_out_qk + self.dim_out_v, 1, bias=qkv_bias)
+
+        self.pos_embed = PosEmbedRel(
+            block_size=self.block_size_ds, win_size=self.win_size, dim_head=self.dim_head_qk, scale=self.scale)
+
+        self.pool = nn.AvgPool2d(2, 2) if use_avg_pool else nn.Identity()
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        std = self.q.weight.shape[1] ** -0.5  # fan-in
+        trunc_normal_(self.q.weight, std=std)
+        trunc_normal_(self.kv.weight, std=std)
+        trunc_normal_(self.pos_embed.height_rel, std=self.scale)
+        trunc_normal_(self.pos_embed.width_rel, std=self.scale)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        _assert(H % self.block_size == 0, '')
+        _assert(W % self.block_size == 0, '')
+        num_h_blocks = H // self.block_size
+        num_w_blocks = W // self.block_size
+        num_blocks = num_h_blocks * num_w_blocks
+
+        q = self.q(x)
+        # unfold
+        q = q.reshape(
+            -1, self.dim_head_qk,
+            num_h_blocks, self.block_size_ds, num_w_blocks, self.block_size_ds).permute(0, 1, 3, 5, 2, 4)
+        # B, num_heads * dim_head * block_size ** 2, num_blocks
+        q = q.reshape(B * self.num_heads, self.dim_head_qk, -1, num_blocks).transpose(1, 3)
+        # B * num_heads, num_blocks, block_size ** 2, dim_head
+
+        kv = self.kv(x)
+        # Generate overlapping windows for kv. This approach is good for GPU and CPU. However, unfold() is not
+        # lowered for PyTorch XLA so it will be very slow. See code at bottom of file for XLA friendly approach.
+        # FIXME figure out how to switch impl between this and conv2d if XLA being used.
+        kv = F.pad(kv, [self.halo_size, self.halo_size, self.halo_size, self.halo_size])
+        kv = kv.unfold(2, self.win_size, self.block_size).unfold(3, self.win_size, self.block_size).reshape(
+            B * self.num_heads, self.dim_head_qk + self.dim_head_v, num_blocks, -1).permute(0, 2, 3, 1)
+        k, v = torch.split(kv, [self.dim_head_qk, self.dim_head_v], dim=-1)
+        # B * num_heads, num_blocks, win_size ** 2, dim_head_qk or dim_head_v
+
+        if self.scale_pos_embed:
+            attn = (q @ k.transpose(-1, -2) + self.pos_embed(q)) * self.scale
+        else:
+            attn = (q @ k.transpose(-1, -2)) * self.scale + self.pos_embed(q)
+        # B * num_heads, num_blocks, block_size ** 2, win_size ** 2
+        attn = attn.softmax(dim=-1)
+
+        out = (attn @ v).transpose(1, 3)  # B * num_heads, dim_head_v, block_size ** 2, num_blocks
+        # fold
+        out = out.reshape(-1, self.block_size_ds, self.block_size_ds, num_h_blocks, num_w_blocks)
+        out = out.permute(0, 3, 1, 4, 2).contiguous().view(
+            B, self.dim_out_v, H // self.block_stride, W // self.block_stride)
+        # B, dim_out, H // block_stride, W // block_stride
+        out = self.pool(out)
+        return out
+
+
+""" Three alternatives for overlapping windows.
+
+`.unfold().unfold()` is same speed as stride tricks with similar clarity as F.unfold()
+
+    if is_xla:
+        # This code achieves haloing on PyTorch XLA with reasonable runtime trade-off, it is
+        # EXTREMELY slow for backward on a GPU though so I need a way of selecting based on environment.
+        WW = self.win_size ** 2
+        pw = torch.eye(WW, dtype=x.dtype, device=x.device).reshape(WW, 1, self.win_size, self.win_size)
+        kv = F.conv2d(kv.reshape(-1, 1, H, W), pw, stride=self.block_size, padding=self.halo_size)
+    elif self.stride_tricks:
+        kv = F.pad(kv, [self.halo_size, self.halo_size, self.halo_size, self.halo_size]).contiguous()
+        kv = kv.as_strided((
+            B, self.dim_out_qk + self.dim_out_v, self.win_size, self.win_size, num_h_blocks, num_w_blocks),
+            stride=(kv.stride(0), kv.stride(1), kv.shape[-1], 1, self.block_size * kv.shape[-1], self.block_size))
+    else:
+        kv = F.unfold(kv, kernel_size=self.win_size, stride=self.block_size, padding=self.halo_size)
+
+    kv = kv.reshape(
+       B * self.num_heads, self.dim_head_qk + self.dim_head_v, -1, num_blocks).transpose(1, 3)
+"""

reference_medclipseg/biomedclip/layers/helpers.py

@@ -0,0 +1,43 @@
+""" Layer/Module Helpers
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from itertools import repeat
+import collections.abc
+
+
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+            return tuple(x)
+        return tuple(repeat(x, n))
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple
+
+
+def make_divisible(v, divisor=8, min_value=None, round_limit=.9):
+    min_value = min_value or divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < round_limit * v:
+        new_v += divisor
+    return new_v
+
+
+def extend_tuple(x, n):
+    # pads a tuple to specified n by padding with last value
+    if not isinstance(x, (tuple, list)):
+        x = (x,)
+    else:
+        x = tuple(x)
+    pad_n = n - len(x)
+    if pad_n <= 0:
+        return x[:n]
+    return x + (x[-1],) * pad_n

reference_medclipseg/biomedclip/layers/hybrid_embed.py

@@ -0,0 +1,253 @@
+""" Image to Patch Hybird Embedding Layer
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import logging
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+from torch import nn as nn
+import torch.nn.functional as F
+
+from .format import Format, nchw_to
+from .helpers import to_2tuple
+from .patch_embed import resample_patch_embed
+
+
+_logger = logging.getLogger(__name__)
+
+
+class HybridEmbed(nn.Module):
+    """ CNN Feature Map Embedding
+    Extract feature map from CNN, flatten, project to embedding dim.
+    """
+    output_fmt: Format
+    dynamic_img_pad: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            backbone: nn.Module,
+            img_size: Union[int, Tuple[int, int]] = 224,
+            patch_size: Union[int, Tuple[int, int]] = 1,
+            feature_size: Optional[Union[int, Tuple[int, int]]] = None,
+            feature_ratio: Optional[Union[int, Tuple[int, int]]] = None,
+            in_chans: int = 3,
+            embed_dim: int = 768,
+            bias: bool = True,
+            proj: bool = True,
+            flatten: bool = True,
+            output_fmt: Optional[str] = None,
+            strict_img_size: bool = True,
+            dynamic_img_pad: bool = False,
+    ):
+        super().__init__()
+        assert isinstance(backbone, nn.Module)
+        self.backbone = backbone
+        self.in_chans = in_chans
+        (
+            self.img_size,
+            self.patch_size,
+            self.feature_size,
+            self.feature_ratio,
+            self.feature_dim,
+            self.grid_size,
+            self.num_patches,
+        ) = self._init_backbone(
+            img_size=img_size,
+            patch_size=patch_size,
+            feature_size=feature_size,
+            feature_ratio=feature_ratio,
+        )
+
+        if output_fmt is not None:
+            self.flatten = False
+            self.output_fmt = Format(output_fmt)
+        else:
+            # flatten spatial dim and transpose to channels last, kept for bwd compat
+            self.flatten = flatten
+            self.output_fmt = Format.NCHW
+        self.strict_img_size = strict_img_size
+        self.dynamic_img_pad = dynamic_img_pad
+        if not dynamic_img_pad:
+            assert self.feature_size[0] % self.patch_size[0] == 0 and self.feature_size[1] % self.patch_size[1] == 0
+
+        if proj:
+            self.proj = nn.Conv2d(
+                self.feature_dim,
+                embed_dim,
+                kernel_size=patch_size,
+                stride=patch_size,
+                bias=bias,
+            )
+        else:
+            assert self.feature_dim == embed_dim, \
+                f'The feature dim ({self.feature_dim} must match embed dim ({embed_dim}) when projection disabled.'
+            self.proj = nn.Identity()
+
+    def _init_backbone(
+            self,
+            img_size: Union[int, Tuple[int, int]] = 224,
+            patch_size: Union[int, Tuple[int, int]] = 1,
+            feature_size: Optional[Union[int, Tuple[int, int]]] = None,
+            feature_ratio: Optional[Union[int, Tuple[int, int]]] = None,
+            feature_dim: Optional[int] = None,
+    ):
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        if feature_size is None:
+            with torch.no_grad():
+                # NOTE Most reliable way of determining output dims is to run forward pass
+                training = self.backbone.training
+                if training:
+                    self.backbone.eval()
+                o = self.backbone(torch.zeros(1, self.in_chans, img_size[0], img_size[1]))
+                if isinstance(o, (list, tuple)):
+                    o = o[-1]  # last feature if backbone outputs list/tuple of features
+                feature_size = o.shape[-2:]
+                feature_dim = o.shape[1]
+                self.backbone.train(training)
+            feature_ratio = tuple([s // f for s, f in zip(img_size, feature_size)])
+        else:
+            feature_size = to_2tuple(feature_size)
+            feature_ratio = to_2tuple(feature_ratio or 16)
+            if feature_dim is None:
+                if hasattr(self.backbone, 'feature_info'):
+                    feature_dim = self.backbone.feature_info.channels()[-1]
+                else:
+                    feature_dim = self.backbone.num_features
+        grid_size = tuple([f // p for f, p in zip(feature_size, patch_size)])
+        num_patches = grid_size[0] * grid_size[1]
+        return img_size, patch_size, feature_size, feature_ratio, feature_dim, grid_size, num_patches
+
+    def set_input_size(
+            self,
+            img_size: Optional[Union[int, Tuple[int, int]]] = None,
+            patch_size: Optional[Union[int, Tuple[int, int]]] = None,
+            feature_size: Optional[Union[int, Tuple[int, int]]] = None,
+            feature_ratio: Optional[Union[int, Tuple[int, int]]] = None,
+            feature_dim: Optional[int] = None,
+    ):
+        assert img_size is not None or patch_size is not None
+        img_size = img_size or self.img_size
+        new_patch_size = None
+        if patch_size is not None:
+            new_patch_size = to_2tuple(patch_size)
+        if new_patch_size is not None and new_patch_size != self.patch_size:
+            assert isinstance(self.proj, nn.Conv2d), 'HybridEmbed must have a projection layer to change patch size.'
+            with torch.no_grad():
+                new_proj = nn.Conv2d(
+                    self.proj.in_channels,
+                    self.proj.out_channels,
+                    kernel_size=new_patch_size,
+                    stride=new_patch_size,
+                    bias=self.proj.bias is not None,
+                )
+                new_proj.weight.copy_(resample_patch_embed(self.proj.weight, new_patch_size, verbose=True))
+                if self.proj.bias is not None:
+                    new_proj.bias.copy_(self.proj.bias)
+                self.proj = new_proj
+            patch_size = new_patch_size
+        patch_size = patch_size or self.patch_size
+
+        if img_size != self.img_size or patch_size != self.patch_size:
+            (
+                self.img_size,
+                self.patch_size,
+                self.feature_size,
+                self.feature_ratio,
+                self.feature_dim,
+                self.grid_size,
+                self.num_patches,
+            ) = self._init_backbone(
+                img_size=img_size,
+                patch_size=patch_size,
+                feature_size=feature_size,
+                feature_ratio=feature_ratio,
+                feature_dim=feature_dim,
+            )
+
+    def feat_ratio(self, as_scalar=True) -> Union[Tuple[int, int], int]:
+        total_reduction = (
+            self.feature_ratio[0] * self.patch_size[0],
+            self.feature_ratio[1] * self.patch_size[1]
+        )
+        if as_scalar:
+            return max(total_reduction)
+        else:
+            return total_reduction
+
+    def dynamic_feat_size(self, img_size: Tuple[int, int]) -> Tuple[int, int]:
+        """ Get feature grid size taking account dynamic padding and backbone network feat reduction
+        """
+        feat_size = (img_size[0] // self.feature_ratio[0], img_size[1] // self.feature_ratio[1])
+        if self.dynamic_img_pad:
+            return math.ceil(feat_size[0] / self.patch_size[0]), math.ceil(feat_size[1] / self.patch_size[1])
+        else:
+            return feat_size[0] // self.patch_size[0], feat_size[1] // self.patch_size[1]
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable: bool = True):
+        if hasattr(self.backbone, 'set_grad_checkpointing'):
+            self.backbone.set_grad_checkpointing(enable=enable)
+        elif hasattr(self.backbone, 'grad_checkpointing'):
+            self.backbone.grad_checkpointing = enable
+
+    def forward(self, x):
+        x = self.backbone(x)
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+        _, _, H, W = x.shape
+        if self.dynamic_img_pad:
+            pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0]
+            pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1]
+            x = F.pad(x, (0, pad_w, 0, pad_h))
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # NCHW -> NLC
+        elif self.output_fmt != Format.NCHW:
+            x = nchw_to(x, self.output_fmt)
+        return x
+
+
+class HybridEmbedWithSize(HybridEmbed):
+    """ CNN Feature Map Embedding
+    Extract feature map from CNN, flatten, project to embedding dim.
+    """
+    def __init__(
+            self,
+            backbone: nn.Module,
+            img_size: Union[int, Tuple[int, int]] = 224,
+            patch_size: Union[int, Tuple[int, int]] = 1,
+            feature_size: Optional[Union[int, Tuple[int, int]]] = None,
+            feature_ratio: Optional[Union[int, Tuple[int, int]]] = None,
+            in_chans: int = 3,
+            embed_dim: int = 768,
+            bias=True,
+            proj=True,
+    ):
+        super().__init__(
+            backbone=backbone,
+            img_size=img_size,
+            patch_size=patch_size,
+            feature_size=feature_size,
+            feature_ratio=feature_ratio,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            bias=bias,
+            proj=proj,
+        )
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable: bool = True):
+        if hasattr(self.backbone, 'set_grad_checkpointing'):
+            self.backbone.set_grad_checkpointing(enable=enable)
+        elif hasattr(self.backbone, 'grad_checkpointing'):
+            self.backbone.grad_checkpointing = enable
+
+    def forward(self, x) -> Tuple[torch.Tensor, List[int]]:
+        x = self.backbone(x)
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+        x = self.proj(x)
+        return x.flatten(2).transpose(1, 2), x.shape[-2:]

reference_medclipseg/biomedclip/layers/inplace_abn.py

@@ -0,0 +1,87 @@
+import torch
+from torch import nn as nn
+
+try:
+    from inplace_abn.functions import inplace_abn, inplace_abn_sync
+    has_iabn = True
+except ImportError:
+    has_iabn = False
+
+    def inplace_abn(x, weight, bias, running_mean, running_var,
+                    training=True, momentum=0.1, eps=1e-05, activation="leaky_relu", activation_param=0.01):
+        raise ImportError(
+            "Please install InplaceABN:'pip install git+https://github.com/mapillary/inplace_abn.git@v1.0.12'")
+
+    def inplace_abn_sync(**kwargs):
+        inplace_abn(**kwargs)
+
+
+class InplaceAbn(nn.Module):
+    """Activated Batch Normalization
+
+    This gathers a BatchNorm and an activation function in a single module
+
+    Parameters
+    ----------
+    num_features : int
+        Number of feature channels in the input and output.
+    eps : float
+        Small constant to prevent numerical issues.
+    momentum : float
+        Momentum factor applied to compute running statistics.
+    affine : bool
+        If `True` apply learned scale and shift transformation after normalization.
+    act_layer : str or nn.Module type
+        Name or type of the activation functions, one of: `leaky_relu`, `elu`
+    act_param : float
+        Negative slope for the `leaky_relu` activation.
+    """
+
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, apply_act=True,
+                 act_layer="leaky_relu", act_param=0.01, drop_layer=None):
+        super(InplaceAbn, self).__init__()
+        self.num_features = num_features
+        self.affine = affine
+        self.eps = eps
+        self.momentum = momentum
+        if apply_act:
+            if isinstance(act_layer, str):
+                assert act_layer in ('leaky_relu', 'elu', 'identity', '')
+                self.act_name = act_layer if act_layer else 'identity'
+            else:
+                # convert act layer passed as type to string
+                if act_layer == nn.ELU:
+                    self.act_name = 'elu'
+                elif act_layer == nn.LeakyReLU:
+                    self.act_name = 'leaky_relu'
+                elif act_layer is None or act_layer == nn.Identity:
+                    self.act_name = 'identity'
+                else:
+                    assert False, f'Invalid act layer {act_layer.__name__} for IABN'
+        else:
+            self.act_name = 'identity'
+        self.act_param = act_param
+        if self.affine:
+            self.weight = nn.Parameter(torch.ones(num_features))
+            self.bias = nn.Parameter(torch.zeros(num_features))
+        else:
+            self.register_parameter('weight', None)
+            self.register_parameter('bias', None)
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.constant_(self.running_mean, 0)
+        nn.init.constant_(self.running_var, 1)
+        if self.affine:
+            nn.init.constant_(self.weight, 1)
+            nn.init.constant_(self.bias, 0)
+
+    def forward(self, x):
+        output = inplace_abn(
+            x, self.weight, self.bias, self.running_mean, self.running_var,
+            self.training, self.momentum, self.eps, self.act_name, self.act_param)
+        if isinstance(output, tuple):
+            output = output[0]
+        return output

reference_medclipseg/biomedclip/layers/interpolate.py

@@ -0,0 +1,68 @@
+""" Interpolation helpers for timm layers
+
+RegularGridInterpolator from https://github.com/sbarratt/torch_interpolations
+Copyright Shane Barratt, Apache 2.0 license
+"""
+import torch
+from itertools import product
+
+
+class RegularGridInterpolator:
+    """ Interpolate data defined on a rectilinear grid with even or uneven spacing.
+    Produces similar results to scipy RegularGridInterpolator or interp2d
+    in 'linear' mode.
+
+    Taken from https://github.com/sbarratt/torch_interpolations
+    """
+
+    def __init__(self, points, values):
+        self.points = points
+        self.values = values
+
+        assert isinstance(self.points, tuple) or isinstance(self.points, list)
+        assert isinstance(self.values, torch.Tensor)
+
+        self.ms = list(self.values.shape)
+        self.n = len(self.points)
+
+        assert len(self.ms) == self.n
+
+        for i, p in enumerate(self.points):
+            assert isinstance(p, torch.Tensor)
+            assert p.shape[0] == self.values.shape[i]
+
+    def __call__(self, points_to_interp):
+        assert self.points is not None
+        assert self.values is not None
+
+        assert len(points_to_interp) == len(self.points)
+        K = points_to_interp[0].shape[0]
+        for x in points_to_interp:
+            assert x.shape[0] == K
+
+        idxs = []
+        dists = []
+        overalls = []
+        for p, x in zip(self.points, points_to_interp):
+            idx_right = torch.bucketize(x, p)
+            idx_right[idx_right >= p.shape[0]] = p.shape[0] - 1
+            idx_left = (idx_right - 1).clamp(0, p.shape[0] - 1)
+            dist_left = x - p[idx_left]
+            dist_right = p[idx_right] - x
+            dist_left[dist_left < 0] = 0.
+            dist_right[dist_right < 0] = 0.
+            both_zero = (dist_left == 0) & (dist_right == 0)
+            dist_left[both_zero] = dist_right[both_zero] = 1.
+
+            idxs.append((idx_left, idx_right))
+            dists.append((dist_left, dist_right))
+            overalls.append(dist_left + dist_right)
+
+        numerator = 0.
+        for indexer in product([0, 1], repeat=self.n):
+            as_s = [idx[onoff] for onoff, idx in zip(indexer, idxs)]
+            bs_s = [dist[1 - onoff] for onoff, dist in zip(indexer, dists)]
+            numerator += self.values[as_s] * \
+                torch.prod(torch.stack(bs_s), dim=0)
+        denominator = torch.prod(torch.stack(overalls), dim=0)
+        return numerator / denominator

reference_medclipseg/biomedclip/layers/lambda_layer.py

@@ -0,0 +1,134 @@
+""" Lambda Layer
+
+Paper: `LambdaNetworks: Modeling Long-Range Interactions Without Attention`
+    - https://arxiv.org/abs/2102.08602
+
+@misc{2102.08602,
+Author = {Irwan Bello},
+Title = {LambdaNetworks: Modeling Long-Range Interactions Without Attention},
+Year = {2021},
+}
+
+Status:
+This impl is a WIP. Code snippets in the paper were used as reference but
+good chance some details are missing/wrong.
+
+I've only implemented local lambda conv based pos embeddings.
+
+For a PyTorch impl that includes other embedding options checkout
+https://github.com/lucidrains/lambda-networks
+
+Hacked together by / Copyright 2021 Ross Wightman
+"""
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .grid import ndgrid
+from .helpers import to_2tuple, make_divisible
+from .weight_init import trunc_normal_
+
+
+def rel_pos_indices(size):
+    size = to_2tuple(size)
+    pos = torch.stack(ndgrid(torch.arange(size[0]), torch.arange(size[1]))).flatten(1)
+    rel_pos = pos[:, None, :] - pos[:, :, None]
+    rel_pos[0] += size[0] - 1
+    rel_pos[1] += size[1] - 1
+    return rel_pos  # 2, H * W, H * W
+
+
+class LambdaLayer(nn.Module):
+    """Lambda Layer
+
+    Paper: `LambdaNetworks: Modeling Long-Range Interactions Without Attention`
+        - https://arxiv.org/abs/2102.08602
+
+    NOTE: intra-depth parameter 'u' is fixed at 1. It did not appear worth the complexity to add.
+
+    The internal dimensions of the lambda module are controlled via the interaction of several arguments.
+      * the output dimension of the module is specified by dim_out, which falls back to input dim if not set
+      * the value (v) dimension is set to dim_out // num_heads, the v projection determines the output dim
+      * the query (q) and key (k) dimension are determined by
+        * dim_head = (dim_out * attn_ratio // num_heads) if dim_head is None
+        * q = num_heads * dim_head, k = dim_head
+      * as seen above, attn_ratio determines the ratio of q and k relative to the output if dim_head not set
+
+    Args:
+        dim (int): input dimension to the module
+        dim_out (int): output dimension of the module, same as dim if not set
+        feat_size (Tuple[int, int]): size of input feature_map for relative pos variant H, W
+        stride (int): output stride of the module, avg pool used if stride == 2
+        num_heads (int): parallel attention heads.
+        dim_head (int): dimension of query and key heads, calculated from dim_out * attn_ratio // num_heads if not set
+        r (int): local lambda convolution radius. Use lambda conv if set, else relative pos if not. (default: 9)
+        qk_ratio (float): ratio of q and k dimensions to output dimension when dim_head not set. (default: 1.0)
+        qkv_bias (bool): add bias to q, k, and v projections
+    """
+    def __init__(
+            self, dim, dim_out=None, feat_size=None, stride=1, num_heads=4, dim_head=16, r=9,
+            qk_ratio=1.0, qkv_bias=False):
+        super().__init__()
+        dim_out = dim_out or dim
+        assert dim_out % num_heads == 0, ' should be divided by num_heads'
+        self.dim_qk = dim_head or make_divisible(dim_out * qk_ratio, divisor=8) // num_heads
+        self.num_heads = num_heads
+        self.dim_v = dim_out // num_heads
+
+        self.qkv = nn.Conv2d(
+            dim,
+            num_heads * self.dim_qk + self.dim_qk + self.dim_v,
+            kernel_size=1, bias=qkv_bias)
+        self.norm_q = nn.BatchNorm2d(num_heads * self.dim_qk)
+        self.norm_v = nn.BatchNorm2d(self.dim_v)
+
+        if r is not None:
+            # local lambda convolution for pos
+            self.conv_lambda = nn.Conv3d(1, self.dim_qk, (r, r, 1), padding=(r // 2, r // 2, 0))
+            self.pos_emb = None
+            self.rel_pos_indices = None
+        else:
+            # relative pos embedding
+            assert feat_size is not None
+            feat_size = to_2tuple(feat_size)
+            rel_size = [2 * s - 1 for s in feat_size]
+            self.conv_lambda = None
+            self.pos_emb = nn.Parameter(torch.zeros(rel_size[0], rel_size[1], self.dim_qk))
+            self.register_buffer('rel_pos_indices', rel_pos_indices(feat_size), persistent=False)
+
+        self.pool = nn.AvgPool2d(2, 2) if stride == 2 else nn.Identity()
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        trunc_normal_(self.qkv.weight, std=self.qkv.weight.shape[1] ** -0.5)  # fan-in
+        if self.conv_lambda is not None:
+            trunc_normal_(self.conv_lambda.weight, std=self.dim_qk ** -0.5)
+        if self.pos_emb is not None:
+            trunc_normal_(self.pos_emb, std=.02)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        M = H * W
+        qkv = self.qkv(x)
+        q, k, v = torch.split(qkv, [
+            self.num_heads * self.dim_qk, self.dim_qk, self.dim_v], dim=1)
+        q = self.norm_q(q).reshape(B, self.num_heads, self.dim_qk, M).transpose(-1, -2)  # B, num_heads, M, K
+        v = self.norm_v(v).reshape(B, self.dim_v, M).transpose(-1, -2)  # B, M, V
+        k = F.softmax(k.reshape(B, self.dim_qk, M), dim=-1)  # B, K, M
+
+        content_lam = k @ v  # B, K, V
+        content_out = q @ content_lam.unsqueeze(1)  # B, num_heads, M, V
+
+        if self.pos_emb is None:
+            position_lam = self.conv_lambda(v.reshape(B, 1, H, W, self.dim_v))  # B, H, W, V, K
+            position_lam = position_lam.reshape(B, 1, self.dim_qk, H * W, self.dim_v).transpose(2, 3)  # B, 1, M, K, V
+        else:
+            # FIXME relative pos embedding path not fully verified
+            pos_emb = self.pos_emb[self.rel_pos_indices[0], self.rel_pos_indices[1]].expand(B, -1, -1, -1)
+            position_lam = (pos_emb.transpose(-1, -2) @ v.unsqueeze(1)).unsqueeze(1)  # B, 1, M, K, V
+        position_out = (q.unsqueeze(-2) @ position_lam).squeeze(-2)  # B, num_heads, M, V
+
+        out = (content_out + position_out).transpose(-1, -2).reshape(B, C, H, W)  # B, C (num_heads * V), H, W
+        out = self.pool(out)
+        return out

reference_medclipseg/biomedclip/layers/layer_scale.py

@@ -0,0 +1,38 @@
+import torch
+from torch import nn
+
+
+class LayerScale(nn.Module):
+    """ LayerScale on tensors with channels in last-dim.
+    """
+    def __init__(
+            self,
+            dim: int,
+            init_values: float = 1e-5,
+            inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+class LayerScale2d(nn.Module):
+    """ LayerScale for tensors with torch 2D NCHW layout.
+    """
+    def __init__(
+            self,
+            dim: int,
+            init_values: float = 1e-5,
+            inplace: bool = False,
+    ):
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x):
+        gamma = self.gamma.view(1, -1, 1, 1)
+        return x.mul_(gamma) if self.inplace else x * gamma
+

reference_medclipseg/biomedclip/layers/linear.py

@@ -0,0 +1,19 @@
+""" Linear layer (alternate definition)
+"""
+import torch
+import torch.nn.functional as F
+from torch import nn as nn
+
+
+class Linear(nn.Linear):
+    r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b`
+
+    Wraps torch.nn.Linear to support AMP + torchscript usage by manually casting
+    weight & bias to input.dtype to work around an issue w/ torch.addmm in this use case.
+    """
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if torch.jit.is_scripting():
+            bias = self.bias.to(dtype=input.dtype) if self.bias is not None else None
+            return F.linear(input, self.weight.to(dtype=input.dtype), bias=bias)
+        else:
+            return F.linear(input, self.weight, self.bias)

reference_medclipseg/biomedclip/layers/median_pool.py

@@ -0,0 +1,49 @@
+""" Median Pool
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch.nn as nn
+import torch.nn.functional as F
+from .helpers import to_2tuple, to_4tuple
+
+
+class MedianPool2d(nn.Module):
+    """ Median pool (usable as median filter when stride=1) module.
+
+    Args:
+         kernel_size: size of pooling kernel, int or 2-tuple
+         stride: pool stride, int or 2-tuple
+         padding: pool padding, int or 4-tuple (l, r, t, b) as in pytorch F.pad
+         same: override padding and enforce same padding, boolean
+    """
+    def __init__(self, kernel_size=3, stride=1, padding=0, same=False):
+        super(MedianPool2d, self).__init__()
+        self.k = to_2tuple(kernel_size)
+        self.stride = to_2tuple(stride)
+        self.padding = to_4tuple(padding)  # convert to l, r, t, b
+        self.same = same
+
+    def _padding(self, x):
+        if self.same:
+            ih, iw = x.size()[2:]
+            if ih % self.stride[0] == 0:
+                ph = max(self.k[0] - self.stride[0], 0)
+            else:
+                ph = max(self.k[0] - (ih % self.stride[0]), 0)
+            if iw % self.stride[1] == 0:
+                pw = max(self.k[1] - self.stride[1], 0)
+            else:
+                pw = max(self.k[1] - (iw % self.stride[1]), 0)
+            pl = pw // 2
+            pr = pw - pl
+            pt = ph // 2
+            pb = ph - pt
+            padding = (pl, pr, pt, pb)
+        else:
+            padding = self.padding
+        return padding
+
+    def forward(self, x):
+        x = F.pad(x, self._padding(x), mode='reflect')
+        x = x.unfold(2, self.k[0], self.stride[0]).unfold(3, self.k[1], self.stride[1])
+        x = x.contiguous().view(x.size()[:4] + (-1,)).median(dim=-1)[0]
+        return x

reference_medclipseg/biomedclip/layers/mixed_conv2d.py

@@ -0,0 +1,51 @@
+""" PyTorch Mixed Convolution
+
+Paper: MixConv: Mixed Depthwise Convolutional Kernels (https://arxiv.org/abs/1907.09595)
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+
+import torch
+from torch import nn as nn
+
+from .conv2d_same import create_conv2d_pad
+
+
+def _split_channels(num_chan, num_groups):
+    split = [num_chan // num_groups for _ in range(num_groups)]
+    split[0] += num_chan - sum(split)
+    return split
+
+
+class MixedConv2d(nn.ModuleDict):
+    """ Mixed Grouped Convolution
+
+    Based on MDConv and GroupedConv in MixNet impl:
+      https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mixnet/custom_layers.py
+    """
+    def __init__(self, in_channels, out_channels, kernel_size=3,
+                 stride=1, padding='', dilation=1, depthwise=False, **kwargs):
+        super(MixedConv2d, self).__init__()
+
+        kernel_size = kernel_size if isinstance(kernel_size, list) else [kernel_size]
+        num_groups = len(kernel_size)
+        in_splits = _split_channels(in_channels, num_groups)
+        out_splits = _split_channels(out_channels, num_groups)
+        self.in_channels = sum(in_splits)
+        self.out_channels = sum(out_splits)
+        for idx, (k, in_ch, out_ch) in enumerate(zip(kernel_size, in_splits, out_splits)):
+            conv_groups = in_ch if depthwise else 1
+            # use add_module to keep key space clean
+            self.add_module(
+                str(idx),
+                create_conv2d_pad(
+                    in_ch, out_ch, k, stride=stride,
+                    padding=padding, dilation=dilation, groups=conv_groups, **kwargs)
+            )
+        self.splits = in_splits
+
+    def forward(self, x):
+        x_split = torch.split(x, self.splits, 1)
+        x_out = [c(x_split[i]) for i, c in enumerate(self.values())]
+        x = torch.cat(x_out, 1)
+        return x

reference_medclipseg/biomedclip/layers/ml_decoder.py

@@ -0,0 +1,146 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from torch import nn, Tensor
+from torch.nn.modules.transformer import _get_activation_fn
+
+
+def add_ml_decoder_head(model):
+    if hasattr(model, 'global_pool') and hasattr(model, 'fc'):  # most CNN models, like Resnet50
+        model.global_pool = nn.Identity()
+        del model.fc
+        num_classes = model.num_classes
+        num_features = model.num_features
+        model.fc = MLDecoder(num_classes=num_classes, initial_num_features=num_features)
+    elif hasattr(model, 'global_pool') and hasattr(model, 'classifier'):  # EfficientNet
+        model.global_pool = nn.Identity()
+        del model.classifier
+        num_classes = model.num_classes
+        num_features = model.num_features
+        model.classifier = MLDecoder(num_classes=num_classes, initial_num_features=num_features)
+    elif 'RegNet' in model._get_name() or 'TResNet' in model._get_name():  # hasattr(model, 'head')
+        del model.head
+        num_classes = model.num_classes
+        num_features = model.num_features
+        model.head = MLDecoder(num_classes=num_classes, initial_num_features=num_features)
+    else:
+        print("Model code-writing is not aligned currently with ml-decoder")
+        exit(-1)
+    if hasattr(model, 'drop_rate'):  # Ml-Decoder has inner dropout
+        model.drop_rate = 0
+    return model
+
+
+class TransformerDecoderLayerOptimal(nn.Module):
+    def __init__(self, d_model, nhead=8, dim_feedforward=2048, dropout=0.1, activation="relu",
+                 layer_norm_eps=1e-5) -> None:
+        super(TransformerDecoderLayerOptimal, self).__init__()
+        self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps)
+        self.dropout = nn.Dropout(dropout)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps)
+        self.norm3 = nn.LayerNorm(d_model, eps=layer_norm_eps)
+
+        self.activation = _get_activation_fn(activation)
+
+    def __setstate__(self, state):
+        if 'activation' not in state:
+            state['activation'] = torch.nn.functional.relu
+        super(TransformerDecoderLayerOptimal, self).__setstate__(state)
+
+    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
+        tgt = tgt + self.dropout1(tgt)
+        tgt = self.norm1(tgt)
+        tgt2 = self.multihead_attn(tgt, memory, memory)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+
+# class ExtrapClasses(object):
+#     def __init__(self, num_queries: int, group_size: int):
+#         self.num_queries = num_queries
+#         self.group_size = group_size
+#
+#     def __call__(self, h: torch.Tensor, class_embed_w: torch.Tensor, class_embed_b: torch.Tensor, out_extrap:
+#     torch.Tensor):
+#         # h = h.unsqueeze(-1).expand(-1, -1, -1, self.group_size)
+#         h = h[..., None].repeat(1, 1, 1, self.group_size) # torch.Size([bs, 5, 768, groups])
+#         w = class_embed_w.view((self.num_queries, h.shape[2], self.group_size))
+#         out = (h * w).sum(dim=2) + class_embed_b
+#         out = out.view((h.shape[0], self.group_size * self.num_queries))
+#         return out
+
+class MLDecoder(nn.Module):
+    def __init__(self, num_classes, num_of_groups=-1, decoder_embedding=768, initial_num_features=2048):
+        super(MLDecoder, self).__init__()
+        embed_len_decoder = 100 if num_of_groups < 0 else num_of_groups
+        if embed_len_decoder > num_classes:
+            embed_len_decoder = num_classes
+        self.embed_len_decoder = embed_len_decoder
+
+        # switching to 768 initial embeddings
+        decoder_embedding = 768 if decoder_embedding < 0 else decoder_embedding
+        self.embed_standart = nn.Linear(initial_num_features, decoder_embedding)
+
+        # decoder
+        decoder_dropout = 0.1
+        num_layers_decoder = 1
+        dim_feedforward = 2048
+        layer_decode = TransformerDecoderLayerOptimal(d_model=decoder_embedding,
+                                                      dim_feedforward=dim_feedforward, dropout=decoder_dropout)
+        self.decoder = nn.TransformerDecoder(layer_decode, num_layers=num_layers_decoder)
+
+        # non-learnable queries
+        self.query_embed = nn.Embedding(embed_len_decoder, decoder_embedding)
+        self.query_embed.requires_grad_(False)
+
+        # group fully-connected
+        self.num_classes = num_classes
+        self.duplicate_factor = int(num_classes / embed_len_decoder + 0.999)
+        self.duplicate_pooling = torch.nn.Parameter(
+            torch.Tensor(embed_len_decoder, decoder_embedding, self.duplicate_factor))
+        self.duplicate_pooling_bias = torch.nn.Parameter(torch.Tensor(num_classes))
+        torch.nn.init.xavier_normal_(self.duplicate_pooling)
+        torch.nn.init.constant_(self.duplicate_pooling_bias, 0)
+
+    def forward(self, x):
+        if len(x.shape) == 4:  # [bs,2048, 7,7]
+            embedding_spatial = x.flatten(2).transpose(1, 2)
+        else:  # [bs, 197,468]
+            embedding_spatial = x
+        embedding_spatial_786 = self.embed_standart(embedding_spatial)
+        embedding_spatial_786 = torch.nn.functional.relu(embedding_spatial_786, inplace=True)
+
+        bs = embedding_spatial_786.shape[0]
+        query_embed = self.query_embed.weight
+        # tgt = query_embed.unsqueeze(1).repeat(1, bs, 1)
+        tgt = query_embed.unsqueeze(1).expand(-1, bs, -1)  # no allocation of memory with expand
+        h = self.decoder(tgt, embedding_spatial_786.transpose(0, 1))  # [embed_len_decoder, batch, 768]
+        h = h.transpose(0, 1)
+
+        out_extrap = torch.zeros(h.shape[0], h.shape[1], self.duplicate_factor, device=h.device, dtype=h.dtype)
+        for i in range(self.embed_len_decoder):  # group FC
+            h_i = h[:, i, :]
+            w_i = self.duplicate_pooling[i, :, :]
+            out_extrap[:, i, :] = torch.matmul(h_i, w_i)
+        h_out = out_extrap.flatten(1)[:, :self.num_classes]
+        h_out += self.duplicate_pooling_bias
+        logits = h_out
+        return logits