wasserth / TotalSegmentator

Tool for robust segmentation of >100 important anatomical structures in CT and MR images
Apache License 2.0
1.53k stars 255 forks source link

TotalSegmentator

Tool for segmentation of most major anatomical structures in any CT or MR image. It was trained on a wide range of different CT and MR images (different scanners, institutions, protocols,...) and therefore should work well on most images. A large part of the training dataset can be downloaded here: CT dataset (1228 subjects) and MR dataset (298 subjects). You can also try the tool online at totalsegmentator.com or as 3D Slicer extension.

ANNOUNCEMENT: We recently added support for MR images. Try out by using the task -ta total_mr or see more details in our paper.

Main classes for CT: Alt text

Main classes for MR: Alt text

See here for all available structures.

Created by the department of Research and Analysis at University Hospital Basel. If you use it please cite our Radiology AI paper (free preprint). If you use it for MR images please cite the TotalSegmentator MRI paper. Please also cite nnUNet since TotalSegmentator is heavily based on it.

Installation

TotalSegmentator works on Ubuntu, Mac, and Windows and on CPU and GPU.

Install dependencies:

Optionally:

Install Totalsegmentator

pip install TotalSegmentator

Usage

For CT images:

TotalSegmentator -i ct.nii.gz -o segmentations

For MR images:

TotalSegmentator -i mri.nii.gz -o segmentations --task total_mr

Note: A Nifti file or a folder (or zip file) with all DICOM slices of one patient is allowed as input.

Note: If you run on CPU use the option --fast or --roi_subset to greatly improve runtime.

Note: This is not a medical device and is not intended for clinical usage.

Subtasks

Alt text

Next to the default task (total) there are more subtasks with more classes. If the taskname ends with _mr it works for MR images, otherwise for CT images.

Openly available for any usage:

*: These models are not trained on the full totalsegmentator dataset but on some small other datasets. Therefore, expect them to work less robustly.

Available with a license (free licenses available for non-commercial usage here. For a commercial license contact jakob.wasserthal@usb.ch):

Usage:

TotalSegmentator -i ct.nii.gz -o segmentations -ta <task_name>

Confused by all the structures and tasks? Check this to search through available structures and tasks.

The mapping from label ID to class name can be found here.

Advanced settings

Run via docker

We also provide a docker container which can be used the following way

docker run --gpus 'device=0' --ipc=host -v /absolute/path/to/my/data/directory:/tmp wasserth/totalsegmentator:2.2.1 TotalSegmentator -i /tmp/ct.nii.gz -o /tmp/segmentations

Running v1

If you want to keep on using TotalSegmentator v1 (e.g. because you do not want to change your pipeline) you can install it with the following command:

pip install TotalSegmentator==1.5.7

The documentation for v1 can be found here. Bugfixes for v1 are developed in the branch v1_bugfixes. Our Radiology AI publication refers to TotalSegmentator v1.

Resource Requirements

Totalsegmentator has the following runtime and memory requirements (using an Nvidia RTX 3090 GPU): (1.5mm is the normal model and 3mm is the --fast model. With v2 the runtimes have increased a bit since we added more classes.)

Alt text

If you want to reduce memory consumption you can use the following options:

Python API

You can run totalsegmentator via Python:

import nibabel as nib
from totalsegmentator.python_api import totalsegmentator

if __name__ == "__main__":
    # option 1: provide input and output as file paths
    totalsegmentator(input_path, output_path)

    # option 2: provide input and output as nifti image objects
    input_img = nib.load(input_path)
    output_img = totalsegmentator(input_img)
    nib.save(output_img, output_path)

You can see all available arguments here. Running from within the main environment should avoid some multiprocessing issues.

The segmentation image contains the names of the classes in the extended header. If you want to load this additional header information you can use the following code (requires pip install xmltodict):

from totalsegmentator.nifti_ext_header import load_multilabel_nifti

segmentation_nifti_img, label_map_dict = load_multilabel_nifti(image_path)

Install latest master branch (contains latest bug fixes)

pip install git+https://github.com/wasserth/TotalSegmentator.git

