add kidney_cyst, breast, intervertebral_discs models

wasserth · wasserth · commit 47bad14db28c · 2025-01-13T15:34:23.000+01:00
diff --git a/README.md b/README.md
@@ -62,6 +62,7 @@ Openly available for any usage:
 * **lung_vessels**: lung_vessels (cite [paper](https://www.sciencedirect.com/science/article/pii/S0720048X22001097)), lung_trachea_bronchia
 * **body**: body, body_trunc, body_extremities, skin
 * **body_mr**: body_trunc, body_extremities (for MR images)
+* **vertebrae_mr**: sacrum, vertebrae_L5, vertebrae_L4, vertebrae_L3, vertebrae_L2, vertebrae_L1, vertebrae_T12, vertebrae_T11, vertebrae_T10, vertebrae_T9, vertebrae_T8, vertebrae_T7, vertebrae_T6, vertebrae_T5, vertebrae_T4, vertebrae_T3, vertebrae_T2, vertebrae_T1, vertebrae_C7, vertebrae_C6, vertebrae_C5, vertebrae_C4, vertebrae_C3, vertebrae_C2, vertebrae_C1 (for CT this is part of the `total` task)
 * **cerebral_bleed**: intracerebral_hemorrhage (cite [paper](https://www.mdpi.com/2077-0383/12/7/2631))*
 * **hip_implant**: hip_implant*
 * **coronary_arteries**: coronary_arteries*
@@ -72,17 +73,20 @@ Openly available for any usage:
 * **headneck_muscles**: sternocleidomastoid_right, sternocleidomastoid_left, superior_pharyngeal_constrictor, middle_pharyngeal_constrictor, inferior_pharyngeal_constrictor, trapezius_right, trapezius_left, platysma_right, platysma_left, levator_scapulae_right, levator_scapulae_left, anterior_scalene_right, anterior_scalene_left, middle_scalene_right, middle_scalene_left, posterior_scalene_right, posterior_scalene_left, sterno_thyroid_right, sterno_thyroid_left, thyrohyoid_right, thyrohyoid_left, prevertebral_right, prevertebral_left (cite [paper](https://www.mdpi.com/2072-6694/16/2/415))
 * **liver_vessels**: liver_vessels, liver_tumor (cite [paper](https://arxiv.org/abs/1902.09063))*
 * **oculomotor_muscles**: skull, eyeball_right, lateral_rectus_muscle_right, superior_oblique_muscle_right, levator_palpebrae_superioris_right, superior_rectus_muscle_right, medial_rectus_muscle_left, inferior_oblique_muscle_right, inferior_rectus_muscle_right, optic_nerve_left, eyeball_left, lateral_rectus_muscle_left, superior_oblique_muscle_left, levator_palpebrae_superioris_left, superior_rectus_muscle_left, medial_rectus_muscle_right, inferior_oblique_muscle_left, inferior_rectus_muscle_left, optic_nerve_right*
+* **lung_nodules**: lung, lung_nodules (this model is kindly provided by [BLUEMIND AI](https://bluemind.co/))
+* **kidney_cysts**: kidney_cyst_left, kidney_cyst_right (strongly improved accuracy compared to kidney_cysts inside of `total` task)
+* **breasts**: breast
 
 *: 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](https://backend.totalsegmentator.com/license-academic/). For a commercial license contact jakob.wasserthal@usb.ch):
 * **heartchambers_highres**: myocardium, atrium_left, ventricle_left, atrium_right, ventricle_right, aorta, pulmonary_artery (trained on sub-millimeter resolution)
 * **appendicular_bones**: patella, tibia, fibula, tarsal, metatarsal, phalanges_feet, ulna, radius, carpal, metacarpal, phalanges_hand
-* **apendicular_bones_mr**: patella, tibia, fibula, tarsal, metatarsal, phalanges_feet, ulna, radius (for MR images)
+* **appendicular_bones_mr**: patella, tibia, fibula, tarsal, metatarsal, phalanges_feet, ulna, radius (for MR images)
 * **tissue_types**: subcutaneous_fat, torso_fat, skeletal_muscle
 * **tissue_types_mr**: subcutaneous_fat, torso_fat, skeletal_muscle (for MR images)
 * **brain_structures**: brainstem, subarachnoid_space, venous_sinuses, septum_pellucidum, cerebellum, caudate_nucleus, lentiform_nucleus, insular_cortex, internal_capsule, ventricle, central_sulcus, frontal_lobe, parietal_lobe, occipital_lobe, temporal_lobe, thalamus (NOTE: this is for CT) (cite [paper](https://doi.org/10.1148/ryai.2020190183) as our model is partly based on this)
-* **vertebrae_body**: vertebral body of all vertebrae (without the vertebral arch)
+* **vertebrae_discs**: vertebral body of all vertebrae (without the vertebral arch), intervertebral_discs (for MR this is part of the `total_mr` task)
 * **face**: face_region (for anonymization)
 * **face_mr**: face_region (for anonymization)
 * **thigh_shoulder_muscles**: thigh_posterior_compartment_left, thigh_posterior_compartment_right, sartorius_left, sartorius_right, pectoralis_minor, serratus_anterior, teres_major, triceps_brachii (WIP)
diff --git a/totalsegmentator/bin/TotalSegmentator.py b/totalsegmentator/bin/TotalSegmentator.py
@@ -54,14 +54,15 @@ def main():
 
     # cerebral_bleed: Intracerebral hemorrhage
     # liver_vessels: hepatic vessels
-    parser.add_argument("-ta", "--task", choices=["total", "body", "body_mr",
+    parser.add_argument("-ta", "--task", choices=["total", "body", "body_mr", "vertebrae_mr",
                         "lung_vessels", "cerebral_bleed", "hip_implant", "coronary_arteries",
                         "pleural_pericard_effusion", "test",
                         "appendicular_bones", "appendicular_bones_mr", "tissue_types", "heartchambers_highres",
-                        "face", "vertebrae_body", "total_mr", "tissue_types_mr", "face_mr",
+                        "face", "vertebrae_discs", "total_mr", "tissue_types_mr", "face_mr",
                         "head_glands_cavities", "head_muscles", "headneck_bones_vessels", "headneck_muscles",
                         "brain_structures", "liver_vessels", "oculomotor_muscles",
-                        "thigh_shoulder_muscles", "thigh_shoulder_muscles_mr", "lung_nodules"],
+                        "thigh_shoulder_muscles", "thigh_shoulder_muscles_mr", "lung_nodules", "kidney_cysts", 
+                        "breasts"],
                         help="Select which model to use. This determines what is predicted.",
                         default="total")
 
diff --git a/totalsegmentator/bin/totalseg_download_weights.py b/totalsegmentator/bin/totalseg_download_weights.py
@@ -20,11 +20,13 @@ def main():
     parser.add_argument("-t", "--task", choices=["total", "total_fast", "total_mr", "total_fast_mr",
                                                  "lung_vessels", "cerebral_bleed",
                                                  "hip_implant", "coronary_arteries", "pleural_pericard_effusion",
-                                                 "body", "body_fast", "body_mr", "body_mr_fast", "vertebrae_body",
+                                                 "body", "body_fast", "body_mr", "body_mr_fast", "vertebrae_mr",
+                                                 "vertebrae_discs",
                                                  "heartchambers_highres", "appendicular_bones", 
                                                  "tissue_types", "tissue_types_mr", "face", "face_mr",
                                                  "head_glands_cavities", "head_muscles", "headneck_bones_vessels",
-                                                 "headneck_muscles", "liver_vessels", "brain_structures"],
+                                                 "headneck_muscles", "liver_vessels", "brain_structures",
+                                                 "lung_nodules", "kidney_cysts", "breasts"],
                         help="Task for which to download the weights", default="total")
 
     args = parser.parse_args()
@@ -43,17 +45,21 @@ def main():
         "body_fast": [300],
         "body_mr": [597],
         "body_mr_fast": [598],
+        "vertebrae_mr": [756],
         "head_glands_cavities": [775],
         "headneck_bones_vessels": [776],
         "head_muscles": [777],
         "headneck_muscles": [778, 779],
         "liver_vessels": [8],
+        "lung_nodules": [913],
+        "kidney_cysts": [789],
+        "breasts": [527],
 
         "heartchambers_highres": [301],
         "appendicular_bones": [304],
         "tissue_types": [481],
         "tissue_types_mr": [734],
-        "vertebrae_body": [302],
+        "vertebrae_discs": [305],
         "face": [303],
         "face_mr": [737],
         "brain_structures": [409]
diff --git a/totalsegmentator/libs.py b/totalsegmentator/libs.py
@@ -136,9 +136,6 @@ def download_pretrained_weights(task_id):
 
     config_dir = get_weights_dir()
     config_dir.mkdir(exist_ok=True, parents=True)
-    # (config_dir / "3d_fullres").mkdir(exist_ok=True, parents=True)
-    # (config_dir / "3d_lowres").mkdir(exist_ok=True, parents=True)
-    # (config_dir / "2d").mkdir(exist_ok=True, parents=True)
 
     old_weights = [
         "nnUNet/3d_fullres/Task251_TotalSegmentator_part1_organs_1139subj",
@@ -163,6 +160,7 @@ def download_pretrained_weights(task_id):
         # "Dataset297_TotalSegmentator_total_3mm_1559subj",  # for >= v2.0.4
         "Dataset297_TotalSegmentator_total_3mm_1559subj_v204",  # for >= v2.0.5
         # "Dataset298_TotalSegmentator_total_6mm_1559subj",  # for >= v2.0.5
+        "Dataset302_vertebrae_body_1559subj"
     ]
 
     # url = "http://backend.totalsegmentator.com"
@@ -233,6 +231,12 @@ def download_pretrained_weights(task_id):
     elif task_id == 351:
         weights_path = config_dir / "Dataset351_oculomotor_muscles_18subj"
         WEIGHTS_URL = url + "/v2.4.0-weights/Dataset351_oculomotor_muscles_18subj.zip"
+    elif task_id == 789:
+        weights_path = config_dir / "Dataset789_kidney_cyst_501subj"
+        WEIGHTS_URL = url + "/v2.5.0-weights/Dataset789_kidney_cyst_501subj.zip"
+    elif task_id == 527:
+        weights_path = config_dir / "Dataset527_breasts_1559subj"
+        WEIGHTS_URL = url + "/v2.5.0-weights/Dataset527_breasts_1559subj.zip"
                 
     # MR models
     elif task_id == 850:
@@ -248,11 +252,14 @@ def download_pretrained_weights(task_id):
         weights_path = config_dir / "Dataset853_TotalSegMRI_total_6mm_1088subj"
         WEIGHTS_URL = url + "/v2.5.0-weights/Dataset853_TotalSegMRI_total_6mm_1088subj.zip"
     elif task_id == 597:
-        weights_path = config_dir / "Dataset597_body_mr_139subj"
-        WEIGHTS_URL = url + "/v2.5.0-weights/Dataset597_body_mr_139subj.zip"
+        weights_path = config_dir / "Dataset597_mri_body_139subj"
+        WEIGHTS_URL = url + "/v2.5.0-weights/Dataset597_mri_body_139subj.zip"
     elif task_id == 598:
-        weights_path = config_dir / "Dataset598_body_mr_6mm_139subj"
-        WEIGHTS_URL = url + "/v2.5.0-weights/Dataset598_body_mr_6mm_139subj.zip"
+        weights_path = config_dir / "Dataset598_mri_body_6mm_139subj"
+        WEIGHTS_URL = url + "/v2.5.0-weights/Dataset598_mri_body_6mm_139subj.zip"
+    elif task_id == 756:
+        weights_path = config_dir / "Dataset756_mri_vertebrae_1076subj"
+        WEIGHTS_URL = url + "/v2.5.0-weights/Dataset756_mri_vertebrae_1076subj.zip"
 
     # Models from other projects
     elif task_id == 258:
@@ -307,8 +314,8 @@ def download_pretrained_weights(task_id):
         weights_path = config_dir / "Dataset303_face_1559subj"
     elif task_id == 481:
         weights_path = config_dir / "Dataset481_tissue_1559subj"
-    elif task_id == 302:
-        weights_path = config_dir / "Dataset302_vertebrae_body_1559subj"
+    elif task_id == 305:
+        weights_path = config_dir / "Dataset305_vertebrae_discs_1559subj"
     elif task_id == 854:
         weights_path = config_dir / "Dataset854_TotalSegMRI_tissue_1088subj"
     elif task_id == 856:
diff --git a/totalsegmentator/map_to_binary.py b/totalsegmentator/map_to_binary.py
@@ -328,6 +328,33 @@
         1: "body_trunc",
         2: "body_extremities",
     },
+    "vertebrae_mr": {
+        1: "sacrum",
+        2: "vertebrae_L5",
+        3: "vertebrae_L4", 
+        4: "vertebrae_L3",
+        5: "vertebrae_L2",
+        6: "vertebrae_L1",
+        7: "vertebrae_T12",
+        8: "vertebrae_T11",
+        9: "vertebrae_T10",
+        10: "vertebrae_T9",
+        11: "vertebrae_T8",
+        12: "vertebrae_T7",
+        13: "vertebrae_T6",
+        14: "vertebrae_T5",
+        15: "vertebrae_T4",
+        16: "vertebrae_T3",
+        17: "vertebrae_T2",
+        18: "vertebrae_T1",
+        19: "vertebrae_C7",
+        20: "vertebrae_C6",
+        21: "vertebrae_C5",
+        22: "vertebrae_C4",
+        23: "vertebrae_C3",
+        24: "vertebrae_C2",
+        25: "vertebrae_C1",
+    },
     "pleural_pericard_effusion": {
         # 1: "lung_pleural",
         2: "pleural_effusion",
@@ -337,8 +364,9 @@
         1: "liver_vessels",
         2: "liver_tumor"
     },
-    "vertebrae_body": {
-        1: "vertebrae_body"
+    "vertebrae_discs": {
+        1: "vertebrae_body",
+        2: "intervertebral_discs"
     },
     "heartchambers_highres": {
         1: "heart_myocardium",
@@ -539,6 +567,18 @@
         1: "lung",
         2: "lung_nodules"
     },
+    "kidney_cysts": {
+        1: "kidney_cyst_left",
+        2: "kidney_cyst_right"
+    },
+    # those classes need to be removed
+    "kidney_cysts_auxiliary": {
+        3: "kidney_left",
+        4: "kidney_right"
+    },
+    "breasts": {
+        1: "breast"
+    },
     "test": {
         1: "carpal"
     }
@@ -551,7 +591,7 @@
     "appendicular_bones_mr": 855,
     "tissue_types": 481,
     "tissue_types_mr": 854,
-    "vertebrae_body": 302,
+    "vertebrae_body": 305,
     "face": 303,
     "face_mr": 856,
     "brain_structures": 409,
diff --git a/totalsegmentator/preview.py b/totalsegmentator/preview.py
@@ -97,14 +97,19 @@
     "body_mr": [
         ["body_trunc", "body_extremities"]
     ],
+    "vertebrae_mr": [
+        ["sacrum", "vertebrae_L5", "vertebrae_L4", "vertebrae_L3", "vertebrae_L2", "vertebrae_L1", "vertebrae_T12", "vertebrae_T11", "vertebrae_T10", 
+         "vertebrae_T9", "vertebrae_T8", "vertebrae_T7", "vertebrae_T6", "vertebrae_T5", "vertebrae_T4", "vertebrae_T3", "vertebrae_T2", "vertebrae_T1", 
+         "vertebrae_C7", "vertebrae_C6", "vertebrae_C5", "vertebrae_C4", "vertebrae_C3", "vertebrae_C2", "vertebrae_C1"]
+    ],
     "pleural_pericard_effusion": [
         ["pleural_effusion", "pericardial_effusion"]
     ],
     "liver_vessels": [
         ["liver_vessels", "liver_tumor"]
     ],
-    "vertebrae_body": [
-        ["vertebrae_body"]
+    "vertebrae_discs": [
+        ["vertebrae_body", "intervertebral_discs"]
     ],
     "heartchambers_highres": [
         ["heart_myocardium"],
@@ -201,6 +206,12 @@
     "lung_nodules": [
         ["lung_nodules"]
     ],
+    "kidney_cysts": [
+        ["kidney_cyst_left", "kidney_cyst_right"]
+    ],
+    "breasts": [
+        ["breast"]
+    ],
     "test": [
         ["ulna"]
     ]
diff --git a/totalsegmentator/python_api.py b/totalsegmentator/python_api.py
@@ -245,18 +245,25 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         if fast:
             task_id = 598  # todo: train
             resample = 6.0
-            trainer = "nnUNetTrainer"
+            trainer = "nnUNetTrainer_DASegOrd0"
             crop = None
             model = "3d_fullres"
             folds = [0]
             if not quiet: print("Using 'fast' option: resampling to lower resolution (6mm)")
         else:
             task_id = 597
             resample = 1.5
-            trainer = "nnUNetTrainer"
+            trainer = "nnUNetTrainer_DASegOrd0"
             crop = None
             model = "3d_fullres"
             folds = [0]
+    elif task == "vertebrae_mr":
+        task_id = 756
+        resample = None
+        trainer = "nnUNetTrainer_DASegOrd0_NoMirroring"
+        crop = None
+        model = "3d_fullres"
+        folds = [0]
     elif task == "pleural_pericard_effusion":
         task_id = 315
         resample = None
@@ -335,17 +342,34 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         model = "3d_fullres"
         folds = [0]
         if fast: raise ValueError("task lung_nodules does not work with option --fast")
+    elif task == "kidney_cysts":
+        task_id = 789
+        resample = [1.5, 1.5, 1.5]
+        trainer = "nnUNetTrainer_DASegOrd0_NoMirroring"
+        crop = ["kidney_left", "kidney_right", "liver", "spleen", "colon"]
+        crop_addon = [10, 10, 10]
+        model = "3d_fullres"
+        folds = [0]
+        if fast: raise ValueError("task kidney_cysts does not work with option --fast")
+    elif task == "breasts":
+        task_id = 527
+        resample = [1.5, 1.5, 1.5]
+        trainer = "nnUNetTrainer_DASegOrd0_NoMirroring"
+        crop = None
+        model = "3d_fullres"
+        folds = [0]
+        if fast: raise ValueError("task breasts does not work with option --fast")
 
         
     # Commercial models
-    elif task == "vertebrae_body":
-        task_id = 302
+    elif task == "vertebrae_discs":
+        task_id = 305
         resample = 1.5
-        trainer = "nnUNetTrainer"
+        trainer = "nnUNetTrainer_DASegOrd0"
         crop = None
         model = "3d_fullres"
         folds = [0]
-        if fast: raise ValueError("task vertebrae_body does not work with option --fast")
+        if fast: raise ValueError("task vertebrae_discs does not work with option --fast")
         show_license_info()
     elif task == "heartchambers_highres":
         task_id = 301
