Pytorch Code for the paper "Improved 3D abdominal multi-organ segmentation via boundary-constrained models"

This repository contains pytorch implementation corresponding to single-task (baseline) and multi-task (boundary-constrained) models.

1. 3D UNet
2. 3D UNet++
3. 3D Attention-UNet

Implementation for baseline models is in script here

1. 3D UNet-MTL-TSOL
2. 3D UNet++-MTL-TSOL
3. 3D UNet-MTL-TSD
4. 3D UNet++-MTL-TSD
5. 3D Attention-UNet-MTL-TSD
6. 3D Attention-UNet-MTL-TSOL

Implementation for boundary-constrained models is in script here

First download the datasets from links below:

1. Pancreas-CT dataset {For images, Use Link1 , For labels: Use Link2}
2. BTCV dataset {For images: Use Link1 , For labels: Use Link2}

1. To prepare the data, we use the pipeline utilized in Obelisk-Net paper
2. Use the ITK-snap to crop the scans and labels according to the bounding box coordinates given on Link2.

Organize the CT scans and their corresponding labels according to the format below:

Data Folder:
     --data:
            --images1:
                     --pancreas_ct1.nii.gz
                     --pancreas_ct2.nii.gz
                     .....................
            --labels1:
                     --label_ct1.nii.gz
                     --label_ct2.nii.gz
                     .....................

Organize the CT scans and their corresponding labels according to the format below:

Data Folder:
     --data:
            --images1:
                     --pancreas_ct1.nii.gz
                     --pancreas_ct2.nii.gz
                     .....................
            --labels1:
                     --label_ct1.nii.gz
                     --label_ct2.nii.gz
                     .....................
            --contours:
                     --label_ct1.nii.gz
                     --label_ct2.nii.gz
                     .....................

To train the baseline models, use the following command:

python train_baseline.py --data_folder 'path where the data is stored' --output_folder 'path to save the results' --batch 4 --epochs 300 --model unet --lr 0.001

To train the boundary-constrained models, use the following command:

python train_3d_boundary_constrained.py --data_folder 'path where the data is stored' --output_folder 'path to save the results' --batch 4 --epochs 300 \ --model unet --conf tsol --lr 0.001 --lambda_edge 0.5

The Models have been trained using two P100 GPUs, you will need to reduce the size of the batch if you use one GPU.

To train the baseline models on a different dataset, organize the data in the format described above and then use the training commands.