Single-cancer Space: Liver (Dataset103) with bundled weights

PanCancerSeg-Specialized-weights/Dataset103_Liver/nnUNetTrainerWandb2000__nnUNetResEncUNetMPlans__3d_fullres/dataset.json

@@ -0,0 +1,11 @@
+{
+    "channel_names": {
+        "0": "CT"
+    },
+    "labels": {
+        "background": 0,
+        "tumor": 1
+    },
+    "numTraining": 866,
+    "file_ending": ".nii.gz"
+}

PanCancerSeg-Specialized-weights/Dataset103_Liver/nnUNetTrainerWandb2000__nnUNetResEncUNetMPlans__3d_fullres/fold_0/checkpoint_best.pth

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

PanCancerSeg-Specialized-weights/Dataset103_Liver/nnUNetTrainerWandb2000__nnUNetResEncUNetMPlans__3d_fullres/plans.json

@@ -0,0 +1,532 @@
+{
+    "dataset_name": "Dataset103_Liver",
+    "plans_name": "nnUNetResEncUNetMPlans",
+    "original_median_spacing_after_transp": [
+        2.5,
+        0.775390625,
+        0.775390625
+    ],
+    "original_median_shape_after_transp": [
+        102,
+        512,
+        512
+    ],
+    "image_reader_writer": "SimpleITKIO",
+    "transpose_forward": [
+        0,
+        1,
+        2
+    ],
+    "transpose_backward": [
+        0,
+        1,
+        2
+    ],
+    "configurations": {
+        "2d": {
+            "data_identifier": "nnUNetPlans_2d",
+            "preprocessor_name": "DefaultPreprocessor",
+            "batch_size": 13,
+            "patch_size": [
+                512,
+                512
+            ],
+            "median_image_size_in_voxels": [
+                512.0,
+                512.0
+            ],
+            "spacing": [
+                0.775390625,
+                0.775390625
+            ],
+            "normalization_schemes": [
+                "CTNormalization"
+            ],
+            "use_mask_for_norm": [
+                false
+            ],
+            "resampling_fn_data": "resample_data_or_seg_to_shape",
+            "resampling_fn_seg": "resample_data_or_seg_to_shape",
+            "resampling_fn_data_kwargs": {
+                "is_seg": false,
+                "order": 3,
+                "order_z": 0,
+                "force_separate_z": null
+            },
+            "resampling_fn_seg_kwargs": {
+                "is_seg": true,
+                "order": 1,
+                "order_z": 0,
+                "force_separate_z": null
+            },
+            "resampling_fn_probabilities": "resample_data_or_seg_to_shape",
+            "resampling_fn_probabilities_kwargs": {
+                "is_seg": false,
+                "order": 1,
+                "order_z": 0,
+                "force_separate_z": null
+            },
+            "architecture": {
+                "network_class_name": "dynamic_network_architectures.architectures.unet.ResidualEncoderUNet",
+                "arch_kwargs": {
+                    "n_stages": 8,
+                    "features_per_stage": [
+                        32,
+                        64,
+                        128,
+                        256,
+                        512,
+                        512,
+                        512,
+                        512
+                    ],
+                    "conv_op": "torch.nn.modules.conv.Conv2d",
+                    "kernel_sizes": [
+                        [
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3
+                        ]
+                    ],
+                    "strides": [
+                        [
+                            1,
+                            1
+                        ],
+                        [
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2
+                        ]
+                    ],
+                    "n_blocks_per_stage": [
+                        1,
+                        3,
+                        4,
+                        6,
+                        6,
+                        6,
+                        6,
+                        6
+                    ],
+                    "n_conv_per_stage_decoder": [
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1
+                    ],
+                    "conv_bias": true,
+                    "norm_op": "torch.nn.modules.instancenorm.InstanceNorm2d",
+                    "norm_op_kwargs": {
+                        "eps": 1e-05,
+                        "affine": true
+                    },
+                    "dropout_op": null,
+                    "dropout_op_kwargs": null,
+                    "nonlin": "torch.nn.LeakyReLU",
+                    "nonlin_kwargs": {
+                        "inplace": true
+                    }
+                },
+                "_kw_requires_import": [
+                    "conv_op",
+                    "norm_op",
+                    "dropout_op",
+                    "nonlin"
+                ]
+            },
+            "batch_dice": true
+        },
+        "3d_lowres": {
+            "data_identifier": "nnUNetResEncUNetMPlans_3d_lowres",
+            "preprocessor_name": "DefaultPreprocessor",
+            "batch_size": 2,
+            "patch_size": [
+                96,
+                160,
+                160
+            ],
+            "median_image_size_in_voxels": [
+                135,
+                267,
+                267
+            ],
+            "spacing": [
+                1.9161034088607827,
+                1.4857286197611925,
+                1.4857286197611925
+            ],
+            "normalization_schemes": [
+                "CTNormalization"
+            ],
+            "use_mask_for_norm": [
+                false
+            ],
+            "resampling_fn_data": "resample_data_or_seg_to_shape",
+            "resampling_fn_seg": "resample_data_or_seg_to_shape",
+            "resampling_fn_data_kwargs": {
+                "is_seg": false,
+                "order": 3,
+                "order_z": 0,
+                "force_separate_z": null
+            },
+            "resampling_fn_seg_kwargs": {
+                "is_seg": true,
+                "order": 1,
+                "order_z": 0,
+                "force_separate_z": null
+            },
+            "resampling_fn_probabilities": "resample_data_or_seg_to_shape",
+            "resampling_fn_probabilities_kwargs": {
+                "is_seg": false,
+                "order": 1,
+                "order_z": 0,
+                "force_separate_z": null
+            },
+            "architecture": {
+                "network_class_name": "dynamic_network_architectures.architectures.unet.ResidualEncoderUNet",
+                "arch_kwargs": {
+                    "n_stages": 6,
+                    "features_per_stage": [
+                        32,
+                        64,
+                        128,
+                        256,
+                        320,
+                        320
+                    ],
+                    "conv_op": "torch.nn.modules.conv.Conv3d",
+                    "kernel_sizes": [
+                        [
+                            3,
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3,
+                            3
+                        ]
+                    ],
+                    "strides": [
+                        [
+                            1,
+                            1,
+                            1
+                        ],
+                        [
+                            2,
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2,
+                            2
+                        ],
+                        [
+                            1,
+                            2,
+                            2
+                        ]
+                    ],
+                    "n_blocks_per_stage": [
+                        1,
+                        3,
+                        4,
+                        6,
+                        6,
+                        6
+                    ],
+                    "n_conv_per_stage_decoder": [
+                        1,
+                        1,
+                        1,
+                        1,
+                        1
+                    ],
+                    "conv_bias": true,
+                    "norm_op": "torch.nn.modules.instancenorm.InstanceNorm3d",
+                    "norm_op_kwargs": {
+                        "eps": 1e-05,
+                        "affine": true
+                    },
+                    "dropout_op": null,
+                    "dropout_op_kwargs": null,
+                    "nonlin": "torch.nn.LeakyReLU",
+                    "nonlin_kwargs": {
+                        "inplace": true
+                    }
+                },
+                "_kw_requires_import": [
+                    "conv_op",
+                    "norm_op",
+                    "dropout_op",
+                    "nonlin"
+                ]
+            },
+            "batch_dice": false,
+            "next_stage": "3d_cascade_fullres"
+        },
+        "3d_fullres": {
+            "data_identifier": "nnUNetPlans_3d_fullres",
+            "preprocessor_name": "DefaultPreprocessor",
+            "batch_size": 2,
+            "patch_size": [
+                96,
+                160,
+                160
+            ],
+            "median_image_size_in_voxels": [
+                258.0,
+                512.0,
+                512.0
+            ],
+            "spacing": [
+                1.0,
+                0.775390625,
+                0.775390625
+            ],
+            "normalization_schemes": [
+                "CTNormalization"
+            ],
+            "use_mask_for_norm": [
+                false
+            ],
+            "resampling_fn_data": "resample_data_or_seg_to_shape",
+            "resampling_fn_seg": "resample_data_or_seg_to_shape",
+            "resampling_fn_data_kwargs": {
+                "is_seg": false,
+                "order": 3,
+                "order_z": 0,
+                "force_separate_z": null
+            },
+            "resampling_fn_seg_kwargs": {
+                "is_seg": true,
+                "order": 1,
+                "order_z": 0,
+                "force_separate_z": null
+            },
+            "resampling_fn_probabilities": "resample_data_or_seg_to_shape",
+            "resampling_fn_probabilities_kwargs": {
+                "is_seg": false,
+                "order": 1,
+                "order_z": 0,
+                "force_separate_z": null
+            },
+            "architecture": {
+                "network_class_name": "dynamic_network_architectures.architectures.unet.ResidualEncoderUNet",
+                "arch_kwargs": {
+                    "n_stages": 6,
+                    "features_per_stage": [
+                        32,
+                        64,
+                        128,
+                        256,
+                        320,
+                        320
+                    ],
+                    "conv_op": "torch.nn.modules.conv.Conv3d",
+                    "kernel_sizes": [
+                        [
+                            3,
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3,
+                            3
+                        ],
+                        [
+                            3,
+                            3,
+                            3
+                        ]
+                    ],
+                    "strides": [
+                        [
+                            1,
+                            1,
+                            1
+                        ],
+                        [
+                            2,
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2,
+                            2
+                        ],
+                        [
+                            2,
+                            2,
+                            2
+                        ],
+                        [
+                            1,
+                            2,
+                            2
+                        ]
+                    ],
+                    "n_blocks_per_stage": [
+                        1,
+                        3,
+                        4,
+                        6,
+                        6,
+                        6
+                    ],
+                    "n_conv_per_stage_decoder": [
+                        1,
+                        1,
+                        1,
+                        1,
+                        1
+                    ],
+                    "conv_bias": true,
+                    "norm_op": "torch.nn.modules.instancenorm.InstanceNorm3d",
+                    "norm_op_kwargs": {
+                        "eps": 1e-05,
+                        "affine": true
+                    },
+                    "dropout_op": null,
+                    "dropout_op_kwargs": null,
+                    "nonlin": "torch.nn.LeakyReLU",
+                    "nonlin_kwargs": {
+                        "inplace": true
+                    }
+                },
+                "_kw_requires_import": [
+                    "conv_op",
+                    "norm_op",
+                    "dropout_op",
+                    "nonlin"
+                ]
+            },
+            "batch_dice": true
+        },
+        "3d_cascade_fullres": {
+            "inherits_from": "3d_fullres",
+            "previous_stage": "3d_lowres"
+        }
+    },
+    "experiment_planner_used": "nnUNetPlannerResEncM",
+    "label_manager": "LabelManager",
+    "foreground_intensity_properties_per_channel": {
+        "0": {
+            "max": 3525.0,
+            "mean": 77.27534484863281,
+            "median": 77.0,
+            "min": -1024.0,
+            "percentile_00_5": -40.0,
+            "percentile_99_5": 188.0,
+            "std": 47.19807434082031
+        }
+    }
+}

