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…

354 ·GitHub
NnUNet_att_position_correction
Что это такое?Это официальный PyTorch-код команды sdkxd для MICCAI 2023 STS Challenge (3D-сегментация зубов по CBCT-снимкам). Он расширяет 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 и предполагает, что он уже установлен и соответствует его требованиям; никаких дополнительных зависимостей сверх этого не нужно.
Это пайплайн в одну команду?Нет, это многоэтапный ручной процесс: сначала обучение на размеченных данных, затем генерация псевдометок для неразмеченных данных, итеративное дообучение и отсеивание некачественных псевдометок, потом обучение nnUNet-att и в конце запуск шага позиционной коррекции на предсказаниях.
Насколько хорошо это решение выступило в челлендже?На финальном закрытом тестовом наборе решение команды sdkxd заняло 5-е место в общем зачёте с результатом 81,47%.
Что делает модуль позиционной коррекции?Это шаг постобработки: после инференса он находит и исправляет аномальные точки сегментации в предсказанных масках зубов.
Зачем вообще генерировать псевдометки?В челлендже неразмеченных CBCT-снимков намного больше, чем размеченных, поэтому пайплайн сначала обучается на размеченных данных, а затем итеративно генерирует и отфильтровывает псевдометки на неразмеченных снимках, расширяя обучающую выборку в полуавтоматическом режиме.
Что такое nnUNet-att?Это архитектура nn-UNet, дополненная слоями axial attention (через исправленный genericUNet.py) и обученная с 5-кратной кросс-валидацией на объединённых данных: размеченных плюс отфильтрованных псевдоразмеченных.
What are similar projects?VisionHOPE, Attention-Residuals, pointdit, External-Attention-pytorch, awesome-semantic-segmentation, Engram, deep-learning-for-image-processing, CVPR2026-Papers-with-Code
Which topics is it in?medical-imaging, convolutional-neural-networks, semantic-segmentation, image-segmentation, model-training, attention-mechanism, paper, neural-network, 3d, academic-research, transformer, data-science
Creator
qpuchen
Size
0.5 MB
Stars
354
Updated
CommitOct 12
Age
3y
Commits · 90 Days