Other commands

If you want to know which contrast phase a CT image is you can use the following command (requires pip install xgboost). More details can be found here:

totalseg_get_phase -i ct.nii.gz -o contrast_phase.json

If you want to know which modality (CT or MR) a image is you can use the following command (requires pip install xgboost).

totalseg_get_modality -i image.nii.gz -o modality.json

If you want to combine some subclasses (e.g. lung lobes) into one binary mask (e.g. entire lung) you can use the following command:

totalseg_combine_masks -i totalsegmentator_output_dir -o combined_mask.nii.gz -m lungcomm 

Normally weights are automatically downloaded when running TotalSegmentator. If you want to download the weights with an extra command (e.g. when building a docker container) use this:

totalseg_download_weights -t <task_name>

This will download them to ~/.totalsegmentator/nnunet/results. You can change this path by doing export TOTALSEG_HOME_DIR=/new/path/.totalsegmentator. If your machine has no internet, then download on another machine with internet and copy ~/.totalsegmentator to the machine without internet.

After acquiring a license number for the non-open tasks you can set it with the following command:

totalseg_set_license -l aca_12345678910

Train/validation/test split

The exact split of the dataset can be found in the file meta.csv inside of the dataset. This was used for the validation in our paper. The exact numbers of the results for the high-resolution model (1.5mm) can be found here. The paper shows these numbers in the supplementary materials Figure 11.

Retrain model and run evaluation

See here for more info on how to train a nnU-Net yourself on the TotalSegmentator dataset, how to split the data into train/validation/test set as in our paper, and how to run the same evaluation as in our paper.

Typical problems

ITK loading Error When you get the following error message

ITK ERROR: ITK only supports orthonormal direction cosines. No orthonormal definition was found!

you should do

pip install SimpleITK==2.0.2

Alternatively you can try

fslorient -copysform2qform input_file
fslreorient2std input_file output_file

Bad segmentations When you get bad segmentation results check the following:

Other

Reference

For more details see our Radiology AI paper (freely available preprint). If you use this tool please cite it as follows

Wasserthal, J., Breit, H.-C., Meyer, M.T., Pradella, M., Hinck, D., Sauter, A.W., Heye, T., Boll, D., Cyriac, J., Yang, S., Bach, M., Segeroth, M., 2023. TotalSegmentator: Robust Segmentation of 104 Anatomic Structures in CT Images. Radiology: Artificial Intelligence. https://doi.org/10.1148/ryai.230024

Please also cite nnUNet since TotalSegmentator is heavily based on it. Moreover, we would really appreciate it if you let us know what you are using this tool for. You can also tell us what classes we should add in future releases. You can do so here.

Class details

The following table shows a list of all classes for task total.

TA2 is a standardized way to name anatomy. Mostly the TotalSegmentator names follow this standard. For some classes they differ which you can see in the table below.

Here you can find a mapping of the TotalSegmentator classes to SNOMED-CT codes.

