train_nnunet.md

Guide how to train your own nnU-Net on the TotalSegmentator dataset

1. Setup nnU-Net as described here
2. Download the data (v2.0.1)
3. Convert the data to nnU-Net format using resources/convert_dataset_to_nnunet.py (see resources/train_nnunet.sh for usage example)
4. Preprocess nnUNetv2_plan_and_preprocess -d <your_dataset_id> -pl ExperimentPlanner -c 3d_fullres -np 2
5. Train nnUNetv2_train <your_dataset_id> 3d_fullres 0 -tr nnUNetTrainerNoMirroring (takes several days) (See meta.csv inside of the dataset for train/val/test split. nnU-Net needs train and val split; use test split in step 6.)
6. Predict test set nnUNetv2_predict -i path/to/imagesTs -o path/to/labelsTs_predicted -d <your_dataset_id> -c 3d_fullres -tr nnUNetTrainerNoMirroring --disable_tta -f 0
7. Evaluate python resources/evaluate.py path/to/labelsTs path/to/labelsTs_predicted (requires pip install git+https://github.com/google-deepmind/surface-distance.git and pip install p_tqdm). The resulting numbers should be similar to the ones in resources/evaluate_results.txt (since training is not deterministic the mean dice score across all classes can vary by up to one dice point)
8. Done

Note: This will not give you the same results as TotalSegmentator for two reasons:

1. TotalSegmentator v2 uses a bigger dataset which is not completely public
2. TotalSegmentator is trained on images without blurred faces. Your dataset contains blurred faces.