add mr all

CHANGELOG.md

@@ -1,8 +1,12 @@
 ## Master
+
+
+## Release 2.2.0
 * also return statistics from python api
 * add `totalseg_get_phase`
 * major bugfix: rib labels were in wrong order
 * hide nnunetv2 2.3.1 warning: `Detected old nnU-Net plans format. Attempting to reconstruct network architecture...`
+* add mr models
 
 
 ## Release 2.1.0

README.md

@@ -47,6 +47,7 @@ Next to the default task (`total`) there are more subtasks with more classes:
 
 Openly available for any usage:
 * **total**: default task containing 117 main classes (see [here](https://github.com/wasserth/TotalSegmentator#class-details) for a list of classes)
+* **total_mr**: default task containing 56 main classes on MR images (see [here](https://github.com/wasserth/TotalSegmentator#class-details) for a list of classes)
 * **lung_vessels**: lung_vessels (cite [paper](https://www.sciencedirect.com/science/article/pii/S0720048X22001097)), lung_trachea_bronchia
 * **body**: body, body_trunc, body_extremities, skin
 * **cerebral_bleed**: intracerebral_hemorrhage (cite [paper](https://www.mdpi.com/2077-0383/12/7/2631))*
@@ -60,6 +61,7 @@ Available with a license (free licenses available for non-commercial usage [here
 * **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
 * **tissue_types**: subcutaneous_fat, torso_fat, skeletal_muscle
+* **tissue_types_mr**: subcutaneous_fat, torso_fat, skeletal_muscle (for MR images)
 * **vertebrae_body**: vertebral body of all vertebrae (without the vertebral arch)
 * **face**: face_region
 
@@ -210,7 +212,7 @@ Moreover, we would really appreciate it if you let us know what you are using th
 
 ### Class details
 
-The following table shows a list of all classes.
+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.
@@ -336,3 +338,66 @@ For some classes they differ which you can see in the table below.
 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 ||

tests/test_end_to_end.py

@@ -45,6 +45,13 @@ def test_prediction_multilabel(self):
         images_equal = dice > 0.99
         self.assertTrue(images_equal, f"multilabel prediction not correct (dice: {dice:.6f})")
 
+    def test_prediction_multilabel_mr(self):
+        img_ref = nib.load("tests/reference_files/example_seg_mr.nii.gz").get_fdata()
+        img_new = nib.load("tests/unittest_prediction_mr.nii.gz").get_fdata()
+        dice = dice_score_multilabel(img_ref, img_new)
+        images_equal = dice > 0.99
+        self.assertTrue(images_equal, f"multilabel prediction MR not correct (dice: {dice:.6f})")
+
     def test_prediction_liver_roi_subset(self):
         img_ref = nib.load("tests/reference_files/example_seg_roi_subset.nii.gz").get_fdata()
         img_new = nib.load("tests/unittest_prediction_roi_subset.nii.gz").get_fdata()

tests/tests_subtasks.sh

@@ -8,6 +8,10 @@ set -e
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/unittest_prediction -ta lung_vessels -d cpu  # ~1min
 pytest -v tests/test_end_to_end.py::test_end_to_end::test_lung_vessels
 
+# Test total_mr
+TotalSegmentator -i tests/reference_files/example_mr_sm.nii.gz -o tests/unittest_prediction_mr.nii.gz -ta total_mr --ml -d cpu
+pytest -v tests/test_end_to_end.py::test_end_to_end::test_prediction_multilabel_mr
+
 # Test tissue types (without license)
 #   (use "|| true" to not abort if this command returns exit code 1, which it should do)
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/unittest_no_license.nii.gz -ta tissue_types -d cpu --ml || true
@@ -27,3 +31,4 @@ pytest -v tests/test_end_to_end.py::test_end_to_end::test_appendicular_bones
 
 # Cleanup generated files and directories
 rm -rf tests/unittest_prediction
+rm tests/unittest_prediction_mr.nii.gz

tests/update_test_files.sh

@@ -6,6 +6,7 @@ set -e
 # ./tests/update_test_files.sh <license_key>
 
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/reference_files/example_seg.nii.gz -bs --ml -d cpu
+TotalSegmentator -i tests/reference_files/example_mr_sm.nii.gz -o tests/reference_files/example_seg_mr.nii.gz -ta total_mr --ml -d cpu
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/reference_files/example_seg_roi_subset.nii.gz --ml -rs liver brain -d cpu
 # TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/reference_files/example_seg_fast --fast --statistics -sii -p -d cpu
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/reference_files/example_seg_fast.nii.gz --fast --ml -d cpu

totalsegmentator/bin/TotalSegmentator.py

@@ -52,7 +52,7 @@ def main():
                         "lung_vessels", "cerebral_bleed", "hip_implant", "coronary_arteries",
                         "pleural_pericard_effusion", "test",
                         "appendicular_bones", "tissue_types", "heartchambers_highres",
-                        "face", "vertebrae_body"],
+                        "face", "vertebrae_body", "total_mr", "tissue_types_mr"],
                         # future: liver_vessels, head,
                         help="Select which model to use. This determines what is predicted.",
                         default="total")

totalsegmentator/bin/totalseg_download_weights.py

@@ -17,17 +17,21 @@ def main():
     parser = argparse.ArgumentParser(description="Import manually downloaded weights.",
                                      epilog="Written by Jakob Wasserthal.")
 
-    parser.add_argument("-t", "--task", choices=["total", "total_fast", "lung_vessels", "cerebral_bleed",
+    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", "vertebrae_body",
-                                                 "heartchambers_highres", "appendicular_bones", "tissue_types", "face"],
+                                                 "heartchambers_highres", "appendicular_bones", 
+                                                 "tissue_types", "tissue_types_mr", "face"],
                         help="Task for which to download the weights", default="total")
 
     args = parser.parse_args()
 
     task_to_id = {
         "total": [291, 292, 293, 294, 295, 298],
         "total_fast": [297, 298],
+        "total_mr": [730, 731],
+        "total_fast_mr": [732, 733],
         "lung_vessels": [258],
         "cerebral_bleed": [150],
         "hip_implant": [260],
@@ -39,6 +43,7 @@ def main():
         "heartchambers_highres": [301],
         "appendicular_bones": [304],
         "tissue_types": [481],
+        "tissue_types_mr": [734],
         "vertebrae_body": [302],
         "face": [303],
 

totalsegmentator/download_pretrained_weights.py

@@ -7,6 +7,6 @@
     Download all pretrained weights
     """
     for task_id in [291, 292, 293, 294, 295, 297, 298, 258, 150, 260, 503,
-                    315, 299, 300]:
+                    315, 299, 300, 730, 731, 732, 733, 734]:
         download_pretrained_weights(task_id)
         sleep(5)

totalsegmentator/libs.py

@@ -215,6 +215,20 @@ def download_pretrained_weights(task_id):
         # WEIGHTS_URL = "https://zenodo.org/record/7334272/files/Task269_Body_extrem_6mm_1200subj.zip?download=1"
         # WEIGHTS_URL = url + "/static/totalseg_v2/Dataset300_body_6mm_1559subj.zip"
         WEIGHTS_URL = url + "/v2.0.0-weights/Dataset300_body_6mm_1559subj.zip"
+
+    # MR models
+    elif task_id == 730:
+        weights_path = config_dir / "Dataset730_TotalSegmentatorMRI_part1_organs_495subj"
+        WEIGHTS_URL = url + "/v2.2.0-weights/Dataset730_TotalSegmentatorMRI_part1_organs_495subj.zip"
+    elif task_id == 731:
+        weights_path = config_dir / "Dataset731_TotalSegmentatorMRI_part2_muscles_495subj"
+        WEIGHTS_URL = url + "/v2.2.0-weights/Dataset731_TotalSegmentatorMRI_part2_muscles_495subj.zip"
+    elif task_id == 732:
+        weights_path = config_dir / "Dataset732_TotalSegmentatorMRI_total_3mm_495subj"
+        WEIGHTS_URL = url + "/v2.2.0-weights/Dataset732_TotalSegmentatorMRI_total_3mm_495subj.zip"
+    elif task_id == 733:
+        weights_path = config_dir / "Dataset733_TotalSegmentatorMRI_total_6mm_495subj"
+        WEIGHTS_URL = url + "/v2.2.0-weights/Dataset733_TotalSegmentatorMRI_total_6mm_495subj.zip"
 
     # Models from other projects
     elif task_id == 258:
@@ -267,7 +281,8 @@ def download_pretrained_weights(task_id):
         weights_path = config_dir / "Dataset481_tissue_1559subj"
     elif task_id == 302:
         weights_path = config_dir / "Dataset302_vertebrae_body_1559subj"
-        # WEIGHTS_URL = url + "/v2.0.0-weights/Dataset302_vertebrae_body_1559subj.zip"
+    elif task_id == 734:
+        weights_path = config_dir / "Dataset734_TotalSegmentatorMRI_tissue_495subj"
 
     else:
         raise ValueError(f"For task_id {task_id} no download path was found.")

totalsegmentator/map_to_binary.py

@@ -373,6 +373,7 @@
     "heartchambers_highres": 301,
     "appendicular_bones": 304,
     "tissue_types": 481,
+    "tissue_types_mr": 734,
     "vertebrae_body": 302,
     "face": 303
 }
@@ -607,6 +608,9 @@
     294: "class_map_part_muscles",
     295: "class_map_part_ribs",
 
+    730: "class_map_part_organs",
+    731: "class_map_part_muscles",
+
     517: "test",
 }
 

totalsegmentator/nnunet.py

@@ -31,7 +31,7 @@
 
 from nnunetv2.utilities.file_path_utilities import get_output_folder
 
-from totalsegmentator.map_to_binary import class_map, class_map_5_parts, map_taskid_to_partname
+from totalsegmentator.map_to_binary import class_map, class_map_5_parts, map_taskid_to_partname, class_map_parts_mr
 from totalsegmentator.alignment import as_closest_canonical_nifti, undo_canonical_nifti
 from totalsegmentator.alignment import as_closest_canonical, undo_canonical
 from totalsegmentator.resampling import change_spacing
@@ -307,6 +307,11 @@ def nnUNet_predict_image(file_in: Union[str, Path, Nifti1Image], file_out, task_
     if img_type == "nifti" and output_type == "dicom":
         raise ValueError("To use output type dicom you also have to use a Dicom image as input.")
 
+    if task_name == "total":
+        class_map_parts = class_map_5_parts
+    elif task_name == "total_mr":
+        class_map_parts = class_map_parts_mr
+
     # for debugging
     # tmp_dir = file_in.parent / ("nnunet_tmp_" + ''.join(random.Random().choices(string.ascii_uppercase + string.digits, k=8)))
     # (tmp_dir).mkdir(exist_ok=True)
@@ -416,7 +421,7 @@ def nnUNet_predict_image(file_in: Union[str, Path, Nifti1Image], file_out, task_
             if roi_subset is not None:
                 part_names = []
                 new_task_id = []
-                for part_name, part_map in class_map_5_parts.items():
+                for part_name, part_map in class_map_parts.items():
                     if any(organ in roi_subset for organ in part_map.values()):
                         # get taskid associated to model part_name
                         map_partname_to_taskid = {v:k for k,v in map_taskid_to_partname.items()}
@@ -447,7 +452,7 @@ def nnUNet_predict_image(file_in: Union[str, Path, Nifti1Image], file_out, task_
                     for img_part in img_parts:
                         (tmp_dir / f"{img_part}.nii.gz").rename(tmp_dir / "parts" / f"{img_part}_{tid}.nii.gz")
                         seg = nib.load(tmp_dir / "parts" / f"{img_part}_{tid}.nii.gz").get_fdata()
-                        for jdx, class_name in class_map_5_parts[map_taskid_to_partname[tid]].items():
+                        for jdx, class_name in class_map_parts[map_taskid_to_partname[tid]].items():
                             seg_combined[img_part][seg == jdx] = class_map_inv[class_name]
                 # iterate over subparts of image
                 for img_part in img_parts:

totalsegmentator/preview.py

@@ -58,6 +58,27 @@
         "pulmonary_vein",
          "superior_vena_cava", "brachiocephalic_vein_left", "brachiocephalic_vein_right"]
     ],
+    "total_mr": [
+        ["humerus_left", "humerus_right", "femur_left", "femur_right", 
+        "hip_left", "hip_right", "sacrum",
+        "vertebrae", "intervertebral_discs",
+        "fibula", "tibia",
+        "autochthon_left", "autochthon_right", "iliopsoas_left", "iliopsoas_right",
+        "gluteus_medius_left", "gluteus_medius_right", "gluteus_minimus_left", "gluteus_minimus_right",
+        "gluteus_maximus_left", "gluteus_maximus_right",
+        "quadriceps_femoris_left", "quadriceps_femoris_right", 
+        "thigh_medial_compartment_left", "thigh_medial_compartment_right", 
+        "thigh_posterior_compartment_left", "thigh_posterior_compartment_right", 
+        "sartorius_left", "sartorius_right"],
+        ["iliac_artery_left", "iliac_artery_right", "iliac_vena_left", "iliac_vena_right",
+        "aorta", "inferior_vena_cava", "portal_vein_and_splenic_vein",
+        "heart", "esophagus", "stomach", "duodenum", "colon", "small_bowel", "urinary_bladder"],
+        ["lung_left", "lung_right", "liver",
+        "spleen", "gallbladder", "pancreas", 
+        "kidney_right", "kidney_left",
+        "adrenal_gland_right", "adrenal_gland_left", 
+        "brain", "prostate", "spinal_cord"],
+    ],
     "lung_vessels": [
         ["lung_trachea_bronchia"],
         ["lung_vessels"]
@@ -101,6 +122,11 @@
         ["torso_fat"],
         ["skeletal_muscle"]
     ],
+    "tissue_types_mr": [
+        ["subcutaneous_fat"],
+        ["torso_fat"],
+        ["skeletal_muscle"]
+    ],
     "face": [
         ["face"]
     ],

totalsegmentator/python_api.py

@@ -110,6 +110,27 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
             crop = None
         model = "3d_fullres"
         folds = [0]
+    elif task == "total_mr":
+        if fast:
+            task_id = 732
+            resample = 3.0
+            trainer = "nnUNetTrainer_DASegOrd0_NoMirroring"
+            # trainer = "nnUNetTrainerNoMirroring"
+            crop = None
+            if not quiet: print("Using 'fast' option: resampling to lower resolution (3mm)")
+        elif fastest:
+            task_id = 733
+            resample = 6.0
+            trainer = "nnUNetTrainer_DASegOrd0_NoMirroring"
+            crop = None
+            if not quiet: print("Using 'fastest' option: resampling to lower resolution (6mm)")
+        else:
+            task_id = [730, 731]
+            resample = 1.5
+            trainer = "nnUNetTrainer_DASegOrd0_NoMirroring"
+            crop = None
+        model = "3d_fullres"
+        folds = [0]
     elif task == "lung_vessels":
         task_id = 258
         resample = None
@@ -229,6 +250,15 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         folds = [0]
         if fast: raise ValueError("task tissue_types does not work with option --fast")
         show_license_info()
+    elif task == "tissue_types_mr":
+        task_id = 734
+        resample = 1.5
+        trainer = "nnUNetTrainer_DASegOrd0_NoMirroring"
+        crop = None
+        model = "3d_fullres"
+        folds = [0]
+        if fast: raise ValueError("task tissue_types_mr does not work with option --fast")
+        show_license_info()
     elif task == "face":
         task_id = 303
         resample = 1.5
@@ -277,6 +307,10 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
     if roi_subset is not None and task != "total":
         raise ValueError("roi_subset only works with task 'total'")
 
+    if task == "total_mr" or task == "tissue_types_mr":
+        body_seg = False
+        print("INFO: For MR models the argument '--body_seg' is not supported and will be ignored.")
+
     # Generate rough organ segmentation (6mm) for speed up if crop or roi_subset is used
     # (for "fast" on GPU it makes no big difference, but on CPU it can help even for "fast")
     if crop is not None or roi_subset is not None:
@@ -286,14 +320,14 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         st = time.time()
         if not quiet: print("Generating rough body segmentation...")
         if robust_rs:
-            crop_model_task = 297
+            crop_model_task = 732 if task == "total_mr" else 297
             crop_spacing = 3.0
         else:
-            crop_model_task = 298
+            crop_model_task = 733 if task == "total_mr" else 298
             crop_spacing = 6.0
         organ_seg, _, _ = nnUNet_predict_image(input, None, crop_model_task, model="3d_fullres", folds=[0],
                             trainer="nnUNetTrainer_4000epochs_NoMirroring", tta=False, multilabel_image=True, resample=crop_spacing,
-                            crop=None, crop_path=None, task_name="total", nora_tag="None", preview=False,
+                            crop=None, crop_path=None, task="total", nora_tag="None", preview=False,
                             save_binary=False, nr_threads_resampling=nr_thr_resamp, nr_threads_saving=1,
                             crop_addon=crop_addon, output_type=output_type, statistics=False,
                             quiet=quiet, verbose=verbose, test=0, skip_saving=False, device=device)