multimodalart HF Staff
Gradio 6 theme/css on launch(); tighten ZeroGPU duration to 45s
7be29ab verified | """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 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 | |
| JOLIA_ID = "raidium/Jolia" | |
| TEXT_ID = "Qwen/Qwen3-Embedding-8B" | |
| MAX_PROMPTS = 10 | |
| MAX_ORGANS = 16 | |
| CACHE_VERSION = "v1" | |
| # The Jolia repo ships its own preprocessing / text-encoder helpers. | |
| _repo = snapshot_download(JOLIA_ID) | |
| 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 | |
| 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("[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, | |
| ) | |
| .eval() | |
| .to("cuda") | |
| ) | |
| 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_ORGAN_PROMPTS = "\n".join(["looks normal", "a lesion", "a mass", "an enlarged organ"]) | |
| # ---------------------------------------------------------------------------- | |
| # CT loading / preprocessing | |
| # ---------------------------------------------------------------------------- | |
| 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 _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 | |
| # ---------------------------------------------------------------------------- | |
| 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: | |
| 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: | |
| Orthogonal previews of the model input, a volume summary, the global | |
| zero-shot table and the per-organ probability matrix. | |
| """ | |
| 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() | |
| image, tiles, info, window = _prep(ct_file, preview_window) | |
| t_prep = time.perf_counter() - t0 | |
| t0 = time.perf_counter() | |
| with torch.no_grad(): | |
| 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"], | |
| ) | |
| 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: 1250px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="Jolia — zero-shot CT") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# 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=4): | |
| ct_file = gr.File( | |
| label="CT volume (NIfTI .nii / .nii.gz)", | |
| file_types=[".nii", ".gz"], | |
| type="filepath", | |
| ) | |
| 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, | |
| ) | |
| 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, | |
| ) | |
| 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, | |
| ) | |
| 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/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=[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( | |
| "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`." | |
| ) | |
| 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( | |
| 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) | |