Add PanCancerSeg Gradio inference app

.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
+example/FLARE23Ts_0005_overlay.mp4 filter=lfs diff=lfs merge=lfs -text
+example/FLARE23Ts_0005_slice_centroid.png filter=lfs diff=lfs merge=lfs -text
+example/FLARE23Ts_0005_slice_extent25.png filter=lfs diff=lfs merge=lfs -text
+example/FLARE23Ts_0005_slice_extent75.png filter=lfs diff=lfs merge=lfs -text
+example/FLARE23Ts_0005_slice_max_area.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+
+PanCancerSeg-Specialized-weights/
+
+output_kidney_test/
+
+sample_input/

README.md

@@ -1,13 +1,131 @@
 ---
 title: PanCancerSeg Specialist Inference
-emoji: 🌍
-colorFrom: pink
-colorTo: green
+emoji: 🩻
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.16.0
-python_version: '3.13'
+sdk_version: 6.12.0
 app_file: app.py
 pinned: false
+license: apache-2.0
+short_description: CT tumour segmentation for 4 cancer types
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# PanCancerSeg Inference
+
+Run one cancer-specific PanCancerSeg nnUNet model on a single CT NIfTI image and save a segmentation mask, slice PNG previews, and an MP4 overlay video.
+
+> On Hugging Face Spaces the trained weights are downloaded automatically at first
+> run from [`KS987/PanCancerSeg-Specialized-weights`](https://huggingface.co/KS987/PanCancerSeg-Specialized-weights).
+> Inference is GPU-recommended; the free CPU tier may be slow or run out of memory on large volumes.
+
+## Model Weights
+
+Download the trained nnUNet weights from Hugging Face: [KS987/PanCancerSeg-Specialized-weights](https://huggingface.co/KS987/PanCancerSeg-Specialized-weights)
+
+```bash
+git lfs install
+git clone https://huggingface.co/KS987/PanCancerSeg-Specialized-weights
+```
+
+## Setup
+
+Create an environment and install the Python dependencies:
+
+```bash
+pip install -r requirements.txt
+```
+
+Download the trained nnUNet model weights to a local directory. Inference resampling can require about 64 GB RAM for large 3D volumes.
+
+Expected model layout:
+
+```text
+nnUNet_results/
+|-- Dataset102_Kidney/
+|   `-- nnUNetTrainerWandb2000__nnUNetResEncUNetMPlans__3d_fullres/
+|       `-- fold_0/
+|           `-- checkpoint_best.pth
+|-- Dataset103_Liver/
+|-- Dataset104_Pancreas/
+`-- Dataset105_Lung/
+```
+
+## Usage
+
+Input images can be named either `case.nii.gz` or `case_0000.nii.gz`; the script handles both.
+
+`--cancer_type` values are `kidney_cancer`, `liver_cancer`, `pancreatic_cancer`, and `lung_cancer`.
+
+Kidney cancer:
+
+```bash
+python predict.py --input /path/to/case.nii.gz --cancer_type kidney_cancer --model_dir /path/to/nnUNet_results --output_dir ./output
+```
+
+Liver cancer:
+
+```bash
+python predict.py --input /path/to/case.nii.gz --cancer_type liver_cancer --model_dir /path/to/nnUNet_results --output_dir ./output
+```
+
+Pancreatic cancer:
+
+```bash
+python predict.py --input /path/to/case.nii.gz --cancer_type pancreatic_cancer --model_dir /path/to/nnUNet_results --output_dir ./output
+```
+
+Lung cancer:
+
+```bash
+python predict.py --input /path/to/case.nii.gz --cancer_type lung_cancer --model_dir /path/to/nnUNet_results --output_dir ./output
+```
+
+Use CPU only when CUDA is unavailable:
+
+```bash
+python predict.py --input /path/to/case.nii.gz --cancer_type kidney_cancer --model_dir /path/to/nnUNet_results --output_dir ./output --device cpu
+```
+
+## Output Files
+
+The output directory contains:
+
+- `{case_id}_seg.nii.gz`: predicted segmentation mask
+- `{case_id}_slice_centroid.png`: centroid slice preview
+- `{case_id}_slice_max_area.png`: max predicted area slice preview
+- `{case_id}_slice_extent25.png`: 25% through predicted z-extent preview
+- `{case_id}_slice_extent75.png`: 75% through predicted z-extent preview
+- `{case_id}_overlay.mp4`: scroll-through overlay video
+
+## Supported Cancer Types
+
+|  `--cancer_type` | Dataset | Window level | Window width |
+|---------------------------|---------|-------------:|-------------:|
+| kidney_cancer | Dataset102_Kidney | 40 | 400 |
+| liver_cancer | Dataset103_Liver | 40 | 400 |
+| pancreatic_cancer | Dataset104_Pancreas | 40 | 400 |
+| lung_cancer | Dataset105_Lung | -600 | 1500 |
+
+## Example Output
+
+The `example/` folder contains sample output from running the kidney cancer model on a validation case (`FLARE23Ts_0005`), including slice PNGs, an overlay video, and the segmentation mask.
+
+## Troubleshooting
+
+CUDA unavailable: run with `--device cpu` or install CUDA-enabled PyTorch.
+
+Missing checkpoint: check that `--model_dir` points to the directory containing the `DatasetXXX_*` model folders and that `fold_0/checkpoint_best.pth` exists.
+
+Missing custom trainer: make sure the cloned repository still has this layout:
+
+```text
+PanCancerSeg-Inference/
+|-- predict.py
+`-- trainers/
+    `-- nnUNetTrainerWandb2000.py
+```
+
+Do not move the `trainers/` directory out of the repository. `predict.py` automatically registers the trainer with nnUNet when inference starts.
+
+MP4 video creation failed: the segmentation mask and PNG files may still have been created. Try running the command on another machine or ask technical support to install video support for OpenCV.

app.py

@@ -0,0 +1,264 @@
+"""Gradio UI for PanCancerSeg single-case CT tumour segmentation."""
+
+import shutil
+import tempfile
+from pathlib import Path
+
+import gradio as gr
+
+from predict import (
+    CANCER_CONFIGS,
+    install_custom_trainer,
+    resolve_case_id,
+    resolve_model_folder,
+    run_nnunet_prediction,
+    summarize_segmentation,
+)
+from visualize import generate_outputs
+
+# ── Constants ──────────────────────────────────────────────────────────────────
+
+CANCER_TYPE_CHOICES = {
+    "Kidney Cancer": "kidney_cancer",
+    "Liver Cancer": "liver_cancer",
+    "Pancreatic Cancer": "pancreatic_cancer",
+    "Lung Cancer": "lung_cancer",
+}
+
+DEFAULT_MODEL_DIR = str(Path(__file__).parent / "PanCancerSeg-Specialized-weights")
+DEFAULT_DEVICE = "cuda"
+
+# Hugging Face Hub repo that hosts the trained nnUNet weights. On Spaces (where the
+# local weights folder is absent) we download them on first use.
+MODEL_REPO_ID = "KS987/PanCancerSeg-Specialized-weights"
+
+
+def resolve_weights_dir() -> Path:
+    """Return a directory containing the DatasetXXX_* model folders.
+
+    Prefer a local checkout (fast local dev); otherwise download the weights
+    from the Hugging Face Hub and cache them.
+    """
+    local_dir = Path(DEFAULT_MODEL_DIR).expanduser().resolve()
+    if local_dir.exists() and any(local_dir.glob("Dataset*")):
+        return local_dir
+
+    from huggingface_hub import snapshot_download
+
+    downloaded = snapshot_download(
+        repo_id=MODEL_REPO_ID,
+        repo_type="model",
+        allow_patterns=["Dataset*/**"],
+    )
+    return Path(downloaded)
+
+_SAMPLE_DIR = Path(__file__).parent / "sample_input"
+_CANCER_TYPE_TO_FOLDER = {
+    "Kidney Cancer": "kidney",
+    "Liver Cancer": "liver",
+    "Pancreatic Cancer": "pancreas",
+    "Lung Cancer": "lung",
+}
+
+def load_example(cancer_type_label: str, index: int) -> str:
+    """Return the index-th (1-based) example _0000.nii.gz for the given cancer type."""
+    folder = _SAMPLE_DIR / _CANCER_TYPE_TO_FOLDER[cancer_type_label]
+    files = sorted(folder.glob("*_0000.nii.gz"))
+    if len(files) < index:
+        raise gr.Error(f"Example {index} not found for {cancer_type_label} in {folder}")
+    return str(files[index - 1])
+
+
+# ── Inference ──────────────────────────────────────────────────────────────────
+
+def run_inference(
+    input_file,
+    cancer_type_label,
+    fps,
+    progress=gr.Progress(track_tqdm=True),
+):
+    import torch
+
+    if input_file is None:
+        raise gr.Error("Please upload a .nii.gz CT image first.")
+
+    input_path = Path(input_file)
+    if not input_path.name.endswith(".nii.gz"):
+        raise gr.Error(f"File must be .nii.gz format. Got: {input_path.name}")
+
+    device = DEFAULT_DEVICE if torch.cuda.is_available() else "cpu"
+
+    progress(0.02, desc="Resolving model weights...")
+    try:
+        model_dir_path = resolve_weights_dir()
+    except Exception as e:
+        raise gr.Error(f"Failed to obtain model weights from '{MODEL_REPO_ID}': {e}")
+
+    cancer_key = CANCER_TYPE_CHOICES[cancer_type_label]
+    config = CANCER_CONFIGS[cancer_key]
+    case_id = resolve_case_id(input_path)
+
+    progress(0.05, desc="Installing custom trainer...")
+    install_custom_trainer()
+
+    progress(0.10, desc="Loading model weights...")
+    model_folder = resolve_model_folder(model_dir_path, config["dataset_name"])
+
+    output_dir = Path(tempfile.mkdtemp(prefix="pancancerseg_out_"))
+
+    try:
+        with tempfile.TemporaryDirectory(prefix="pancancerseg_in_") as tmp:
+            tmp_path = Path(tmp)
+            tmp_input_dir = tmp_path / "input"
+            tmp_output_dir = tmp_path / "prediction"
+            tmp_input_dir.mkdir()
+            tmp_output_dir.mkdir()
+
+            nnunet_input = tmp_input_dir / f"{case_id}_0000.nii.gz"
+            try:
+                nnunet_input.symlink_to(input_path.resolve())
+            except (OSError, NotImplementedError):
+                shutil.copy2(input_path, nnunet_input)
+
+            progress(0.20, desc="Running nnUNet inference (this may take a few minutes)...")
+            run_nnunet_prediction(
+                model_folder=model_folder,
+                input_dir=tmp_input_dir,
+                output_dir=tmp_output_dir,
+                device=device,
+            )
+
+            raw_seg = tmp_output_dir / f"{case_id}.nii.gz"
+            if not raw_seg.exists():
+                produced = [p.name for p in tmp_output_dir.glob("*.nii.gz")]
+                raise RuntimeError(
+                    f"nnUNet did not produce the expected segmentation. Found: {produced}"
+                )
+
+            seg_path = output_dir / f"{case_id}_seg.nii.gz"
+            shutil.copy2(raw_seg, seg_path)
+
+        progress(0.80, desc="Generating slice images and overlay video...")
+        viz = generate_outputs(
+            image_path=input_path,
+            mask_path=seg_path,
+            output_dir=output_dir,
+            case_name=case_id,
+            cancer_type=config["display_name"],
+            wl=config["wl"],
+            ww=config["ww"],
+            color=config["color"],
+            alpha=0.5,
+            fps=int(fps),
+        )
+
+        progress(0.95, desc="Computing tumour volume...")
+        positive_voxels, tumor_volume_ml = summarize_segmentation(seg_path)
+
+        stats = (
+            f"Case ID        : {case_id}\n"
+            f"Cancer type    : {config['display_name']}\n"
+            f"Positive voxels: {positive_voxels:,}\n"
+            f"Tumour volume  : {tumor_volume_ml:.3f} mL"
+        )
+
+        slices = viz["slices"]
+        video_path = viz["video"]
+        video_out = (
+            str(video_path)
+            if video_path.exists() and video_path.stat().st_size > 0
+            else None
+        )
+
+        progress(1.0, desc="Done!")
+        return (
+            stats,
+            str(seg_path),
+            str(slices.get("centroid")),
+            str(slices.get("max_area")),
+            str(slices.get("extent25")),
+            str(slices.get("extent75")),
+            video_out,
+        )
+
+    except Exception as e:
+        shutil.rmtree(output_dir, ignore_errors=True)
+        raise gr.Error(str(e))
+
+
+# ── UI ─────────────────────────────────────────────────────────────────────────
+
+def build_ui():
+    with gr.Blocks(title="PanCancerSeg Inference") as demo:
+        gr.Markdown(
+            """
+# PanCancerSeg — Specialist CT Tumour Segmentation
+Upload a `.nii.gz` CT image, select the cancer type, and click **Run Inference** to obtain
+a segmentation mask and visualisations.
+"""
+        )
+
+        with gr.Row():
+            # ── Left panel: inputs ─────────────────────────────────────────────
+            with gr.Column(scale=1, min_width=300):
+                input_file = gr.File(
+                    label="CT Image (.nii.gz)",
+                    file_types=[".gz"],
+                )
+                cancer_type = gr.Dropdown(
+                    choices=list(CANCER_TYPE_CHOICES.keys()),
+                    value="Kidney Cancer",
+                    label="Cancer Type",
+                )
+                fps = gr.Slider(
+                    minimum=1,
+                    maximum=30,
+                    value=10,
+                    step=1,
+                    label="Video FPS",
+                )
+                with gr.Row():
+                    load_btn_1 = gr.Button("Load Example 1", size="lg")
+                    load_btn_2 = gr.Button("Load Example 2", size="lg")
+                run_btn = gr.Button("Run Inference", variant="primary", size="lg")
+                video_out = gr.Video(label="Overlay Video")
+
+            # ── Right panel: outputs ───────────────────────────────────────────
+            with gr.Column(scale=2):
+                with gr.Row():
+                    stats_box = gr.Textbox(
+                        label="Inference Summary",
+                        lines=4,
+                        interactive=False,
+                    )
+                    seg_file = gr.File(label="Download Segmentation Mask (.nii.gz)")
+                with gr.Row():
+                    img_centroid = gr.Image(label="Centroid Slice", type="filepath")
+                    img_max_area = gr.Image(label="Max Area Slice", type="filepath")
+                with gr.Row():
+                    img_ext25 = gr.Image(label="Extent 25% Slice", type="filepath")
+                    img_ext75 = gr.Image(label="Extent 75% Slice", type="filepath")
+
+        load_btn_1.click(fn=lambda ct: load_example(ct, 1), inputs=[cancer_type], outputs=[input_file])
+        load_btn_2.click(fn=lambda ct: load_example(ct, 2), inputs=[cancer_type], outputs=[input_file])
+
+        run_btn.click(
+            fn=run_inference,
+            inputs=[input_file, cancer_type, fps],
+            outputs=[
+                stats_box,
+                seg_file,
+                img_centroid,
+                img_max_area,
+                img_ext25,
+                img_ext75,
+                video_out,
+            ],
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = build_ui()
+    demo.launch( server_name="0.0.0.0", server_port=7869, share=False, theme=gr.themes.Soft())

example/FLARE23Ts_0005_overlay.mp4

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

example/FLARE23Ts_0005_seg.nii.gz

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

predict.py

@@ -0,0 +1,285 @@
+"""Run single-case PanCancerSeg nnUNet CT inference and visualization."""
+
+import argparse
+import shutil
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+
+from visualize import generate_outputs
+
+
+CANCER_CONFIGS = {
+    "kidney_cancer": {
+        "dataset_id": 102,
+        "dataset_name": "Dataset102_Kidney",
+        "display_name": "Kidney cancer",
+        "wl": 40,
+        "ww": 400,
+        "color": (255, 0, 0),
+    },
+    "liver_cancer": {
+        "dataset_id": 103,
+        "dataset_name": "Dataset103_Liver",
+        "display_name": "Liver cancer",
+        "wl": 40,
+        "ww": 400,
+        "color": (255, 0, 0),
+    },
+    "pancreatic_cancer": {
+        "dataset_id": 104,
+        "dataset_name": "Dataset104_Pancreas",
+        "display_name": "Pancreatic cancer",
+        "wl": 40,
+        "ww": 400,
+        "color": (255, 0, 0),
+    },
+    "lung_cancer": {
+        "dataset_id": 105,
+        "dataset_name": "Dataset105_Lung",
+        "display_name": "Lung cancer",
+        "wl": -600,
+        "ww": 1500,
+        "color": (255, 0, 0),
+    },
+}
+
+CANCER_TYPE_ALIASES = {
+    "kidney": "kidney_cancer",
+    "liver": "liver_cancer",
+    "pancreas": "pancreatic_cancer",
+    "lung": "lung_cancer",
+}
+
+TRAINER_NAME = "nnUNetTrainerWandb2000"
+PLANS_NAME = "nnUNetResEncUNetMPlans"
+CONFIGURATION = "3d_fullres"
+CHECKPOINT_NAME = "checkpoint_best.pth"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Run one PanCancerSeg cancer-specific nnUNet model on a single NIfTI image."
+    )
+    parser.add_argument("--input", required=True, help="Path to a single .nii.gz image")
+    parser.add_argument(
+        "--cancer_type",
+        required=True,
+        help=(
+            "Cancer-specific model to use. "
+            f"Canonical values: {', '.join(sorted(CANCER_CONFIGS))}. "
+            f"Legacy aliases still accepted: {', '.join(sorted(CANCER_TYPE_ALIASES))}."
+        ),
+    )
+    parser.add_argument(
+        "--model_dir",
+        required=True,
+        help="Path to nnUNet results directory containing DatasetXXX_* folders",
+    )
+    parser.add_argument("--output_dir", default="./output", help="Where to save results")
+    parser.add_argument("--fps", type=int, default=10, help="Video frames per second")
+    parser.add_argument("--device", choices=["cuda", "cpu"], default="cuda")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    args.cancer_type = normalize_cancer_type(args.cancer_type)
+    input_path = Path(args.input).expanduser().resolve()
+    model_dir = Path(args.model_dir).expanduser().resolve()
+    output_dir = Path(args.output_dir).expanduser().resolve()
+
+    if not input_path.exists():
+        raise FileNotFoundError(f"Input image not found: {input_path}")
+    if input_path.name.startswith("._") or not input_path.name.endswith(".nii.gz"):
+        raise ValueError(f"Expected a .nii.gz image, got: {input_path.name}")
+    if not model_dir.exists():
+        raise FileNotFoundError(f"Model directory not found: {model_dir}")
+    if args.device == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(
+            "CUDA was requested but torch.cuda.is_available() is False. "
+            "Use --device cpu or install CUDA-enabled PyTorch."
+        )
+    if args.fps <= 0:
+        raise ValueError("--fps must be a positive integer")
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    config = CANCER_CONFIGS[args.cancer_type]
+    case_id = resolve_case_id(input_path)
+
+    install_custom_trainer()
+    model_folder = resolve_model_folder(model_dir, config["dataset_name"])
+
+    with tempfile.TemporaryDirectory(prefix="pancancerseg_") as tmp:
+        tmp_path = Path(tmp)
+        tmp_input_dir = tmp_path / "input"
+        tmp_output_dir = tmp_path / "prediction"
+        tmp_input_dir.mkdir()
+        tmp_output_dir.mkdir()
+
+        nnunet_input = tmp_input_dir / f"{case_id}_0000.nii.gz"
+        symlink_or_copy(input_path, nnunet_input)
+
+        run_nnunet_prediction(
+            model_folder=model_folder,
+            input_dir=tmp_input_dir,
+            output_dir=tmp_output_dir,
+            device=args.device,
+        )
+
+        raw_seg = tmp_output_dir / f"{case_id}.nii.gz"
+        if not raw_seg.exists():
+            produced = sorted(tmp_output_dir.glob("*.nii.gz"))
+            raise FileNotFoundError(
+                f"nnUNet did not write the expected segmentation {raw_seg}. "
+                f"Found: {[p.name for p in produced]}"
+            )
+
+        seg_path = output_dir / f"{case_id}_seg.nii.gz"
+        shutil.copy2(raw_seg, seg_path)
+
+    viz_outputs = generate_outputs(
+        image_path=input_path,
+        mask_path=seg_path,
+        output_dir=output_dir,
+        case_name=case_id,
+        cancer_type=config["display_name"],
+        wl=config["wl"],
+        ww=config["ww"],
+        color=config["color"],
+        alpha=0.5,
+        fps=args.fps,
+    )
+
+    positive_voxels, tumor_volume_ml = summarize_segmentation(seg_path)
+    print_summary(seg_path, viz_outputs, positive_voxels, tumor_volume_ml)
+
+
+def resolve_case_id(input_path):
+    name = input_path.name
+    if not name.endswith(".nii.gz"):
+        raise ValueError(f"Expected a .nii.gz image, got: {name}")
+    case_id = name[: -len(".nii.gz")]
+    if case_id.endswith("_0000"):
+        case_id = case_id[: -len("_0000")]
+    if not case_id:
+        raise ValueError(f"Could not resolve a case ID from: {input_path}")
+    return case_id
+
+
+def normalize_cancer_type(cancer_type):
+    cancer_type = cancer_type.strip().lower()
+    normalized = CANCER_TYPE_ALIASES.get(cancer_type, cancer_type)
+    if normalized not in CANCER_CONFIGS:
+        valid = sorted(list(CANCER_CONFIGS) + list(CANCER_TYPE_ALIASES))
+        raise ValueError(
+            f"Unsupported --cancer_type '{cancer_type}'. Valid values: {', '.join(valid)}"
+        )
+    return normalized
+
+
+def install_custom_trainer():
+    import nnunetv2
+
+    src = Path(__file__).resolve().parent / "trainers" / f"{TRAINER_NAME}.py"
+    if not src.exists():
+        raise FileNotFoundError(f"Custom trainer file is missing: {src}")
+
+    variants_dir = Path(nnunetv2.__path__[0]) / "training" / "nnUNetTrainer" / "variants"
+    variants_dir.mkdir(parents=True, exist_ok=True)
+    dst = variants_dir / src.name
+
+    if dst.exists() or dst.is_symlink():
+        try:
+            if dst.resolve() == src.resolve():
+                return dst
+        except OSError:
+            pass
+        dst.unlink()
+
+    try:
+        dst.symlink_to(src.resolve())
+    except (OSError, NotImplementedError):
+        shutil.copy2(src, dst)
+    print(f"Installed custom trainer: {dst}")
+    return dst
+
+
+def resolve_model_folder(model_dir, dataset_name):
+    model_folder = (
+        model_dir
+        / dataset_name
+        / f"{TRAINER_NAME}__{PLANS_NAME}__{CONFIGURATION}"
+    )
+    checkpoint = model_folder / "fold_0" / CHECKPOINT_NAME
+    if not checkpoint.exists():
+        raise FileNotFoundError(
+            f"Expected checkpoint not found: {checkpoint}. "
+            "Check --model_dir and make sure the trained weights are downloaded."
+        )
+    return model_folder
+
+
+def symlink_or_copy(src, dst):
+    try:
+        dst.symlink_to(src.resolve())
+    except (OSError, NotImplementedError):
+        shutil.copy2(src, dst)
+
+
+def run_nnunet_prediction(model_folder, input_dir, output_dir, device):
+    from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
+
+    predictor = nnUNetPredictor(
+        tile_step_size=0.5,
+        use_gaussian=True,
+        use_mirroring=False,
+        perform_everything_on_device=(device == "cuda"),
+        device=torch.device(device),
+        verbose=False,
+        verbose_preprocessing=False,
+        allow_tqdm=True,
+    )
+    predictor.initialize_from_trained_model_folder(
+        str(model_folder),
+        use_folds=(0,),
+        checkpoint_name=CHECKPOINT_NAME,
+    )
+    predictor.predict_from_files(
+        str(input_dir),
+        str(output_dir),
+        save_probabilities=False,
+        overwrite=True,
+        num_processes_preprocessing=1,
+        num_processes_segmentation_export=1,
+        folder_with_segs_from_prev_stage=None,
+        num_parts=1,
+        part_id=0,
+    )
+
+
+def summarize_segmentation(seg_path):
+    seg = sitk.ReadImage(str(seg_path))
+    seg_arr = sitk.GetArrayFromImage(seg)
+    positive_voxels = int(np.count_nonzero(seg_arr))
+    spacing_x, spacing_y, spacing_z = seg.GetSpacing()
+    tumor_volume_ml = positive_voxels * spacing_x * spacing_y * spacing_z / 1000.0
+    return positive_voxels, tumor_volume_ml
+
+
+def print_summary(seg_path, viz_outputs, positive_voxels, tumor_volume_ml):
+    print("\nPanCancerSeg inference complete")
+    print(f"Segmentation mask : {seg_path}")
+    print("Slice PNGs        :")
+    for label, path in viz_outputs["slices"].items():
+        print(f"  {label:9s} : {path}")
+    print(f"Overlay video     : {viz_outputs['video']}")
+    print(f"Positive voxels   : {positive_voxels}")
+    print(f"Tumor volume      : {tumor_volume_ml:.3f} mL")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,11 @@
+nnunetv2>=2.6
+torch>=2.0
+nibabel>=5.0
+numpy
+opencv-python>=4.8
+matplotlib>=3.7
+scipy
+SimpleITK>=2.3
+wandb
+huggingface_hub>=1.18.0
+gradio>=6.12.0

sample_input/kidney/FLARE_02510.nii.gz

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

sample_input/kidney/FLARE_02510_0000.nii.gz

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

sample_input/kidney/FLARE_10805.nii.gz

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

sample_input/kidney/FLARE_10805_0000.nii.gz

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

sample_input/liver/FLARE_04886.nii.gz

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

sample_input/liver/FLARE_04886_0000.nii.gz

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

sample_input/liver/FLARE_04936.nii.gz

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

sample_input/liver/FLARE_04936_0000.nii.gz

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

sample_input/lung/FLARE_00673.nii.gz

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

sample_input/lung/FLARE_00673_0000.nii.gz

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

sample_input/lung/FLARE_01254.nii.gz

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

sample_input/lung/FLARE_01254_0000.nii.gz

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

sample_input/pancreas/FLARE_04341.nii.gz

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

sample_input/pancreas/FLARE_04341_0000.nii.gz

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

sample_input/pancreas/FLARE_04346.nii.gz

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

sample_input/pancreas/FLARE_04346_0000.nii.gz

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

trainers/nnUNetTrainerWandb2000.py

@@ -0,0 +1,72 @@
+"""nnUNet trainer with Weights & Biases logging and 2000 epochs.
+
+Place this file (or symlink it) into the nnUNet trainer variants directory
+so that nnUNet can discover it via the -tr flag:
+
+    VARIANTS_DIR=$(python -c "import nnunetv2; print(nnunetv2.__path__[0])")/training/nnUNetTrainer/variants
+    ln -sf $(realpath nnUNetTrainerWandb2000.py) "$VARIANTS_DIR/nnUNetTrainerWandb2000.py"
+
+Then train with:
+    nnUNetv2_train DATASET CONFIG FOLD -tr nnUNetTrainerWandb2000 ...
+"""
+
+import os
+
+import torch
+import wandb
+from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer
+
+
+class nnUNetTrainerWandb2000(nnUNetTrainer):
+    def __init__(self, plans: dict, configuration: str, fold: int,
+                 dataset_json: dict,
+                 device: torch.device = torch.device("cuda")):
+        super().__init__(plans, configuration, fold, dataset_json, device)
+        self.num_epochs = 2000
+
+    def on_train_start(self):
+        super().on_train_start()
+        wandb.init(
+            project=os.environ.get("WANDB_PROJECT", "CVPR2026-PanCancerSeg"),
+            name=f"{self.plans_manager.dataset_name}_fold{self.fold}",
+            config={
+                "dataset": self.plans_manager.dataset_name,
+                "configuration": self.configuration_name,
+                "fold": self.fold,
+                "num_epochs": self.num_epochs,
+                "batch_size": self.batch_size,
+                "patch_size": list(self.configuration_manager.patch_size),
+            },
+            resume="allow",
+        )
+
+    def on_epoch_end(self):
+        super().on_epoch_end()
+
+        # Save periodic checkpoint every 200 epochs after epoch 1000
+        if self.current_epoch > 1000 and self.current_epoch % 200 == 0:
+            self.save_checkpoint(
+                os.path.join(self.output_folder, f"checkpoint_epoch{self.current_epoch}.pth")
+            )
+
+        logs = self.logger.my_fantastic_logging
+        metrics = {"epoch": self.current_epoch}
+        if logs["train_losses"]:
+            metrics["train_loss"] = logs["train_losses"][-1]
+        if logs["val_losses"]:
+            metrics["val_loss"] = logs["val_losses"][-1]
+        if logs["ema_fg_dice"]:
+            metrics["ema_fg_dice"] = logs["ema_fg_dice"][-1]
+        if logs["dice_per_class_or_region"]:
+            latest = logs["dice_per_class_or_region"][-1]
+            for i, d in enumerate(latest):
+                metrics[f"dice_class_{i}"] = d
+        metrics["learning_rate"] = self.optimizer.param_groups[0]["lr"]
+        metrics["epoch_time"] = self.logger.my_fantastic_logging.get(
+            "epoch_end_timestamps", [0])[-1] - self.logger.my_fantastic_logging.get(
+            "epoch_start_timestamps", [0])[-1]
+        wandb.log(metrics, step=self.current_epoch)
+
+    def on_train_end(self):
+        super().on_train_end()
+        wandb.finish()

visualize.py

@@ -0,0 +1,287 @@
+"""Visualization helpers for single-case PanCancerSeg inference."""
+
+from pathlib import Path
+
+import cv2
+import numpy as np
+import SimpleITK as sitk
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+
+DEFAULT_OVERLAY_COLOR = (255, 0, 0)
+
+
+def preprocess_volume(volume, wl, ww):
+    """Apply CT windowing and return uint8 data in [0, 255]."""
+    volume = volume.astype(np.float32, copy=False)
+    lower_bound = wl - ww / 2
+    upper_bound = wl + ww / 2
+    clipped = np.clip(volume, lower_bound, upper_bound)
+    return _normalize_to_uint8(clipped)
+
+
+def overlay_mask(gray_slice, mask_slice, color=DEFAULT_OVERLAY_COLOR, alpha=0.5):
+    """Apply a semi-transparent RGB overlay to one grayscale slice."""
+    gray_slice = np.asarray(gray_slice, dtype=np.uint8)
+    if gray_slice.ndim != 2:
+        raise ValueError(f"Expected a 2D grayscale slice, got shape {gray_slice.shape}")
+
+    rgb = np.stack([gray_slice] * 3, axis=-1)
+    mask = mask_slice > 0
+    if not np.any(mask):
+        return rgb
+
+    out = rgb.copy()
+    color_arr = np.asarray(color, dtype=np.float32)
+    blended = out[mask].astype(np.float32) * (1 - alpha) + color_arr * alpha
+    out[mask] = np.clip(blended, 0, 255).astype(np.uint8)
+    return out
+
+
+def find_key_slices(mask_vol):
+    """Return named representative z-slices for a mask in [z, y, x] order."""
+    if mask_vol.ndim != 3:
+        raise ValueError(f"Expected a 3D mask volume, got shape {mask_vol.shape}")
+
+    depth = mask_vol.shape[0]
+    if depth == 0:
+        raise ValueError("Cannot select key slices from an empty z-dimension")
+
+    mask = mask_vol > 0
+    if np.any(mask):
+        z_indices = np.where(np.any(mask, axis=(1, 2)))[0]
+        areas = mask.reshape(depth, -1).sum(axis=1)
+        coords = np.argwhere(mask)
+        centroid_z = int(round(float(coords[:, 0].mean())))
+        min_z = int(z_indices.min())
+        max_z = int(z_indices.max())
+        return {
+            "centroid": _clip_slice(centroid_z, depth),
+            "max_area": int(areas.argmax()),
+            "extent25": _clip_slice(round(min_z + 0.25 * (max_z - min_z)), depth),
+            "extent75": _clip_slice(round(min_z + 0.75 * (max_z - min_z)), depth),
+        }
+
+    middle = depth // 2
+    offset = max(1, depth // 10)
+    return {
+        "centroid": middle,
+        "max_area": _clip_slice(middle - offset, depth),
+        "extent25": _clip_slice(middle + offset, depth),
+        "extent75": _clip_slice(middle + 2 * offset, depth),
+    }
+
+
+def generate_slice_images(
+    image_uint8,
+    mask_vol,
+    output_dir,
+    case_name,
+    color=DEFAULT_OVERLAY_COLOR,
+    alpha=0.5,
+):
+    """Save side-by-side PNGs for representative slices."""
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    key_slices = find_key_slices(mask_vol)
+    outputs = {}
+
+    for label, z_idx in key_slices.items():
+        gray_slice = image_uint8[z_idx]
+        mask_slice = mask_vol[z_idx] > 0
+        overlay = overlay_mask(gray_slice, mask_slice, color=color, alpha=alpha)
+
+        fig, axes = plt.subplots(1, 2, figsize=(10, 5), dpi=150)
+        axes[0].imshow(gray_slice, cmap="gray", vmin=0, vmax=255)
+        axes[0].set_title("Image")
+        axes[0].axis("off")
+        axes[1].imshow(overlay)
+        axes[1].set_title("Segmentation overlay")
+        axes[1].axis("off")
+        fig.suptitle(f"{case_name} | z={z_idx}")
+        fig.tight_layout()
+
+        out_path = output_dir / f"{case_name}_slice_{label}.png"
+        fig.savefig(out_path, dpi=150, bbox_inches="tight", facecolor="white")
+        plt.close(fig)
+        outputs[label] = out_path
+
+    return outputs
+
+
+def generate_video(
+    image_uint8,
+    mask_vol,
+    output_dir,
+    case_name,
+    cancer_type,
+    color=DEFAULT_OVERLAY_COLOR,
+    alpha=0.5,
+    fps=10,
+):
+    """Generate an MP4 scroll-through overlay video."""
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    video_path = output_dir / f"{case_name}_overlay.mp4"
+
+    start_z, end_z = _video_z_range(mask_vol)
+    first_frame = _make_video_frame(
+        image_uint8[start_z],
+        mask_vol[start_z],
+        color,
+        alpha,
+        start_z,
+        image_uint8.shape[0],
+        cancer_type,
+    )
+    height, width = first_frame.shape[:2]
+
+    writer = _open_video_writer(video_path, fps, width, height)
+    # Frame annotations are drawn in RGB space; convert only when writing to OpenCV.
+    writer.write(cv2.cvtColor(first_frame, cv2.COLOR_RGB2BGR))
+
+    for z_idx in range(start_z + 1, end_z + 1):
+        frame = _make_video_frame(
+            image_uint8[z_idx],
+            mask_vol[z_idx],
+            color,
+            alpha,
+            z_idx,
+            image_uint8.shape[0],
+            cancer_type,
+        )
+        writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+
+    writer.release()
+    return video_path
+
+
+def generate_outputs(
+    image_path,
+    mask_path,
+    output_dir,
+    case_name,
+    cancer_type,
+    wl,
+    ww,
+    color=DEFAULT_OVERLAY_COLOR,
+    alpha=0.5,
+    fps=10,
+):
+    """Read image and mask volumes, then write PNG previews and MP4 video."""
+    image = sitk.ReadImage(str(image_path))
+    mask = sitk.ReadImage(str(mask_path))
+    image_vol = sitk.GetArrayFromImage(image)
+    mask_vol = sitk.GetArrayFromImage(mask)
+
+    if image_vol.shape != mask_vol.shape:
+        raise ValueError(
+            "Image and segmentation shapes do not match: "
+            f"image={image_vol.shape}, segmentation={mask_vol.shape}. "
+            "Both arrays are expected in [z, y, x] order."
+        )
+
+    image_uint8 = preprocess_volume(image_vol, wl, ww)
+    slice_paths = generate_slice_images(
+        image_uint8,
+        mask_vol,
+        output_dir,
+        case_name,
+        color,
+        alpha,
+    )
+    video_path = generate_video(
+        image_uint8,
+        mask_vol,
+        output_dir,
+        case_name,
+        cancer_type,
+        color,
+        alpha,
+        fps,
+    )
+    return {"slices": slice_paths, "video": video_path}
+
+
+def _normalize_to_uint8(volume):
+    v_min = float(np.min(volume))
+    v_max = float(np.max(volume))
+    if not np.isfinite(v_min) or not np.isfinite(v_max) or v_max <= v_min:
+        return np.zeros(volume.shape, dtype=np.uint8)
+    normalized = (volume - v_min) / (v_max - v_min) * 255.0
+    return np.clip(normalized, 0, 255).astype(np.uint8)
+
+
+def _clip_slice(index, depth):
+    return int(np.clip(index, 0, depth - 1))
+
+
+def _video_z_range(mask_vol, padding=10, empty_window=80):
+    depth = mask_vol.shape[0]
+    mask = mask_vol > 0
+    if np.any(mask):
+        z_indices = np.where(np.any(mask, axis=(1, 2)))[0]
+        return (
+            max(0, int(z_indices.min()) - padding),
+            min(depth - 1, int(z_indices.max()) + padding),
+        )
+
+    if depth <= empty_window:
+        return 0, depth - 1
+    middle = depth // 2
+    half = empty_window // 2
+    return max(0, middle - half), min(depth - 1, middle + half)
+
+
+def _make_video_frame(gray_slice, mask_slice, color, alpha, z_idx, depth, cancer_type):
+    frame = overlay_mask(gray_slice, mask_slice, color=color, alpha=alpha)
+    frame = _upscale_if_small(frame)
+
+    annotation = f"Slice {z_idx + 1}/{depth} | {cancer_type}"
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    font_scale = max(0.6, min(frame.shape[:2]) / 900)
+    thickness = max(1, int(round(font_scale * 2)))
+    text_size, baseline = cv2.getTextSize(annotation, font, font_scale, thickness)
+    x, y = 12, 12 + text_size[1]
+    cv2.rectangle(
+        frame,
+        (x - 6, y - text_size[1] - 6),
+        (x + text_size[0] + 6, y + baseline + 6),
+        (0, 0, 0),
+        thickness=-1,
+    )
+    cv2.putText(frame, annotation, (x, y), font, font_scale, (255, 255, 255), thickness, cv2.LINE_AA)
+    return frame
+
+
+def _upscale_if_small(frame, min_short_side=512):
+    height, width = frame.shape[:2]
+    short_side = min(height, width)
+    if short_side >= min_short_side:
+        return frame
+    scale = min_short_side / short_side
+    new_size = (int(round(width * scale)), int(round(height * scale)))
+    return cv2.resize(frame, new_size, interpolation=cv2.INTER_LINEAR)
+
+
+def _open_video_writer(video_path, fps, width, height):
+    attempts = [
+        ("avc1", "H.264/avc1"),
+        ("mp4v", "MPEG-4/mp4v"),
+    ]
+    for fourcc_text, label in attempts:
+        fourcc = cv2.VideoWriter_fourcc(*fourcc_text)
+        writer = cv2.VideoWriter(str(video_path), fourcc, float(fps), (width, height))
+        if writer.isOpened():
+            return writer
+        writer.release()
+    raise RuntimeError(
+        f"Could not open MP4 writer at {video_path}. Tried "
+        + ", ".join(label for _, label in attempts)
+        + ". Install an OpenCV build with MP4 codec support or try another machine."
+    )