3D nnU-Net: Osteosarcoma Lesion Segmentation on CE-MRI

This repository provides the trained model weights and configuration files for an automated 3D full-resolution nnU-Net (v2) model specialized for segmenting primary extremity osteosarcoma lesions on preoperative contrast-enhanced MRI (T1W_TSE_FS+C).

🔬 Model & Cohort Overview

Task: Automated 3D volumetric semantic segmentation of primary osteosarcoma tumor core.
Framework: nnU-Net v2 (3d_fullres configuration, PlainConvUNet 6-stage 3D architecture).
Modality: Preoperative axial T1-weighted fat-suppressed contrast-enhanced MRI (T1W_TSE_FS+C).
Labels:
- 0: Background / normal tissue
- 1: Osteosarcoma lesion (intratumoral core)
Cohort Composition ($N = 355$ total patients across two centers):
- Training Cohort ($n = 168$): Manually delineated on axial T1W_TSE_FS+C scans and randomly split 8:2 into training ($n = 134$) and validation ($n = 34$) subsets to train the 3D nnU-Net using a 2,000-epoch schedule (achieving a tuning DSC of $0.907$; best validation EMA Dice of $0.9229$, checkpointed at epoch 1,851).
- Internal Test Cohort ($n = 112$): Automated segmentations were quantitatively evaluated against senior radiologist manual reference standards on a randomly selected benchmark subset of $N = 70$ cases (table below).
- Temporal Validation Cohort ($n = 40$) & External Validation Cohort ($n = 35$): Segmented end-to-end by this locked nnU-Net model for downstream prognostic validation.

📊 Segmentation Performance

Validated against expert radiologist manual delineations on the held-out evaluation benchmark ($N = 70$):

Metric	Mean ± SD
Dice Similarity Coefficient (DSC)	0.893 ± 0.127
Jaccard Index (IoU)	0.827 ± 0.178
Sensitivity (Recall)	0.933 ± 0.099
Positive Predictive Value (PPV)	0.881 ± 0.165
Absolute Volume Difference	32,614 ± 56,570 mm³

📦 Repository Structure & Preprocessing Note

checkpoint_best.pth: Trained PyTorch checkpoint containing network weights (network_weights), best EMA score (_best_ema = 0.9229), and embedded init_args.
plans.json: nnU-Net v2 architecture and preprocessing configuration (target spacing [3.0, 0.758, 0.758] mm, patch size [32, 192, 288], and foreground intensity statistics).
dataset.json: Dataset metadata defining the CE-MRI (T1W_TSE_FS+C) input channel and label map.

Important Note on Intensity Normalization (plans.json): During initial nnU-Net experiment planning and throughout the 2,000-epoch training, intensity normalization was configured using foreground percentile clipping (0.5%–99.5% intensity range: [59.0, 1637.0]) followed by global foreground standardization (mean = 612.63, std = 317.29), which corresponds to the "CTNormalization" class entry in plans.json. Do not modify "normalization_schemes" in plans.json, as nnUNetv2_predict relies on this exact locked scheme to match training-time preprocessing.

🚀 Quick Start / Inference

⚠️ Hardware Requirement (GPU Required)

A CUDA-enabled NVIDIA GPU is required for standard CLI inference.
nnUNetv2_predict defaults to -device cuda and enables CUDA pin_memory in its data loader. Running nnUNetv2_predict on a CPU-only machine without a GPU may trigger an upstream PyTorch pin_memory runtime error.

1. Installation

pip install torch torchvision nnunetv2

2. Prediction via `nnU-Net v2` CLI

Option A: Linux / macOS (Bash)

# 1. Set nnUNet_results environment variable and create directory structure
export nnUNet_results="./nnUNet_results"
MODEL_DIR="${nnUNet_results}/Dataset001_Osteosarcoma_CEMRI/nnUNetTrainer__nnUNetPlans__3d_fullres/fold_0"
mkdir -p "${MODEL_DIR}"

# 2. Copy repository files into the expected hierarchy
cp checkpoint_best.pth "${MODEL_DIR}/"
cp dataset.json "${MODEL_DIR}/../"
cp plans.json "${MODEL_DIR}/../"

# 3. Run 3D full-resolution inference (input MRI files must end with _0000.nii.gz)
nnUNetv2_predict -i /path/to/input_mri_dir \
                 -o /path/to/output_predictions \
                 -d 001 \
                 -c 3d_fullres \
                 -f 0 \
                 -chk checkpoint_best.pth \
                 -device cuda

Option B: Windows (PowerShell)

# 1. Set nnUNet_results environment variable and create directory structure
$env:nnUNet_results = "$PWD\nnUNet_results"
$MODEL_DIR = "$env:nnUNet_results\Dataset001_Osteosarcoma_CEMRI\nnUNetTrainer__nnUNetPlans__3d_fullres\fold_0"
New-Item -ItemType Directory -Force -Path $MODEL_DIR

# 2. Copy repository files into the expected hierarchy
Copy-Item checkpoint_best.pth "$MODEL_DIR\"
Copy-Item dataset.json "$MODEL_DIR\..\"
Copy-Item plans.json "$MODEL_DIR\..\"

# 3. Run 3D full-resolution inference (input MRI files must end with _0000.nii.gz)
nnUNetv2_predict -i C:\path\to\input_mri_dir `
                 -o C:\path\to\output_predictions `
                 -d 001 `
                 -c 3d_fullres `
                 -f 0 `
                 -chk checkpoint_best.pth `
                 -device cuda

(For Windows CMD, use set nnUNet_results=%cd%\nnUNet_results prior to running nnUNetv2_predict.)

3. Inspecting the Checkpoint in Python

import torch

checkpoint = torch.load("checkpoint_best.pth", map_location="cpu", weights_only=False)

print("Trainer:", checkpoint.get("trainer_name"))
print("Configuration:", checkpoint["init_args"].get("configuration"))
print("Trained Epochs:", checkpoint.get("current_epoch"))
print("Best Validation EMA Dice:", checkpoint.get("_best_ema"))

⚕️ Intended Use & Limitations

Intended Use: Research use only. This model is intended for automated tumor segmentation in retrospective osteosarcoma imaging research (e.g., lesion quantification and extraction of radiomic/prognostic features). It is not a diagnostic or treatment-planning device and must not be used for clinical decision-making without qualified radiologist review.
Expected Input: Preoperative axial T1W_TSE_FS+C MRI. Performance may degrade on other pulse sequences, scanners, or acquisition protocols not represented in the development cohort ($N = 355$, two centers).
No Patient Data: This repository distributes model weights and configuration files only — no imaging or clinical data are included. The underlying data were not distributed with this repository due to patient privacy considerations.

📖 Citation

If you use this model in your research, please cite:

@article{jin2026osteosarcoma,
  title={Prognostic Value of Peritumoral and Intratumoral Features on Contrast-Enhanced MRI for Early Recurrence After Limb-Salvage Surgery in Osteosarcoma},
  author={Jin, Xuan-chen and Jiang, Meng-dan and Zhang, Rui-peng and Li, Yue-Hua},
  year={2026}
}

📜 License

This model and its configuration are released under the Apache 2.0 License.

Downloads last month: -; Downloads are not tracked for this model. How to track