add liver vessels and brain structures (WIP)

CHANGELOG.md

@@ -1,4 +1,6 @@
 ## Master
+* add brain structures
+* add liver vessels
 
 
 ## Release 2.3.0

totalsegmentator/bin/TotalSegmentator.py

@@ -59,8 +59,8 @@ def main():
                         "pleural_pericard_effusion", "test",
                         "appendicular_bones", "tissue_types", "heartchambers_highres",
                         "face", "vertebrae_body", "total_mr", "tissue_types_mr", "face_mr",
-                        "head_glands_cavities", "head_muscles", "headneck_bones_vessels", "headneck_muscles"],
-                        # future: liver_vessels, head,
+                        "head_glands_cavities", "head_muscles", "headneck_bones_vessels", "headneck_muscles",
+                        "brain_structures", "liver_vessels"],
                         help="Select which model to use. This determines what is predicted.",
                         default="total")
 

totalsegmentator/bin/totalseg_download_weights.py

@@ -24,7 +24,7 @@ def main():
                                                  "heartchambers_highres", "appendicular_bones", 
                                                  "tissue_types", "tissue_types_mr", "face", "face_mr",
                                                  "head_glands_cavities", "head_muscles", "headneck_bones_vessels",
-                                                 "headneck_muscles"],
+                                                 "headneck_muscles", "liver_vessels", "brain_structures"],
                         help="Task for which to download the weights", default="total")
 
     args = parser.parse_args()
@@ -41,11 +41,11 @@ def main():
         "pleural_pericard_effusion": [315],
         "body": [299],
         "body_fast": [300],
-
         "head_glands_cavities": [775],
         "headneck_bones_vessels": [776],
         "head_muscles": [777],
         "headneck_muscles": [778, 779],
+        "liver_vessels": [8],
 
         "heartchambers_highres": [301],
         "appendicular_bones": [304],
@@ -54,7 +54,7 @@ def main():
         "vertebrae_body": [302],
         "face": [303],
         "face_mr": [737],
-        # "liver_vessels": [8]
+        "brain_structures": [409]
     }
 
     setup_totalseg()

totalsegmentator/libs.py

@@ -282,8 +282,7 @@ def download_pretrained_weights(task_id):
         WEIGHTS_URL = url + "/v2.0.0-weights/Dataset503_cardiac_motion.zip"
     elif task_id == 8:
         weights_path = config_dir / "Task008_HepaticVessel"
-        # WEIGHTS_URL = "https://zenodo.org/record/7573746/files/Task008_HepaticVessel.zip?download=1"
-        WEIGHTS_URL = "todo"
+        WEIGHTS_URL = url + "/v2.0.0-weights/Dataset008_HepaticVessel.zip"
 
     # Commercial models
     elif task_id == 304:
@@ -300,6 +299,8 @@ def download_pretrained_weights(task_id):
         weights_path = config_dir / "Dataset734_TotalSegmentatorMRI_tissue_495subj"
     elif task_id == 737:
         weights_path = config_dir / "Dataset737_TotalSegmentatorMRI_face_495subj"
+    elif task_id == 409:
+        weights_path = config_dir / "Dataset409_neuro_550subj"
 
     else:
         raise ValueError(f"For task_id {task_id} no download path was found.")

totalsegmentator/map_to_binary.py

@@ -371,6 +371,24 @@
         2: "brain",
         3: "liver"
     },
+    "brain_structures": {
+        1: "brainstem",
+        2: "subarachnoid_space",
+        3: "venous_sinuses",
+        4: "septum_pellucidum",
+        5: "cerebellum",
+        6: "caudate_nucleus",
+        7: "lentiform_nucleus",
+        8: "insular_cortex",
+        9: "internal_capsule",
+        10: "ventricle",
+        11: "central_sulcus",
+        12: "frontal_lobe",
+        13: "parietal_lobe",
+        14: "occipital_lobe",
+        15: "temporal_lobe",
+        16: "thalamus"
+    },
     "head_glands_cavities": {
         1: "eye_left",
         2: "eye_right",
@@ -458,11 +476,8 @@
     "vertebrae_body": 302,
     "face": 303,
     "face_mr": 737,
+    "brain_structures": 409
 }
-# future
-# - brain subparts
-# - head subparts
-# - everything highres?
 
 
 class_map_5_parts = {

totalsegmentator/preview.py

@@ -133,6 +133,12 @@
     "face_mr": [
         ["face"]
     ],
+    "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"]
+    ],
     "head_glands_cavities": [
         ["eye_left", "eye_right", "eye_lens_left", "eye_lens_right",
          "optic_nerve_left", "optic_nerve_right", "parotid_gland_left", "parotid_gland_right",

totalsegmentator/python_api.py

@@ -247,7 +247,7 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         crop = ["liver"]
         crop_addon = [20, 20, 20]
         model = "3d_fullres"
-        folds = None
+        folds = [0]
         if fast: raise ValueError("task liver_vessels does not work with option --fast")
     elif task == "head_glands_cavities":
         task_id = 775
@@ -256,7 +256,7 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         crop = ["skull"]
         crop_addon = [10, 10, 10]
         model = "3d_fullres_high"
-        folds = None
+        folds = [0]
         if fast: raise ValueError("task head_glands_cavities does not work with option --fast")
     elif task == "headneck_bones_vessels":
         task_id = 776
@@ -267,7 +267,7 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         crop = ["clavicula_left", "clavicula_right", "vertebrae_C1", "vertebrae_C5", "vertebrae_T1", "vertebrae_T4"]
         crop_addon = [40, 40, 40]
         model = "3d_fullres_high"
-        folds = None
+        folds = [0]
         if fast: raise ValueError("task headneck_bones_vessels does not work with option --fast")
     elif task == "head_muscles":
         task_id = 777
@@ -276,7 +276,7 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         crop = ["skull"]
         crop_addon = [10, 10, 10]
         model = "3d_fullres_high"
-        folds = None
+        folds = [0]
         if fast: raise ValueError("task head_muscles does not work with option --fast")
     elif task == "headneck_muscles":
         task_id = [778, 779]
@@ -287,7 +287,7 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         crop = ["clavicula_left", "clavicula_right", "vertebrae_C1", "vertebrae_C5", "vertebrae_T1", "vertebrae_T4"]
         crop_addon = [40, 40, 40]
         model = "3d_fullres_high"
-        folds = None
+        folds = [0]
         if fast: raise ValueError("task headneck_muscles does not work with option --fast")
 
     # Commercial models
@@ -307,7 +307,7 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         crop = ["heart"]
         crop_addon = [5, 5, 5]
         model = "3d_fullres"
-        folds = None
+        folds = [0]
         if fast: raise ValueError("task heartchambers_highres does not work with option --fast")
         show_license_info()
     elif task == "appendicular_bones":
@@ -355,6 +355,16 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         folds = [0]
         if fast: raise ValueError("task face_mr does not work with option --fast")
         show_license_info()
+    elif task == "brain_structures":
+        task_id = 409
+        resample = [1.0, 0.5, 0.5]
+        trainer = "nnUNetTrainer_DASegOrd0"
+        crop = ["brain"]
+        crop_addon = [10, 10, 10]
+        model = "3d_fullres_high"
+        folds = [0]
+        if fast: raise ValueError("task brain_structures does not work with option --fast")
+        show_license_info()
     elif task == "test":
         task_id = [517]
         resample = None