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.gzformat - 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.gzfiles 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 |