Jolia zero-shot CT demo

README.md

@@ -1,38 +1,52 @@
 ---
-title: Jolia 3D CT Zero-Shot
+title: Jolia
 emoji: 🫁
 colorFrom: red
 colorTo: gray
 sdk: gradio
 sdk_version: 6.22.0
 app_file: app.py
-short_description: Zero-shot 3D CT classification with Jolia
+short_description: Zero-shot CT findings with the Jolia 3D CT foundation model
 python_version: "3.12"
 startup_duration_timeout: 1h
 pinned: false
+license: other
 ---
 
-# 🫁 Jolia — 3D CT Foundation Model (Zero-Shot)
+# Jolia — zero-shot CT analysis
 
-Interactive demo of [`raidium/Jolia`](https://huggingface.co/raidium/Jolia), a 3D CT foundation
-model that encodes a whole CT volume into a global embedding plus 102 per-organ query embeddings,
-all aligned with radiology-report text via CLIP.
+Demo of [`raidium/Jolia`](https://huggingface.co/raidium/Jolia), a 3D CT foundation model that
+encodes a whole chest / abdominal CT volume into a global embedding **and** 102 named organ-query
+embeddings, both aligned with radiology-report text.
 
-Upload a chest/abdominal CT volume (NIfTI `.nii` / `.nii.gz`) and score it against arbitrary text
-prompts — no fine-tuning:
+Upload a NIfTI CT volume and:
 
-- **Global zero-shot** — the whole-volume embedding vs each text prompt (calibrated CLIP logits → match probability).
-- **Per-organ zero-shot** — routes short findings phrases to a specific organ's query embedding via the ParallelOrganCLIP head.
+- score free-text findings against the **whole volume** (global CLIP head), and
+- route short findings phrases to a **single organ query** (ParallelOrganCLIP head, each organ with
+  its own trained temperature and bias).
 
-Paired text encoder: [`Qwen/Qwen3-Embedding-8B`](https://huggingface.co/Qwen/Qwen3-Embedding-8B) (bf16).
+The pipeline follows `example_zero_shot.py` from the model repo exactly: `JoliaPreprocessor`
+(1.5 mm isotropic, 192³ centre crop, 11 CT windowing channels) for the image, and the paired
+[`Qwen/Qwen3-Embedding-8B`](https://huggingface.co/Qwen/Qwen3-Embedding-8B) text encoder
+(last-token pooling, context length 512) for the prompts.
 
-## Example data
+> ⚠️ Research preview. Not a medical device; not for clinical use.
 
-The bundled example CT volumes come from the **TotalSegmentator** dataset
-([MedOtter/totalsegmentator-organs](https://huggingface.co/datasets/MedOtter/totalsegmentator-organs),
-CC-BY-4.0; Wasserthal et al., *Radiology: Artificial Intelligence* 2023).
+## Example volumes
 
-## Disclaimer
+The bundled example CTs come from the **TotalSegmentator dataset**
+(Wasserthal et al., [Zenodo record 10047292](https://zenodo.org/records/10047292), **CC-BY-4.0**),
+downloaded via [`YongchengYAO/TotalSegmentator-CT-Lite`](https://huggingface.co/datasets/YongchengYAO/TotalSegmentator-CT-Lite).
+File names carry that dataset's own study-type / pathology metadata. Attribution:
 
-⚠️ Research preview. **Not a medical device and not for clinical use.** Jolia is a feature extractor
-for research on adult chest/abdominal CT and does not produce diagnoses.
+> Wasserthal, J. et al. *TotalSegmentator: Robust segmentation of 104 anatomic structures in CT
+> images.* Radiology: Artificial Intelligence (2023). Dataset licensed CC-BY-4.0.
+
+## Notes
+
+- Volumes are reoriented to the radiological axial layout (rows anterior→posterior, columns
+  right→left, slices inferior→superior) before `JoliaPreprocessor`, which then flips depth and
+  centre-crops.
+- Probabilities are `sigmoid(calibrated logit)` — a per-pair "is this a match?" score, not a
+  softmax over prompts.
+- DICOM series can be converted to NIfTI with `dcm2niix`.

app.py

@@ -1,280 +1,430 @@
+"""Jolia — zero-shot CT analysis demo.
+
+Upload a chest / abdominal CT volume (NIfTI) and score it against free-text
+findings, either against the whole volume (global CLIP head) or routed to a
+specific organ query (ParallelOrganCLIP head).
+
+Mirrors `example_zero_shot.py` from the raidium/Jolia repo 1:1: same
+preprocessing (`JoliaPreprocessor`), same paired text encoder
+(Qwen3-Embedding-8B with last-token pooling), same calibrated logits.
+"""
+
 import os
+
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
+import spaces  # noqa: F401  # must precede torch / any CUDA-touching import
+
 import sys
 import time
-import numpy as np
-import spaces
-import torch
+
 import gradio as gr
 import nibabel as nib
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
 from huggingface_hub import snapshot_download
 from transformers import AutoModel
 
-# ----------------------------------------------------------------------------
-# Model loading (module scope, eager .to("cuda") for ZeroGPU)
-# ----------------------------------------------------------------------------
 JOLIA_ID = "raidium/Jolia"
 TEXT_ID = "Qwen/Qwen3-Embedding-8B"
+MAX_PROMPTS = 10
+MAX_ORGANS = 16
+CACHE_VERSION = "v1"
 
-# Jolia ships its preprocessing + text-encoder helpers as plain modules in the
-# repo. Grab them via snapshot_download and put them on sys.path.
+# The Jolia repo ships its own preprocessing / text-encoder helpers.
 _repo = snapshot_download(JOLIA_ID)
-sys.path.append(_repo)
+if _repo not in sys.path:
+    sys.path.insert(0, _repo)
+
+from jolia_windowing import get_available_windows  # noqa: E402
 from preprocessing_jolia import JoliaPreprocessor  # noqa: E402
 from text_encoder_jolia import JoliaTextEncoder  # noqa: E402
 
-print("Loading Jolia vision backbone ...")
-jolia = AutoModel.from_pretrained(JOLIA_ID, trust_remote_code=True).eval().to("cuda")
+PRE = JoliaPreprocessor()
+CT_WINDOWS = get_available_windows("CT")  # channel order of the 11 windowing channels
+PREVIEW_WINDOWS = ["auto", "lung", "mediastinum", "abdomen", "liver", "bone", "soft_tissue"]
 
-print("Loading paired text encoder (Qwen3-Embedding-8B, bf16) ...")
-text_encoder = (
+print("[1/2] Loading Jolia vision backbone ...", flush=True)
+JOLIA = AutoModel.from_pretrained(JOLIA_ID, trust_remote_code=True).eval().to("cuda")
+print("[2/2] Loading paired text encoder Qwen3-Embedding-8B (~15 GB) ...", flush=True)
+TEXT = (
     JoliaTextEncoder.from_pretrained(
         TEXT_ID,
         dtype=torch.bfloat16,
-        context_length=jolia.config.text_context_length,
+        context_length=JOLIA.config.text_context_length,
     )
     .eval()
     .to("cuda")
 )
-
-preprocessor = JoliaPreprocessor()
-ORGAN_NAMES = jolia.organ_slot_names  # 102 named organ slots
-print(f"Ready. {len(ORGAN_NAMES)} organ slots available.")
-
-DEFAULT_PROMPTS = (
-    "a chest CT showing pneumonia\n"
-    "a CT with a lung nodule\n"
-    "a CT showing pleural effusion\n"
-    "a normal chest CT\n"
-    "a CT showing cardiomegaly"
+ORGAN_NAMES = list(JOLIA.organ_slot_names)
+print(f"Ready — {len(ORGAN_NAMES)} organ slots available.", flush=True)
+
+DEFAULT_ORGANS = [
+    "lungs",
+    "pleura",
+    "heart",
+    "mediastinum",
+    "liver",
+    "kidneys",
+    "spleen",
+    "pancreas",
+    "spine",
+]
+DEFAULT_VOLUME_PROMPTS = "\n".join(
+    [
+        "a normal chest CT",
+        "a chest CT showing a pulmonary nodule",
+        "a chest CT showing pneumonia",
+        "a chest CT showing pleural effusion",
+        "a CT showing a liver lesion",
+    ]
 )
-DEFAULT_ORGANS = ["lungs", "heart", "liver", "spleen", "kidneys", "pancreas"]
-DEFAULT_ORGAN_PROMPTS = "a lesion\nan enlarged organ\nlooks normal"
+DEFAULT_ORGAN_PROMPTS = "\n".join(["looks normal", "a lesion", "a mass", "an enlarged organ"])
 
 
 # ----------------------------------------------------------------------------
-# NIfTI reading + light preview (CPU, cheap)
+# CT loading / preprocessing
 # ----------------------------------------------------------------------------
-def _read_nifti(path):
-    """Load a NIfTI CT volume -> (volume float32 (H,W,D), spacing (r,c,s) mm)."""
-    img = nib.load(path)
-    vol = np.asarray(img.get_fdata(), dtype=np.float32)
-    if vol.ndim != 3:
-        raise gr.Error(
-            f"Expected a 3D CT volume, got shape {vol.shape}. "
-            "Upload a single-series CT scan in NIfTI (.nii / .nii.gz) format."
-        )
-    zooms = img.header.get_zooms()[:3]
-    spacing = tuple(float(z) if z and z > 0 else 1.0 for z in zooms)
-    return vol, spacing
-
-
-def _window(slice_hu, center=40.0, width=400.0):
-    """Apply a CT window (default soft-tissue) and map to 0-255 uint8."""
-    lo, hi = center - width / 2.0, center + width / 2.0
-    s = np.clip((slice_hu - lo) / (hi - lo), 0.0, 1.0)
-    return (s * 255.0).astype(np.uint8)
-
-
-def preview_nifti(path):
-    """Render 3 mid-slices (axial / coronal / sagittal) of the uploaded CT."""
-    if not path:
-        return None
-    vol, spacing = _read_nifti(path)
-    # vol is (H, W, D). Build a simple montage of mid slices.
-    h, w, d = vol.shape
-    ax = np.rot90(_window(vol[:, :, d // 2]))          # axial
-    cor = np.rot90(_window(vol[:, w // 2, :]))         # coronal
-    sag = np.rot90(_window(vol[h // 2, :, :]))         # sagittal
-    # Pad to same height and concat side by side.
-    hh = max(ax.shape[0], cor.shape[0], sag.shape[0])
-
-    def _pad(a):
-        pad = hh - a.shape[0]
-        if pad > 0:
-            a = np.pad(a, ((0, pad), (0, 0)), constant_values=0)
-        return a
-
-    montage = np.concatenate([_pad(ax), _pad(cor), _pad(sag)], axis=1)
-    return montage
-
-
-def _parse_prompts(text):
-    return [ln.strip() for ln in (text or "").replace(",", "\n").splitlines() if ln.strip()]
+def _load_ct(path: str):
+    """NIfTI file -> (volume (H, W, D) in HU, resolution (row, col, slice) mm, info)."""
+    try:
+        img = nib.load(path)
+    except Exception as exc:  # noqa: BLE001
+        raise gr.Error(f"Could not read this file as NIfTI ({exc}). Convert DICOM with dcm2niix first.")
+    try:
+        img = nib.as_closest_canonical(img)  # reorient to RAS+
+    except Exception:  # noqa: BLE001
+        pass
+
+    arr = np.asanyarray(img.dataobj)
+    while arr.ndim > 3:
+        arr = arr[..., 0]
+    if arr.ndim != 3:
+        raise gr.Error(f"Expected a 3D volume, got shape {tuple(arr.shape)}.")
+    arr = np.nan_to_num(arr.astype(np.float32), nan=-1024.0)
+    zooms = [float(z) for z in img.header.get_zooms()[:3]]
+    zooms = [z if z > 0 else 1.0 for z in zooms]
+
+    # RAS+ (x->Right, y->Anterior, z->Superior) to the radiological axial layout
+    # the checkpoint was trained on: rows anterior->posterior, columns
+    # right->left, slices inferior->superior (PrepareVolume then flips depth).
+    vol = np.ascontiguousarray(arr.transpose(1, 0, 2)[::-1, ::-1, :])
+    resolution = (zooms[1], zooms[0], zooms[2])  # (row, col, slice) mm
+    info = {
+        "shape": tuple(int(s) for s in arr.shape),
+        "spacing": tuple(round(z, 3) for z in zooms),
+        "hu_range": (float(np.percentile(vol, 0.5)), float(np.percentile(vol, 99.5))),
+        "z_coverage_mm": round(arr.shape[2] * zooms[2], 1),
+    }
+    return vol, resolution, info
+
+
+def _auto_window(vol: np.ndarray) -> str:
+    """Pick a sensible display window: lung if there is lung parenchyma, else abdomen."""
+    lung_frac = float(np.mean((vol > -900.0) & (vol < -500.0)))
+    return "lung" if lung_frac > 0.06 else "abdomen"
+
+
+def _u8(plane: np.ndarray) -> np.ndarray:
+    img = (np.clip(plane, 0.0, 1.0) * 255.0).astype(np.uint8)
+    return np.repeat(np.repeat(img, 2, axis=0), 2, axis=1)  # 192 -> 384 px
+
+
+def _preview_tiles(image: torch.Tensor, window: str) -> list:
+    """Orthogonal previews of the exact 192**3 cube the model sees."""
+    vol = image[CT_WINDOWS.index(window)].float().numpy()
+    depth, height, width = vol.shape
+    tiles = []
+    for frac in (0.3, 0.5, 0.7):
+        idx = int(round(frac * (depth - 1)))
+        tiles.append((_u8(vol[idx]), f"axial · slice {idx}/{depth - 1}"))
+    tiles.append((_u8(vol[:, height // 2, :]), "coronal · mid"))
+    tiles.append((_u8(vol[:, :, width // 2]), "sagittal · mid"))
+    return tiles
+
+
+def _prep(path: str, preview_window: str):
+    """Load + preprocess a CT and render previews. CPU only."""
+    vol, resolution, info = _load_ct(path)
+    window = _auto_window(vol) if preview_window == "auto" else preview_window
+    image = PRE(vol, resolution=resolution)  # (11, 192, 192, 192) float32
+    return image, _preview_tiles(image, window), info, window
+
+
+def _volume_summary(info: dict, window: str, extra: str = "") -> str:
+    sx, sy, sz = info["spacing"]
+    lo, hi = info["hu_range"]
+    return (
+        f"**Volume** {info['shape'][0]}×{info['shape'][1]}×{info['shape'][2]} @ "
+        f"{sx}×{sy}×{sz} mm · {info['z_coverage_mm']} mm cranio-caudal coverage · "
+        f"HU p0.5–p99.5 {lo:.0f} → {hi:.0f}  \n"
+        f"**Model input** 11×192×192×192 (1.5 mm isotropic, centre crop) · preview window `{window}`"
+        + (f"  \n{extra}" if extra else "")
+    )
 
 
-def _estimate_duration(nifti_path, prompts_text, organs, organ_prompts_text, *args, **kwargs):
-    n_prompts = max(1, len(_parse_prompts(prompts_text)) + len(_parse_prompts(organ_prompts_text)))
-    n_organs = max(1, len(organs or []))
-    # Measured GPU compute is ~2s; the rest is worker cold-start + I/O headroom.
-    return min(90, 20 + n_prompts + n_organs)
+def _parse_lines(text: str, limit: int) -> list:
+    lines = [ln.strip() for ln in (text or "").splitlines()]
+    return [ln for ln in lines if ln][:limit]
 
 
 # ----------------------------------------------------------------------------
-# Inference (GPU)
+# Inference
 # ----------------------------------------------------------------------------
-@spaces.GPU(duration=_estimate_duration)
-def classify(nifti_path, prompts_text, organs, organ_prompts_text,
-             progress=gr.Progress(track_tqdm=True)):
-    """Zero-shot classification of a 3D CT volume with the Jolia foundation model.
+@spaces.GPU(duration=90)
+def analyze(
+    ct_file: str,
+    volume_prompts: str = DEFAULT_VOLUME_PROMPTS,
+    organ_prompts: str = DEFAULT_ORGAN_PROMPTS,
+    organs: list = DEFAULT_ORGANS,
+    preview_window: str = "auto",
+):
+    """Zero-shot classify a CT volume against free-text findings with Jolia.
 
     Args:
-        nifti_path: path to a CT volume in NIfTI (.nii/.nii.gz) format.
-        prompts_text: newline/comma separated text prompts scored against the
-            whole-volume (global) embedding.
-        organs: list of organ names to run per-organ (query-routed) zero-shot on.
-        organ_prompts_text: short findings phrases scored per organ.
+        ct_file: Path to a chest / abdominal CT volume in NIfTI format (.nii or .nii.gz).
+        volume_prompts: Whole-volume prompts, one per line (global CLIP head).
+        organ_prompts: Short findings phrases, one per line (per-organ CLIP head).
+        organs: Organ query slots to route the findings phrases to.
+        preview_window: CT display window for the preview images ("auto" picks lung or abdomen).
 
     Returns:
-        A global-classification table and a per-organ score table.
+        Orthogonal previews of the model input, a volume summary, the global
+        zero-shot table and the per-organ probability matrix.
     """
-    if not nifti_path:
-        raise gr.Error("Please upload a CT volume (.nii or .nii.gz) first.")
-
-    prompts = _parse_prompts(prompts_text)
-    if not prompts:
-        raise gr.Error("Please enter at least one text prompt for global classification.")
-    organ_prompts = _parse_prompts(organ_prompts_text)
+    if not ct_file:
+        raise gr.Error("Upload a CT volume (NIfTI .nii / .nii.gz) first.")
+    vol_prompts = _parse_lines(volume_prompts, MAX_PROMPTS)
+    org_prompts = _parse_lines(organ_prompts, MAX_PROMPTS)
+    organs = [o for o in (organs or []) if o in ORGAN_NAMES][:MAX_ORGANS]
+    if not vol_prompts and not org_prompts:
+        raise gr.Error("Enter at least one prompt.")
 
     t0 = time.perf_counter()
-    vol, spacing = _read_nifti(nifti_path)
-    image = preprocessor(vol, resolution=spacing).unsqueeze(0).to("cuda")
+    image, tiles, info, window = _prep(ct_file, preview_window)
+    t_prep = time.perf_counter() - t0
 
+    t0 = time.perf_counter()
     with torch.no_grad():
-        # --- Global zero-shot ---
-        text_features = text_encoder(prompts).to("cuda")
-        logits = jolia.zero_shot(image, text_features)          # (1, N) calibrated
-        probs = torch.sigmoid(logits)[0].float().cpu().numpy()
-        cosine = jolia.zero_shot(image, text_features, calibrated=False)[0].float().cpu().numpy()
-
-    order = np.argsort(-probs)
-    global_rows = [
-        [prompts[i], round(float(probs[i]), 4), round(float(cosine[i]), 4)]
-        for i in order
-    ]
-
-    # --- Per-organ zero-shot ---
-    organ_rows = []
-    if organs and organ_prompts:
-        with torch.no_grad():
-            organ_text = text_encoder(organ_prompts).to("cuda")
-            per_organ = jolia.zero_shot_organs(image, organ_text, organs=list(organs))
-        for organ in organs:
-            lg = per_organ[organ][0].float().cpu().numpy()
-            pr = 1.0 / (1.0 + np.exp(-lg))
-            row = [organ] + [round(float(x), 4) for x in pr]
-            organ_rows.append(row)
-
-    elapsed = time.perf_counter() - t0
-    organ_headers = ["organ"] + (organ_prompts if organ_prompts else ["(no prompts)"])
-    info = (
-        f"Volume {vol.shape} @ {tuple(round(s, 2) for s in spacing)} mm  •  "
-        f"{len(prompts)} global prompt(s), {len(organs or [])} organ(s)  •  "
-        f"{elapsed:.1f}s"
+        x = image.unsqueeze(0).to("cuda")
+        cls, organ_queries = JOLIA.forward_with_queries(x)  # (1, 576), (1, slots, 576)
+        image_emb = F.normalize(cls.float(), dim=-1, eps=1e-6)
+
+        global_rows = []
+        if vol_prompts:
+            text_features = TEXT(vol_prompts).to(image_emb.device)  # (N, 4096)
+            text_emb = JOLIA.encode_text(text_features)  # (N, 576)
+            cosine = (image_emb @ text_emb.t())[0]
+            logits = JOLIA.zero_shot_logits(image_emb, text_emb)[0]
+            probs = torch.sigmoid(logits)
+            global_rows = [
+                [p, round(float(lg), 4), round(float(pr), 4), round(float(cs), 4)]
+                for p, lg, pr, cs in zip(vol_prompts, logits, probs, cosine)
+            ]
+            global_rows.sort(key=lambda r: -r[1])
+
+        organ_rows = []
+        if org_prompts and organs:
+            organ_text = TEXT(org_prompts).to(image_emb.device)
+            organ_text_emb = JOLIA.encode_organ_text(organ_text)  # (N, 576)
+            for name in organs:
+                idx = ORGAN_NAMES.index(name)
+                emb = F.normalize(organ_queries[:, idx, :].float(), dim=-1, eps=1e-6)
+                scale = JOLIA.organ_logit_scale[idx].float().exp()
+                bias = JOLIA.organ_text_bias[idx].float()
+                logits = (emb @ organ_text_emb.t())[0] * scale + bias
+                organ_rows.append([name] + [round(float(v), 4) for v in torch.sigmoid(logits)])
+    t_gpu = time.perf_counter() - t0
+
+    global_df = pd.DataFrame(
+        global_rows or [["—", 0.0, 0.0, 0.0]],
+        columns=["prompt", "calibrated logit", "match probability", "cosine"],
     )
-    return (
-        gr.update(value=global_rows),
-        gr.update(value=organ_rows, headers=organ_headers),
+    organ_df = pd.DataFrame(
+        organ_rows or [["—"] + [0.0] * max(1, len(org_prompts))],
+        columns=["organ"] + (org_prompts or ["—"]),
+    )
+    best = f"**Top whole-volume match** · `{global_rows[0][0]}` (p={global_rows[0][2]:.3f})  \n" if global_rows else ""
+    summary = _volume_summary(
         info,
+        window,
+        f"{best}*preprocess {t_prep:.1f}s · encode + score {t_gpu:.1f}s*",
     )
+    return tiles, summary, global_df, organ_df
+
+
+def preview(ct_file: str, preview_window: str = "auto"):
+    """Render orthogonal previews of the preprocessed CT volume (no GPU).
+
+    Args:
+        ct_file: Path to a CT volume in NIfTI format (.nii or .nii.gz).
+        preview_window: CT display window ("auto" picks lung or abdomen).
+
+    Returns:
+        Preview images of the 192**3 model input and a short volume summary.
+    """
+    if not ct_file:
+        return [], "Upload a CT volume (NIfTI `.nii` / `.nii.gz`) to get started."
+    _, tiles, info, window = _prep(ct_file, preview_window)
+    return tiles, _volume_summary(info, window, "*Preview only — press **Analyze** to score prompts.*")
 
 
 # ----------------------------------------------------------------------------
 # UI
 # ----------------------------------------------------------------------------
 CSS = """
-#col-container { max-width: 1150px; margin: 0 auto; }
+#col-container { max-width: 1250px; margin: 0 auto; }
 .dark .gradio-container { color: var(--body-text-color); }
 """
 
-with gr.Blocks(title="Jolia — 3D CT Zero-Shot") as demo:
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="Jolia — zero-shot CT") as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(
-            """
-            # 🫁 Jolia — 3D CT Foundation Model (Zero-Shot)
-
-            [`raidium/Jolia`](https://huggingface.co/raidium/Jolia) is a 3D CT foundation model that
-            encodes a whole CT volume into a global embedding + 102 per-organ query embeddings, aligned
-            with radiology-report text via CLIP. Upload a chest/abdominal CT volume and score it against
-            **arbitrary text prompts** — no fine-tuning — both globally and per organ.
-
-            Paired text encoder: `Qwen/Qwen3-Embedding-8B`.
-            ⚠️ **Research preview — not a medical device, not for clinical use.**
-            """
+            "# Jolia — zero-shot CT analysis\n"
+            "[`raidium/Jolia`](https://huggingface.co/raidium/Jolia) is a 3D CT foundation model: it "
+            "encodes a whole chest / abdominal CT into one global embedding **and** 102 named "
+            "organ-query embeddings, both aligned with report text. Score any free-text finding "
+            "against the whole volume, or route it to a single organ.\n\n"
+            "⚠️ Research preview — **not a medical device, not for clinical use.**"
         )
-
         with gr.Row():
-            with gr.Column(scale=1):
-                nifti = gr.File(
-                    label="CT volume (NIfTI: .nii / .nii.gz)",
+            with gr.Column(scale=4):
+                ct_file = gr.File(
+                    label="CT volume (NIfTI .nii / .nii.gz)",
                     file_types=[".nii", ".gz"],
                     type="filepath",
                 )
-                preview = gr.Image(
-                    label="Preview (axial · coronal · sagittal, soft-tissue window)",
-                    height=240,
+                volume_prompts = gr.Textbox(
+                    label="Whole-volume prompts (one per line)",
+                    info="Scored against the global CLIP head — full sentences work best.",
+                    value=DEFAULT_VOLUME_PROMPTS,
+                    lines=5,
                 )
-                run = gr.Button("Run zero-shot classification", variant="primary")
-            with gr.Column(scale=1):
-                prompts_box = gr.Textbox(
-                    label="Global prompts (one per line)",
-                    value=DEFAULT_PROMPTS,
-                    lines=6,
+                organ_prompts = gr.Textbox(
+                    label="Per-organ findings phrases (one per line)",
+                    info="Scored against the per-organ head — short phrases work best.",
+                    value=DEFAULT_ORGAN_PROMPTS,
+                    lines=4,
                 )
-                with gr.Accordion("Per-organ zero-shot (query-routed)", open=True):
-                    organs_box = gr.Dropdown(
-                        label="Organs to score",
-                        choices=ORGAN_NAMES,
-                        value=DEFAULT_ORGANS,
-                        multiselect=True,
-                    )
-                    organ_prompts_box = gr.Textbox(
-                        label="Per-organ findings phrases (one per line)",
-                        value=DEFAULT_ORGAN_PROMPTS,
-                        lines=3,
+                organs = gr.Dropdown(
+                    label="Organ query slots",
+                    choices=ORGAN_NAMES,
+                    value=DEFAULT_ORGANS,
+                    multiselect=True,
+                    max_choices=MAX_ORGANS,
+                )
+                run = gr.Button("Analyze", variant="primary")
+                with gr.Accordion("Advanced", open=False):
+                    preview_window = gr.Dropdown(
+                        label="Preview window",
+                        choices=PREVIEW_WINDOWS,
+                        value="auto",
+                        info="Display only — the model always sees all 11 windowing channels.",
                     )
+            with gr.Column(scale=6):
+                gallery = gr.Gallery(
+                    label="Model input (1.5 mm isotropic, 192³ centre crop)",
+                    columns=3,
+                    height=340,
+                    object_fit="contain",
+                )
+                summary = gr.Markdown("Upload a CT volume (NIfTI `.nii` / `.nii.gz`) to get started.")
+                global_df = gr.Dataframe(
+                    label="Whole-volume zero-shot (global CLIP head)",
+                    headers=["prompt", "calibrated logit", "match probability", "cosine"],
+                    wrap=True,
+                )
+                organ_df = gr.Dataframe(
+                    label="Per-organ zero-shot — match probability per (organ, phrase)",
+                    wrap=True,
+                )
 
-        info_out = gr.Markdown()
-        global_out = gr.Dataframe(
-            headers=["prompt", "match probability", "cosine"],
-            label="Global zero-shot — whole-volume embedding vs text (sorted by probability)",
-            wrap=True,
-        )
-        organ_out = gr.Dataframe(
-            label="Per-organ zero-shot — organ-query embedding vs findings text (match probability)",
-            wrap=True,
+        gr.Markdown(
+            "### Examples\n"
+            "Public CT volumes from the [TotalSegmentator dataset](https://zenodo.org/records/10047292) "
+            "(Wasserthal et al., CC-BY-4.0), via "
+            "[`YongchengYAO/TotalSegmentator-CT-Lite`](https://huggingface.co/datasets/YongchengYAO/TotalSegmentator-CT-Lite). "
+            "Radiology labels in the file names come from that dataset's metadata."
         )
-
         gr.Examples(
             examples=[
-                ["examples/s0000_ct.nii.gz", DEFAULT_PROMPTS, DEFAULT_ORGANS, DEFAULT_ORGAN_PROMPTS],
-                ["examples/s0004_ct.nii.gz", DEFAULT_PROMPTS, DEFAULT_ORGANS, DEFAULT_ORGAN_PROMPTS],
-                ["examples/s0010_ct.nii.gz", DEFAULT_PROMPTS, DEFAULT_ORGANS, DEFAULT_ORGAN_PROMPTS],
+                [
+                    "examples/chest_ct_lung_tumor_s1173.nii.gz",
+                    "\n".join(
+                        [
+                            "a normal chest CT",
+                            "a chest CT showing a pulmonary nodule",
+                            "a chest CT showing pneumonia",
+                            "a chest CT showing pleural effusion",
+                            "a chest CT showing emphysema",
+                        ]
+                    ),
+                    "\n".join(["looks normal", "a nodule", "a mass", "an effusion"]),
+                ],
+                [
+                    "examples/chest_ct_inflammation_s1353.nii.gz",
+                    "\n".join(
+                        [
+                            "a normal chest CT",
+                            "a chest CT showing pneumonia",
+                            "a chest CT showing consolidation",
+                            "a chest CT showing a pulmonary nodule",
+                        ]
+                    ),
+                    "\n".join(["looks normal", "consolidation", "an infection", "a nodule"]),
+                ],
+                [
+                    "examples/abdomen_pelvis_ct_normal_s0143.nii.gz",
+                    "\n".join(
+                        [
+                            "a normal abdominal CT",
+                            "an abdominal CT showing a liver lesion",
+                            "an abdominal CT showing hepatic steatosis",
+                            "an abdominal CT showing bowel obstruction",
+                        ]
+                    ),
+                    "\n".join(["looks normal", "a lesion", "an enlarged organ"]),
+                ],
+                [
+                    "examples/abdomen_ct_tumor_s0168.nii.gz",
+                    "\n".join(
+                        [
+                            "a normal abdominal CT",
+                            "an abdominal CT showing a tumour",
+                            "an abdominal CT showing a liver lesion",
+                            "an abdominal CT showing enlarged lymph nodes",
+                        ]
+                    ),
+                    "\n".join(["looks normal", "a lesion", "a mass", "an enlarged organ"]),
+                ],
             ],
-            inputs=[nifti, prompts_box, organs_box, organ_prompts_box],
-            outputs=[global_out, organ_out, info_out],
-            fn=classify,
+            inputs=[ct_file, volume_prompts, organ_prompts],
+            outputs=[gallery, summary, global_df, organ_df],
+            fn=analyze,
             cache_examples=True,
             cache_mode="lazy",
+            label=f"Example CT volumes ({CACHE_VERSION})",
         )
 
         gr.Markdown(
-            """
-            ---
-            Example CT volumes from the **TotalSegmentator** dataset
-            ([MedOtter/totalsegmentator-organs](https://huggingface.co/datasets/MedOtter/totalsegmentator-organs),
-            CC-BY-4.0; Wasserthal et al.). Model: [raidium/Jolia](https://huggingface.co/raidium/Jolia).
-            """
+            "Whole-volume scores use Jolia's global CLIP head; per-organ scores route the phrase to "
+            "one organ query through the ParallelOrganCLIP head (each organ has its own trained "
+            "temperature and bias). Probabilities are `sigmoid(calibrated logit)` — a per-pair "
+            '"is this a match?" score, not a softmax over prompts, so they do not sum to 1. '
+            "Text is encoded with the paired [`Qwen/Qwen3-Embedding-8B`](https://huggingface.co/Qwen/Qwen3-Embedding-8B) "
+            "(last-token pooling, context length 512). DICOM series can be converted with `dcm2niix`."
         )
 
-    nifti.change(fn=preview_nifti, inputs=nifti, outputs=preview)
+    ct_file.change(preview, inputs=[ct_file, preview_window], outputs=[gallery, summary])
+    preview_window.change(preview, inputs=[ct_file, preview_window], outputs=[gallery, summary])
     run.click(
-        fn=classify,
-        inputs=[nifti, prompts_box, organs_box, organ_prompts_box],
-        outputs=[global_out, organ_out, info_out],
-        api_name="classify",
+        analyze,
+        inputs=[ct_file, volume_prompts, organ_prompts, organs, preview_window],
+        outputs=[gallery, summary, global_df, organ_df],
+        api_name="analyze",
     )
 
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)
+    demo.launch(mcp_server=True)

examples/abdomen_ct_tumor_s0168.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7fe41adbbffea3284c6675222bf1d0352efbaee3367574e428ad810c0aa747b
+size 16649551

examples/abdomen_pelvis_ct_normal_s0143.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b61aefa67756e2ff763e5509ab318af93060b18a61c0e63e165449aa6880e83d
+size 15146822

examples/chest_ct_inflammation_s1353.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb23dd1759048ce8992efc636e2dae767987e703f941dfa1f214bb5f3b34bc5e
+size 15963600

examples/chest_ct_lung_tumor_s1173.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7955aa6e0b065c404842be81815d6dbf5c0169a9c31da4514102ceda638ef787
+size 18445659

requirements.txt

@@ -1,7 +1,10 @@
-transformers
+torch
+torchvision==0.26.0
+transformers>=4.57.0
 timm
 einops
-numpy
 safetensors
-nibabel
 accelerate
+numpy
+pandas
+nibabel