README.md

@@ -1,13 +1,131 @@
 ---
-title: PanCancerSeg Liver
-emoji: 📊
-colorFrom: green
+title: PanCancerSeg Specialist Inference
+emoji: 🩻
+colorFrom: blue
 colorTo: indigo
 sdk: gradio
-sdk_version: 6.16.0
-python_version: '3.12'
+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,377 @@
+"""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_single,
+    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"
+
+# Resolved once per process; subsequent inferences reuse it (no re-download).
+_WEIGHTS_DIR: Path | None = None
+
+
+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 once and cache the resolved path in-process so we
+    never hit the Hub again on later inferences.
+    """
+    global _WEIGHTS_DIR
+    if _WEIGHTS_DIR is not None:
+        return _WEIGHTS_DIR
+
+    local_dir = Path(DEFAULT_MODEL_DIR).expanduser().resolve()
+    if local_dir.exists() and any(local_dir.glob("Dataset*")):
+        _WEIGHTS_DIR = local_dir
+        return _WEIGHTS_DIR
+
+    from huggingface_hub import snapshot_download
+
+    downloaded = snapshot_download(
+        repo_id=MODEL_REPO_ID,
+        repo_type="model",
+        allow_patterns=["Dataset*/**"],
+    )
+    _WEIGHTS_DIR = Path(downloaded)
+    return _WEIGHTS_DIR
+
+
+# ── ZeroGPU support ──────────────────────────────────────────────────────────
+# On Hugging Face ZeroGPU Spaces the `spaces` package is available, and any GPU
+# work must run inside a function decorated with `@spaces.GPU`. Locally (or on a
+# dedicated GPU Space) the package is absent, so we fall back to a no-op so the
+# same code keeps working everywhere.
+try:
+    import spaces  # type: ignore
+
+    _HAS_ZEROGPU = True
+except ImportError:
+    spaces = None
+    _HAS_ZEROGPU = False
+
+
+def gpu_task(duration: int = 180):
+    if _HAS_ZEROGPU:
+        return spaces.GPU(duration=duration)
+
+    def _identity(fn):
+        return fn
+
+    return _identity
+
+
+@gpu_task(duration=180)
+def run_gpu_segmentation(model_folder_str: str, input_file_str: str, output_file_str: str) -> None:
+    """Run nnUNet inference on GPU. Executed inside the ZeroGPU worker process.
+
+    Uses the single-case, no-multiprocessing path because ZeroGPU runs this in a
+    daemon process that is not allowed to spawn child processes.
+    """
+    # The custom trainer must be registered inside the GPU worker process so that
+    # nnUNet can discover it when initialising from the trained model folder.
+    install_custom_trainer()
+    run_nnunet_prediction_single(
+        model_folder=model_folder_str,
+        input_file=input_file_str,
+        output_file=output_file_str,
+        device="cuda",
+    )
+
+
+_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])
+
+
+def count_examples(cancer_type_label: str) -> int:
+    """Number of bundled example CT volumes for a cancer type."""
+    folder = _SAMPLE_DIR / _CANCER_TYPE_TO_FOLDER[cancer_type_label]
+    if not folder.exists():
+        return 0
+    return len(sorted(folder.glob("*_0000.nii.gz")))
+
+
+def available_cancer_labels(weights_dir) -> list:
+    """Cancer labels whose DatasetXXX folder is present under ``weights_dir``.
+
+    A single-cancer Space bundles exactly one DatasetXXX folder, so this returns
+    a single label and the UI locks to it. A full checkout with all four datasets
+    returns every label and the UI shows the selector.
+    """
+    weights_dir = Path(weights_dir)
+    found = [
+        label
+        for label, key in CANCER_TYPE_CHOICES.items()
+        if (weights_dir / CANCER_CONFIGS[key]["dataset_name"]).exists()
+    ]
+    return found or list(CANCER_TYPE_CHOICES.keys())
+
+
+# ── Inference ──────────────────────────────────────────────────────────────────
+
+def run_inference(
+    input_file,
+    cancer_type_label,
+    fps,
+    progress=gr.Progress(track_tqdm=True),
+):
+    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}")
+
+    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.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)
+
+            raw_seg = tmp_output_dir / f"{case_id}.nii.gz"
+
+            progress(0.20, desc="Running nnUNet inference on GPU (this may take a few minutes)...")
+            run_gpu_segmentation(
+                str(model_folder),
+                str(nnunet_input),
+                str(raw_seg),
+            )
+
+            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(available_labels=None):
+    labels = available_labels or list(CANCER_TYPE_CHOICES.keys())
+    single = len(labels) == 1
+    default_label = labels[0]
+
+    if single:
+        title = f"# PanCancerSeg — {default_label} CT Segmentation"
+        intro = (
+            f"Upload a `.nii.gz` CT image and click **Run Inference** to segment "
+            f"**{default_label.lower()}** and obtain a mask plus visualisations."
+        )
+    else:
+        title = "# PanCancerSeg — Specialist CT Tumour Segmentation"
+        intro = (
+            "Upload a `.nii.gz` CT image, select the cancer type, and click "
+            "**Run Inference** to obtain a segmentation mask and visualisations."
+        )
+
+    n_examples = count_examples(default_label) if single else 2
+
+    with gr.Blocks(title="PanCancerSeg Inference") as demo:
+        gr.Markdown(f"{title}\n{intro}")
+
+        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=labels,
+                    value=default_label,
+                    label="Cancer Type",
+                    interactive=not single,
+                )
+                fps = gr.Slider(
+                    minimum=1,
+                    maximum=30,
+                    value=10,
+                    step=1,
+                    label="Video FPS",
+                )
+                example_buttons = []
+                if n_examples > 0:
+                    with gr.Row():
+                        for i in range(1, n_examples + 1):
+                            label = "Load Example" if n_examples == 1 else f"Load Example {i}"
+                            example_buttons.append(gr.Button(label, 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")
+
+        for idx, btn in enumerate(example_buttons, start=1):
+            btn.click(
+                fn=(lambda i: lambda ct: load_example(ct, i))(idx),
+                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__":
+    import os
+
+    # Warm the weights cache at startup so the very first inference (and every
+    # later one) does not trigger a download. Failures are non-fatal: we fall
+    # back to lazy download on the first request.
+    labels = None
+    try:
+        weights_dir = resolve_weights_dir()
+        labels = available_cancer_labels(weights_dir)
+        print(f"[startup] available cancer models: {labels}")
+    except Exception as e:
+        print(f"[startup] weight pre-fetch skipped: {e}")
+
+    demo = build_ui(labels)
+    # Hugging Face Spaces expect the app on port 7860 (set via GRADIO_SERVER_PORT).
+    # Locally this falls back to 7860 unless overridden.
+    port = int(os.environ.get("GRADIO_SERVER_PORT", 7860))
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=port,
+        share=False,
+        theme=gr.themes.Soft(),
+        ssr_mode=False,
+    )

predict.py

@@ -0,0 +1,318 @@
+"""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 run_nnunet_prediction_single(model_folder, input_file, output_file, device):
+    """Single-case nnUNet inference without any multiprocessing.
+
+    Uses nnUNet's in-process API so it works inside restricted environments
+    (e.g. ZeroGPU daemon workers) where spawning child processes is forbidden.
+    """
+    from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
+    from nnunetv2.imageio.simpleitk_reader_writer import SimpleITKIO
+
+    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,
+    )
+
+    io = SimpleITKIO()
+    images, properties = io.read_images([str(input_file)])
+    segmentation = predictor.predict_single_npy_array(
+        images, properties, None, None, False
+    )
+    io.write_seg(segmentation, str(output_file), properties)
+
+
+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,12 @@
+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
+spaces
+gradio>=6.12.0

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

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."
+    )