We devise a Multi-granularity Boundary Aggregation (MBA) module to introduce boundary features into the diffusion model for leveraging boundary information to improve the 3D segmentation accuracy. %offering feature supplementation and boundary perception.
We propose a Monte Carlo Diffusion (MC-Diff) module, obtaining an uncertainty map during the diffusion training stage to guide the computation of segmentation loss and improve the segmentation on uncertain pixels. %and increasing focus on uncertain pixels.
During the diffusion inference stage, we design a Progressive Uncertainty-driven REfinement (PURE) strategy to fuse segmentation results from different inference steps for further boosting 3D medical image segmentation. %the robustness of predictions.
- torch 2.1.2
- CUDA 12.2
- monai
- dynamic_network_architectures
- batchgenerator
https://codalab.lisn.upsaclay.fr/competitions/13238#learn_the_details-overview
The data should be placed in "./data/raw_data/AIIB23_Train_T1/"
run the preprocessing script to generate the normalized and resampled data
python 1_preprocessing_aiib2023.py
After running the script, the preprocessed data will be saved in
"./data/fullres/train/"train the model
python 2_train_diffunet_aiib2023.py
The trained model will be saved in
"./logs/diffunet/model/"Meanwhile, you can use tensorboard to monitor the training process.tensorboard --logdir ./logs/diffunet/
inference using the trained model
python 3_predict_aiib2023.py
the inference results will be saved in
"./prediction_results/diffunet/"compute the metrics
python 4_compute_metrics_aiib2023.py --pred=diffunet
https://arxiv.org/pdf/2305.19369
The data should be placed in "./data/raw_data/ASNR-MICCAI-BraTS2023-GLI-Challenge-TrainingData/"
- run the preprocessing script to generate the normalized and resampled data
python 1_rename_mri_data_brats2023.py python 2_preprocessing_brats2023.py
- train the model
python 3_train_diffunet_brats2023.py
- inference using the trained model
python 4_predict_brats2023.py
- compute the metrics
python 5_compute_metrics_brats2023.py --pred=diffunet
To simplify the running process, rather than predict the boundary information, we directly use the segmentation labels in the current version. In the future, we will continue to release more related codes.