Semi-Supervised 3D Tooth Segmentation Using nn-UNet with Axial Attention and Positional Correction

Official PyTorch implementation of Semi-Supervised 3D Tooth Segmentation Using nn-UNet with Axial Attention and Positional Correction (MICCAI 2023 STS).

Environments and Requirements:

Install nnU-Net as below. You should meet the requirements of nnUNet, our method does not need any additional requirements. For more details, please refer to https://github.com/MIC-DKFZ/nnUNet/tree/v1.7.1

git clone https://github.com/MIC-DKFZ/nnUNet.git
cd nnUNet
pip install -e .

1. Training nnUNet for Labeled Data

1.1. Prepare Labeled Data of MICCAI 2023 Challenge

Following nnUNet, give a TaskID (e.g. Task001) to the labeled data and organize them folowing the requirement of nnUNet.

nnUNet_raw_data_base/nnUNet_raw_data/Task01_Tooth/
├── dataset.json
├── imagesTr
├── imagesTs
└── labelsTr

1.2. Conduct automatic preprocessing using nnUNet.

Here we do use the default setting.

nnUNet_plan_and_preprocess -t 1 --verify_dataset_integrity

1.3. Training nnUNet by 5-fold Cross Validation

for FOLD in 0 1 2 3 4
do
CUDA_VISIBLE_DEVICES=0,1 nnUNet_train_DP 3d_fullres nnUNetTrainerV2_DP 1 $FOLD -gpus 2 -c --npz
done

2. Generate Pseudo Labels for Unlabeled Data

2.1. Generate Pseudo Labels

nnUNet_predict -i $INPUTS_FOLDER -o $OUTPUTS_FOLDER -t 2 -m 3d_fullres --save_npz

2.2. Iteratively Train Models and Generate Pseudo Labels

Give a new TaskID (e.g. Task002) and organize the Labeled Data and Pseudo Labeled Data as above.

Conduct automatic preprocessing using nnUNet as above.

nnUNet_plan_and_preprocess -t 2 --verify_dataset_integrity

Training new nnUNet by all training data

for FOLD in 0 1 2 3 4
do
CUDA_VISIBLE_DEVICES=0,1 nnUNet_train_DP 3d_fullres nnUNetTrainerV2_DP 2 $FOLD -gpus 2 -c --npz
done

Generate new pseudo labels for unlabeled data.

2.3. Filter Low-quality Pseudo Labels

We compare Pseudo Labels in different rounds and filter out the labels with high variants.

select_pseudo_label.ipynb

3. Training nnUNet-att by 5-fold Cross Validation

3.1. Copy the following files in this repo to your nnUNet environment.

./nnunet/network_architecture/generic_UNet.py

3.2. Training nnUNet-att by 5-fold Cross Validation

for FOLD in 0 1 2 3 4
do
CUDA_VISIBLE_DEVICES=0,1 nnUNet_train_DP 3d_fullres nnUNetTrainerV2_DP 3 $FOLD -gpus 2 -c --npz
done

3.3. Add the '_000' suffix to the file name of the testset data

pyhton add_suffix.py

3.4. Do Efficient Inference with nnUNet-att

We modify the generic_UNet.py of nnunet source code for efficiency. Please make sure the code backup is done and then copy the whole repo to your nnunet environment.

nnUNet_predict -i $INPUTS_FOLDER -o $OUTPUTS_FOLDER -t 3 -m 3d_fullres --save_npz

4. Correction of abnormal segmentation points through position correction module

python position_correction.py

Chaoqun-Guo / nnUNet_att_position_correction