Index TotalSegmentator name TA2 name
1 spleen
2 kidney_right
3 kidney_left
4 gallbladder
5 liver
6 stomach
7 pancreas
8 adrenal_gland_right suprarenal gland
9 adrenal_gland_left suprarenal gland
10 lung_upper_lobe_left superior lobe of left lung
11 lung_lower_lobe_left inferior lobe of left lung
12 lung_upper_lobe_right superior lobe of right lung
13 lung_middle_lobe_right middle lobe of right lung
14 lung_lower_lobe_right inferior lobe of right lung
15 esophagus
16 trachea
17 thyroid_gland
18 small_bowel small intestine
19 duodenum
20 colon
21 urinary_bladder
22 prostate
23 kidney_cyst_left
24 kidney_cyst_right
25 sacrum
26 vertebrae_S1
27 vertebrae_L5
28 vertebrae_L4
29 vertebrae_L3
30 vertebrae_L2
31 vertebrae_L1
32 vertebrae_T12
33 vertebrae_T11
34 vertebrae_T10
35 vertebrae_T9
36 vertebrae_T8
37 vertebrae_T7
38 vertebrae_T6
39 vertebrae_T5
40 vertebrae_T4
41 vertebrae_T3
42 vertebrae_T2
43 vertebrae_T1
44 vertebrae_C7
45 vertebrae_C6
46 vertebrae_C5
47 vertebrae_C4
48 vertebrae_C3
49 vertebrae_C2
50 vertebrae_C1
51 heart
52 aorta
53 pulmonary_vein
54 brachiocephalic_trunk
55 subclavian_artery_right
56 subclavian_artery_left
57 common_carotid_artery_right
58 common_carotid_artery_left
59 brachiocephalic_vein_left suprarenal gland
60 brachiocephalic_vein_right
61 atrial_appendage_left
62 superior_vena_cava
63 inferior_vena_cava
64 portal_vein_and_splenic_vein hepatic portal vein
65 iliac_artery_left common iliac artery
66 iliac_artery_right common iliac artery
67 iliac_vena_left common iliac vein
68 iliac_vena_right common iliac vein
69 humerus_left
70 humerus_right
71 scapula_left
72 scapula_right
73 clavicula_left clavicle
74 clavicula_right clavicle
75 femur_left
76 femur_right
77 hip_left
78 hip_right
79 spinal_cord
80 gluteus_maximus_left gluteus maximus muscle
81 gluteus_maximus_right gluteus maximus muscle
82 gluteus_medius_left gluteus medius muscle
83 gluteus_medius_right gluteus medius muscle
84 gluteus_minimus_left gluteus minimus muscle
85 gluteus_minimus_right gluteus minimus muscle
86 autochthon_left
87 autochthon_right
88 iliopsoas_left iliopsoas muscle
89 iliopsoas_right iliopsoas muscle
90 brain
91 skull
92 rib_left_1
93 rib_left_2
94 rib_left_3
95 rib_left_4
96 rib_left_5
97 rib_left_6
98 rib_left_7
99 rib_left_8
100 rib_left_9
101 rib_left_10
102 rib_left_11
103 rib_left_12
104 rib_right_1
105 rib_right_2
106 rib_right_3
107 rib_right_4
108 rib_right_5
109 rib_right_6
110 rib_right_7
111 rib_right_8
112 rib_right_9
113 rib_right_10
114 rib_right_11
115 rib_right_12
116 sternum
117 costal_cartilages

Class map for task total_mr:

Index TotalSegmentator name TA2 name
1 spleen
2 kidney_right
3 kidney_left
4 gallbladder
5 liver
6 stomach
7 pancreas
8 adrenal_gland_right suprarenal gland
9 adrenal_gland_left suprarenal gland
10 lung_left
11 lung_right
12 esophagus
13 small_bowel small intestine
14 duodenum
15 colon
16 urinary_bladder
17 prostate
18 sacrum
19 vertebrae
20 intervertebral_discs
21 spinal_cord
22 heart
23 aorta
24 inferior_vena_cava
25 portal_vein_and_splenic_vein hepatic portal vein
26 iliac_artery_left common iliac artery
27 iliac_artery_right common iliac artery
28 iliac_vena_left common iliac vein
29 iliac_vena_right common iliac vein
30 humerus_left
31 humerus_right
32 fibula
33 tibia
34 femur_left
35 femur_right
36 hip_left
37 hip_right
38 gluteus_maximus_left gluteus maximus muscle
39 gluteus_maximus_right gluteus maximus muscle
40 gluteus_medius_left gluteus medius muscle
41 gluteus_medius_right gluteus medius muscle
42 gluteus_minimus_left gluteus minimus muscle
43 gluteus_minimus_right gluteus minimus muscle
44 autochthon_left
45 autochthon_right
46 iliopsoas_left iliopsoas muscle
47 iliopsoas_right iliopsoas muscle
48 quadriceps_femoris_left
49 quadriceps_femoris_right
50 thigh_medial_compartment_left
51 thigh_medial_compartment_right
52 thigh_posterior_compartment_left
53 thigh_posterior_compartment_right
54 sartorius_left
55 sartorius_right
56 brain