Info

This is the code repository of the work Tackling the Class Imbalance Problem of Deep Learning Based Head and Neck Organ Segmentation from Elias Tappeiner, Martin Welk and Rainer Schubert.

The implementation is based on the Monai DynUNet pipeline module, a reimplementation of the dynamic UNet used in the nnU-Net framework of Isensee et. al and further adapted to follow the nnU-Net parameterization.

Requirements

• Python: requirments.txt (the presented results are generated using the given config files and the python environment)
• Libs: plastimatch (additional for the result evaluation)

Code

Dataset preparation

1. Download the MICCAI Head and Neck segmentation challenge dataset
2. Run python src/scripts/combine_dataset_label_files.py --datasetpath path/to/unpacked/datasetfiles/ --outdir data/pddca
3. The dataset with combined labelmaps can now be found under data/pddca
4. In config/data a segmentation decatlon conform json file for the dataset is defined, which is used throughout the code to access the data

Train

Simply run the training script using one of the given config files or use your own. Details about the available configuration options are found in the default config file.

e.g. training the nnU-Net baseline:

python src/inference.py --experiment_config config/experiments/nnunet3d_nnUDice_ce.yaml 

Tensorboard logs are written to the model directory of the experiment, the parameter configuration of all experiments is logged using mlflow and can be found in mlruns after the training.

Inference

1. Pretrained weights are given here (skip next point if you trained the model by your own)
2. Extract the pretrained experiment models to models (the baseline weights should then be found under models/cars22/nnunet3d_caDice_ce/ckpt/checkpoint_final_iteration=125000.pt)
3. e.g. to infer the baseline experiment run (change yaml file to run your own):

python src/inference.py --experiment_config config/experiments/nnunet3d_nnUDice_ce.yaml