From ff50878153342c7b4cb8ae466f7d98aadde4797d Mon Sep 17 00:00:00 2001 From: jakob Date: Thu, 17 Oct 2024 14:54:23 +0200 Subject: [PATCH] add oculomotor muscles subtask --- CHANGELOG.md | 1 + README.md | 1 + totalsegmentator/bin/TotalSegmentator.py | 2 +- totalsegmentator/libs.py | 3 +++ totalsegmentator/map_to_binary.py | 21 +++++++++++++++++++++ totalsegmentator/nnunet.py | 1 + totalsegmentator/preview.py | 12 ++++++++++++ totalsegmentator/python_api.py | 9 +++++++++ 8 files changed, 49 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c92aa03b6..91e772eb7 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index c565357b7..41b29a27e 100644 --- a/README.md +++ b/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. diff --git a/totalsegmentator/bin/TotalSegmentator.py b/totalsegmentator/bin/TotalSegmentator.py index c46b9cd04..ba34ae1b2 100644 --- a/totalsegmentator/bin/TotalSegmentator.py +++ b/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") diff --git a/totalsegmentator/libs.py b/totalsegmentator/libs.py index c3bc829d2..cacc8359a 100644 --- a/totalsegmentator/libs.py +++ b/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: diff --git a/totalsegmentator/map_to_binary.py b/totalsegmentator/map_to_binary.py index d56358551..a04bcb232 100644 --- a/totalsegmentator/map_to_binary.py +++ b/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", }, diff --git a/totalsegmentator/nnunet.py b/totalsegmentator/nnunet.py index a69c325b1..627f490c5 100644 --- a/totalsegmentator/nnunet.py +++ b/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, diff --git a/totalsegmentator/preview.py b/totalsegmentator/preview.py index ea4b9b112..0b0bae058 100644 --- a/totalsegmentator/preview.py +++ b/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"] ] diff --git a/totalsegmentator/python_api.py b/totalsegmentator/python_api.py index eeb12b58a..74ff0c32a 100644 --- a/totalsegmentator/python_api.py +++ b/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":