chadepl / punet-replication

Replication of some of the results of the Probabilistic UNet paper with accompanying PyTorch code.

Probabilistic UNet replication

This repository replicates some of the results of the paper A probabilistic u-net for segmentation of ambiguous images [1]. A write up of the replication can be found at this link.

The original code can be found at this link. We adapted it to improve readability and also ease of use for conducting experiments and doing inference.

Setup

1. conda create --name=punet-pytorch python=3.9
2. conda activate punet-pytorch
3. pip install torch torchvision torchaudio
4. pip install tqdm matplotlib scikit-image
5. pip install wandb for logging
6. pip install seaborn pandas for evaluation

To train the network:

• cd punet-pytorch
• python train_punet_lidc.py

References

[1] Kohl, S., Romera-Paredes, B., Meyer, C., De Fauw, J., Ledsam, J. R., Maier-Hein, K., ... & Ronneberger, O. (2018). A probabilistic u-net for segmentation of ambiguous images. Advances in neural information processing systems, 31.