add oculomotor muscles subtask

CHANGELOG.md

@@ -2,6 +2,7 @@
 * add option to remove small connected components in postprocessing
 * add `totalseg_get_modality`
 * change pi_time threshold for arterial late phase from 50s to 60s
+* add oculomotor muscles model
 
 ## Release 2.4.0
 * add brain structures

README.md

@@ -70,6 +70,7 @@ Openly available for any usage:
 * **headneck_bones_vessels**: larynx_air, thyroid_cartilage, hyoid, cricoid_cartilage, zygomatic_arch_right, zygomatic_arch_left, styloid_process_right, styloid_process_left, internal_carotid_artery_right, internal_carotid_artery_left, internal_jugular_vein_right, internal_jugular_vein_left (cite [paper](https://www.mdpi.com/2072-6694/16/2/415))
 * **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*
 
 *: These models are not trained on the full totalsegmentator dataset but on some small other datasets. Therefore, expect them to work less robustly.
 

totalsegmentator/bin/TotalSegmentator.py

@@ -60,7 +60,7 @@ def main():
                         "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",
-                        "brain_structures", "liver_vessels"],
+                        "brain_structures", "liver_vessels", "oculomotor_muscles"],
                         help="Select which model to use. This determines what is predicted.",
                         default="total")
 

totalsegmentator/libs.py

@@ -230,6 +230,9 @@ def download_pretrained_weights(task_id):
     elif task_id == 779:
         weights_path = config_dir / "Dataset779_headneck_muscles_part2_492subj"
         WEIGHTS_URL = url + "/v2.3.0-weights/Dataset779_headneck_muscles_part2_492subj.zip"
+    elif task_id == 351:
+        weights_path = config_dir / "Dataset351_oculomotor_muscles_18subj"
+        WEIGHTS_URL = url + "/v2.4.0-weights/Dataset351_oculomotor_muscles_18subj.zip"
 
     # MR models
     elif task_id == 730:

totalsegmentator/map_to_binary.py

@@ -293,6 +293,27 @@
         1: "lung_vessels",
         2: "lung_trachea_bronchia"
     },
+    "oculomotor_muscles": {
+        1: "skull",  # auxiliary class
+        2: "eyeball_right",
+        3: "lateral_rectus_muscle_right",
+        4: "superior_oblique_muscle_right",
+        5: "levator_palpebrae_superioris_right",
+        6: "superior_rectus_muscle_right",
+        7: "medial_rectus_muscle_left",
+        8: "inferior_oblique_muscle_right",
+        9: "inferior_rectus_muscle_right",
+        10: "optic_nerve_left",
+        11: "eyeball_left",
+        12: "lateral_rectus_muscle_left",
+        13: "superior_oblique_muscle_left",
+        14: "levator_palpebrae_superioris_left",
+        15: "superior_rectus_muscle_left",
+        16: "medial_rectus_muscle_right",
+        17: "inferior_oblique_muscle_left",
+        18: "inferior_rectus_muscle_left",
+        19: "optic_nerve_right"
+    },
     "covid": {
         1: "lung_covid_infiltrate",
     },

totalsegmentator/nnunet.py

@@ -249,6 +249,7 @@ def nnUNetv2_predict(dir_in, dir_out, task_id, model="3d_fullres", folds=None,
         use_folds=folds,
         checkpoint_name=chk,
     )
+    # new nnunetv2 feature: keep dir_out empty to return predictions as return value
     predictor.predict_from_files(dir_in, dir_out,
                                  save_probabilities=save_probabilities, overwrite=not continue_prediction,
                                  num_processes_preprocessing=npp, num_processes_segmentation_export=nps,

totalsegmentator/preview.py

@@ -168,6 +168,18 @@
         "sterno_thyroid_right", "sterno_thyroid_left", "thyrohyoid_right", "thyrohyoid_left",
         "prevertebral_right", "prevertebral_left"]
     ],
+    "oculomotor_muscles": [
+        ["skull"], 
+        ["eyeball_right", "eyeball_left", 
+        "levator_palpebrae_superioris_right", "levator_palpebrae_superioris_left", 
+        "superior_rectus_muscle_right", "superior_rectus_muscle_left",
+        "inferior_oblique_muscle_right", "inferior_oblique_muscle_left"],
+        ["lateral_rectus_muscle_right", "lateral_rectus_muscle_left", 
+        "superior_oblique_muscle_right", "superior_oblique_muscle_left",
+        "medial_rectus_muscle_right", "medial_rectus_muscle_left",
+        "inferior_rectus_muscle_right", "inferior_rectus_muscle_left", 
+        "optic_nerve_right", "optic_nerve_left"]
+    ],
     "test": [
         ["ulna"]
     ]

totalsegmentator/python_api.py

@@ -290,6 +290,15 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         model = "3d_fullres_high"
         folds = [0]
         if fast: raise ValueError("task headneck_muscles does not work with option --fast")
+    elif task == "oculomotor_muscles":
+        task_id = 351
+        resample = [0.47251562774181366, 0.47251562774181366, 0.8500002026557922]
+        trainer = "nnUNetTrainer_DASegOrd0_NoMirroring"
+        crop = ["skull"]
+        crop_addon = [20, 20, 20]
+        model = "3d_fullres"
+        folds = [0]
+        if fast: raise ValueError("task oculomotor_muscles does not work with option --fast")
 
     # Commercial models
     elif task == "vertebrae_body":