Choroid Plexus Segmentation

PyTorch/MONAI pipeline for binary choroid plexus (ChP) segmentation from T1-weighted MRI. It includes ventricle-centered preprocessing, patch-based UXNet training, multi-rater label handling, evaluation, and single-model or ensemble prediction.

Setup

In a Python virtual environment, install the recorded CUDA 12.1 dependencies:

python -m pip install -r requirements-lock.txt --extra-index-url https://download.pytorch.org/whl/cu121

A CUDA GPU is recommended for training; CPU execution is supported. Paths and experiment settings are configured inside the scripts.

MP2RAGE augmentation: the engine is bundled in mp2rage_3d.py and uses NumPy; no external augmentation package is needed for it. It is enabled by Config.mp2rage_engine_enabled=True. Set the flag to False to disable it. The separate T2 engine remains optional and requires the external augmentation package.

Usage

1. Prepare data. Supply canonical-RAS, ventricle-cropped NIfTI images and matching categorical ChP masks on the same voxel grid. For UXNet, each dimension must be at least 96 voxels and divisible by 16. Use Config.dataset_manifest with dataset_manifest.schema.json, or paired files named <case>_t1_cropped.nii.gz and <case>_chp_cropped.nii.gz in mri_dir and seg_dir. For raw-data preprocessing, configure AugmentChoroidScans in augment_scans.py first: its checked-in cohort list is empty.
2. Train. Edit Config in train_choroid.py, especially input/output paths, annotation settings, and fold_index. Run python train_choroid.py. Defaults use five participant-grouped, cohort-stratified development folds and a separate 20% test holdout; run fold indices 0–4 with the same split_seed.
3. Evaluate. Set Evaluator.run_dirs and a new, nonexistent Evaluator.out_dir in evaluate.py. Run python evaluate.py --evaluation-scope val to select inference settings, then use a fresh output directory and --evaluation-scope test for the frozen test evaluation. Metrics include Dice, surface distances, and multi-rater agreement when available.
4. Predict. Set MODELS (one or more best_model.pth paths), INPUT (a ventricle crop), and OUTPUT in predict_chp.py, then run python predict_chp.py. It saves a segmentation, probability/uncertainty maps, and an inference-settings JSON sidecar.

Training saves config.json, splits.json, training.log, and best/final checkpoints under checkpoint_dir. Checkpoints have SHA-256 sidecars. Tune ensembles, thresholds, and test-time augmentation on validation data only.

See PIPELINE.md and CHEATSHEET.md for the dataset, annotation, and evaluation workflow.

Differences from the original UXNet paper

1. Task: This project segments ChP versus background from T1 ventricle crops; the original paper evaluates multi-class brain-tissue and abdominal-organ segmentation.
2. Loss and labels: A hybrid loss function of dice plus BCE was used, as opposed to the paper's pure dice loss.
3. Normalization: Nonzero-voxel z-score normalization replaces the paper's percentile-based intensity scaling.
4. Augmentation: The paper's rotation and intensity augmentations are adapted for thin ChP structures, with left-right flips, mild deformations, and MRI intensity effects. Our version includes additional augmentations, such as T2 and MP2RAGE augmentations.
5. Training: Weight decay, epochs, and steps vary.