AI that knows when it does not know
FLARE22 codebase of team brananas
This repository contains the codebase of team brananas (Natalia Alves & Bram de Wilde) for the FLARE22 MICCAI challenge. This challenge features 50 labelled and 2000 unlabelled abdominal CT scans, on which 13 organs have to be annotated. We propose an uncertainty-guided self-learning framework to attack this semi-supervised organ segmentation problem. You can read our paper here.
Training
All models are trained with the nnU-Net
framework. Make sure to first install all dependencies in requirements.txt.
Training follows an iterative process:
- Train 5 models (5-fold cross validation) with (pseudo-)labelled data
- Perform inference on unlabelled data
- Calculate uncertainty between the 5 models for each case
- Include pseudolabelled cases that pass the uncertainty threshold
- Repeat from step 1
Note that you have to manually edit some paths in files for this to work. Please refer to the nnU-Net documentation for training and inference of the models.
1. Train 5 models in 5-fold cross validation
nnUNet_train 3d_fullres Task101_FLARE [0, 1, 2, 3, 4]
2. Perform inference on unlabelled data
nnUNet_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -t Task101_FLARE -m 3d_fullres
3. Calculate uncertainty metrics
python uncertainty_metrics/calculate_metrics.py
4. Include pseudolabelled cases that pass the uncertainty threshold
python ensemble.py
python generate_iteration.py
Inference
To speed up nnU-Net inference, we built a custom inference pipeline with some slight changes. The final inference command also includes custom postprocessing.
python docker_inference_final.py
Docker
You can build a Dockerfile for this model as follows:
docker build . --tag brananas:latest
Citation
If you use this framework, or parts of it, please cite us as follows:
@InProceedings{10.1007/978-3-031-23911-3_11,
author="Alves, Nat{\'a}lia
and de Wilde, Bram",
editor="Ma, Jun
and Wang, Bo",
title="Uncertainty-Guided Self-learning Framework for Semi-supervised Multi-organ Segmentation",
booktitle="Fast and Low-Resource Semi-supervised Abdominal Organ Segmentation",
year="2022",
publisher="Springer Nature Switzerland",
address="Cham",
pages="116--127",
isbn="978-3-031-23911-3"
}