nnUNet_att_position_correction README

The project's README from github.com/qpuchen/nnUNet_att_position_correction. As one markdown file: README.md.

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).

In the final ranking on unseen test data, this solution ranked 5th with a score of 81.47%.

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

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

nnUNet_att_position_correction (qpuchen/nnUNet_att_position_correction)

Searching Database…

这是做什么用的?这是 sdkxd 团队参加 MICCAI 2023 STS 挑战赛(从 CBCT 扫描中进行 3D 牙齿分割)的官方 PyTorch 代码。它在 nn-UNet 的基础上加入了轴向注意力(axial attention)和一个位置纠正模块,用来清理分割结果中的异常点。
它是一个独立工具,还是 nn-UNet 的扩展?它不是一个独立的应用,而是一组修改过的文件和脚本,需要复制到已有的 nn-UNet 安装目录里(例如打了补丁的 genericUNet.py),从而加入注意力模块和纠正模块。
我怎么把整个流程跑起来?先安装 nn-UNet,并把标注好的 CBCT 数据按它要求的 Task 文件夹结构整理好(dataset.json、imagesTr、imagesTs、labelsTr),然后运行 nn-UNet 的自动预处理,接着用 5 折交叉验证进行训练,之后再进入伪标签生成和 nnUNet-att 训练阶段。
使用它之前我需要先安装 nn-UNet 吗?需要。这套代码完全基于 nn-UNet 构建,默认你已经安装好 nn-UNet 并满足它的运行要求;除此之外不需要任何额外的依赖。
是不是一条命令就能跑完全流程?不是。这是一个多阶段的手动流程:先用有标注的数据训练,再为无标注数据生成伪标签,然后迭代地重新训练并过滤掉低质量的伪标签,接着训练 nnUNet-att,最后对预测结果执行位置纠正步骤。
这个方案在比赛中成绩如何?在最终的未见测试集上,sdkxd 团队的方案以 81.47% 的分数获得了总排名第 5 名。
位置纠正模块是干什么的?它是一个后处理步骤,在推理完成之后检测并修复预测牙齿掩膜中的异常分割点。
为什么要费劲生成伪标签?这次挑战赛中无标注的 CBCT 扫描远多于有标注的,所以整个流程先用有标注数据训练,然后在无标注扫描上迭代地生成并过滤伪标签,以半监督的方式扩充训练集。
nnUNet-att 是什么?它是修改后的 nn-UNet 架构,通过打了补丁的 genericUNet.py 加入轴向注意力层(axial attention),并在有标注数据加上过滤后伪标签的合并数据上用 5 折交叉验证进行训练。
Creator
qpuchen
Size
0.5 MB
Stars
342
Updated
CommitOct 12
Age
3y
Commits · 90 Days