Probabilistic 3D U-Net

Resources shared as part of the paper - Probabilistic 3D segmentation for aleatoric uncertainty quantification in full 3D medical data. SPIE paper, Arxiv paper.

Network Architecture:

Example predictions:

3D visuals:

Code

Since the repo is based on https://github.com/wolny/pytorch-3dunet, most of the code works the same way.

Setup

install requirements or use the sudochris/3dunet:v3 docker container.

Train

python train.py --config ./resources/probabilistic_3d_unet/train_config_vanilla_0.yaml

Validate

python validate.py --config ./resources/probabilistic_3d_unet/val_config_vanilla_0.yaml

Acknowledgement

The repository is based on the work by https://github.com/wolny/pytorch-3dunet and Kohl et. al., but adapted from Valiuddin et. al. and the PyTorch implemntation.