ATLAS-Net: Multi-Organ Abdominal CT Segmentation

This repository provides the nnUNet v2 checkpoint for multi-organ abdominal CT segmentation.

The model was trained on a large-scale multi-center CT dataset spanning 145 medical centers, with pseudo-labels automatically corrected using organ-specific anatomical rules across 25 organ classes. The strength of this model lies in its segmentation of gastrointenstinal tract organs.

First of all, you need a Linux OS server and basic terminal usage. You also need an NVIDIA GPU with CUDA 11.7 or higher.


0. Before You Start

You need:

  • CT files in .nii.gz format
  • Each CT is a 3D volume (not single slices)
  • File names must end with _0000.nii.gz

1. Installation

git clone https://huggingface.co/Koushik45048545309/ATLAS-Net
conda create -n atlasnet python=3.10 -y
conda activate atlasnet
cd ATLAS-Net
pip install nnunetv2

2. Run Inference on Example CT Scans

Example CT scans are stored in this path: ./nnUNet_eval/Dataset001_ATLASNet/imagesTs/

This is the organization of example CT scans:

nnUNet_eval/
  └── Dataset001_ATLASNet/
        └── imagesTs/
              β”œβ”€β”€ PanTS_00000001_0000.nii.gz
              β”œβ”€β”€ PanTS_00000002_0000.nii.gz
              β”œβ”€β”€ PanTS_00000003_0000.nii.gz
              └── ...

Directly run this command in terminal:

CUDA_VISIBLE_DEVICES=0 bash inference.sh

Replace 0 with whichever GPU you want to use. To check which GPUs are available: nvidia-smi

If you want to run inference on your own CT scans, place them in the same folder. Make sure the file name ends with _0000.nii.gz.

Running on your own CT scans

CT files must be named with the suffix _0000.nii.gz and placed in ./nnUNet_eval/Dataset001_ATLASNet/imagesTs/.

Option A β€” your files are already named correctly (e.g. case_0000.nii.gz):

cp /path/to/your/case_0000.nii.gz ./nnUNet_eval/Dataset001_ATLASNet/imagesTs/
CUDA_VISIBLE_DEVICES=0 bash inference.sh

Replace 0 with whichever GPU you want to use. To check which GPUs are available: nvidia-smi

Option B β€” your dataset is organized as DATASET_NAME/BDMAP_00000001/ct.nii.gz:

Run the provided linking script β€” no files are copied or renamed, symlinks are created automatically in the correct location:

bash create_links.sh /path/to/DATASET_NAME
CUDA_VISIBLE_DEVICES=0 bash inference.sh

Replace 0 with whichever GPU you want to use. To check which GPUs are available: nvidia-smi


3. Output

Predictions are saved to ./nnUNet_predictions/:

nnUNet_predictions/
  β”œβ”€β”€ PanTS_00000001.nii.gz
  β”œβ”€β”€ PanTS_00000002.nii.gz
  β”œβ”€β”€ PanTS_00000003.nii.gz
  └── ...

4. Postprocessing

Single file:

python utils/atlas_postprocess.py \
    -p ./nnUNet_predictions/case.nii.gz \   # path to your model's predicted segmentation mask
    -c ./ct/case_0000.nii.gz \              # path to the original CT scan in Hounsfield units
    -o ./corrected/case.nii.gz               # path where the corrected mask will be saved

Folder of predictions (processes all .nii.gz files, CT matched by filename stem):

python utils/atlas_postprocess.py \
    -p ./nnUNet_predictions/ \   # folder of prediction .nii.gz files
    -c ./ct/ \                   # folder of CT .nii.gz files (matched by filename stem)
    -o ./corrected/               # output folder (created if it doesn't exist)

Custom labels (only needed if your model uses different integers than the ATLAS-Net defaults):

python utils/atlas_postprocess.py \
    -p ./nnUNet_predictions/case.nii.gz \   # path to your model's predicted segmentation mask
    -c ./ct/case_0000.nii.gz \              # path to the original CT scan in Hounsfield units
    -o ./corrected/case.nii.gz \            # path where the corrected mask will be saved
    --labels liver:12 spleen:17 colon:6 stomach:18   # organ_name:label_index pairs

Note: The CT scan must be in Hounsfield units (HU) β€” this is the standard output of CT scanners and is preserved in .nii.gz files unless you have explicitly normalized them. If your CT values range from 0–255 or 0–1, do not use this postprocessing step as the HU-based corrections will not apply correctly.


Predictions will be saved as .nii.gz files. Refer the label map below to extract specific labels.

5. Extract Organs (Optional)

Extracts per-organ binary masks from ATLAS-Net multi-label predictions, one .nii.gz file per organ per case β€” same format as TotalSegmentator output.

# extract all organs from all cases in a folder
python extract_organs.py \
    -i ./predictions \    # folder containing your model's .nii.gz prediction files
    -o ./organs           # folder where per-case organ masks will be saved
# extract specific organs by name
python extract_organs.py \
    -i ./predictions \    # input folder
    -o ./organs \         # output folder
    --organs liver spleen colon pancreas
# extract specific organs by label index
python extract_organs.py \
    -i ./predictions \    # input folder
    -o ./organs \         # output folder
    --labels 12 17 6 13
# list all available organ names and label indices
python extract_organs.py --list

6. Label Map

Label Organ
0 Background
1 Aorta
2 Adrenal gland (L)
3 Adrenal gland (R)
4 Common bile duct
5 Celiac artery
6 Colon
7 Duodenum
8 Gallbladder
9 Inferior vena cava
10 Kidney (L)
11 Kidney (R)
12 Liver
13 Pancreas
14 Pancreatic duct
15 Superior mesenteric artery
16 Small intestine
17 Spleen
18 Stomach
19 Portal/splenic veins
20 Renal vein (L)
21 Renal vein (R)
22 CBD stent
23 Pancreatic PDAC
24 Pancreatic cyst
25 Pancreatic PNET

License

CC BY 4.0

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support