add tissue_4_types model

wasserth · Jan 16, 2025 · 044a1d8 · 044a1d8
1 parent 11350a9
commit 044a1d8
Show file tree

Hide file tree

Showing 7 changed files with 26 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -84,6 +84,7 @@ Available with a license (free licenses available for non-commercial usage [here
 * **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)
+* **tissue_4_types**: subcutaneous_fat, torso_fat, skeletal_muscle, intermuscular_fat (in contrast to `tissue_types` skeletal_muscle is split into two classes: muscle and fat)
 * **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), intervertebral_discs (for MR this is part of the `total_mr` task)
 * **face**: face_region (for anonymization)

diff --git a/resources/prepare_weights_for_release.sh b/resources/prepare_weights_for_release.sh
@@ -22,7 +22,7 @@ print(get_full_task_name_v2($dataset_id))
     echo "Processing $task_name..."
 
     # Copy dataset folder
-    cp -r "$nnUNet_results/$task_name" .
+    # cp -r "$nnUNet_results/$task_name" .
 
     # Get the only folder inside task_name
     trainer_folder=$(ls "$task_name" | head -n 1)

diff --git a/totalsegmentator/bin/TotalSegmentator.py b/totalsegmentator/bin/TotalSegmentator.py
@@ -58,7 +58,7 @@ def main():
                         "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_body", "total_mr", "tissue_types_mr", "tissue_4_types", "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", "kidney_cysts", 

diff --git a/totalsegmentator/bin/totalseg_download_weights.py b/totalsegmentator/bin/totalseg_download_weights.py
@@ -23,7 +23,7 @@ def main():
                                                  "body", "body_fast", "body_mr", "body_mr_fast", "vertebrae_mr",
                                                  "vertebrae_body",
                                                  "heartchambers_highres", "appendicular_bones", 
-                                                 "tissue_types", "tissue_types_mr", "face", "face_mr",
+                                                 "tissue_types", "tissue_types_mr", "tissue_4_types", "face", "face_mr",
                                                  "head_glands_cavities", "head_muscles", "headneck_bones_vessels",
                                                  "headneck_muscles", "liver_vessels", "brain_structures",
                                                  "lung_nodules", "kidney_cysts", "breasts",
@@ -60,6 +60,7 @@ def main():
         "appendicular_bones": [304],
         "tissue_types": [481],
         "tissue_types_mr": [854],
+        "tissue_4_types": [485],
         "vertebrae_body": [305],
         "face": [303],
         "face_mr": [856],

diff --git a/totalsegmentator/map_to_binary.py b/totalsegmentator/map_to_binary.py
@@ -420,6 +420,12 @@
         2: "torso_fat",
         3: "skeletal_muscle"
     },
+    "tissue_4_types": {
+        1: "subcutaneous_fat",
+        2: "torso_fat",
+        3: "skeletal_muscle",
+        4: "intermuscular_fat"
+    },
     "face": {
         1: "face"
     },

diff --git a/totalsegmentator/preview.py b/totalsegmentator/preview.py
@@ -135,6 +135,12 @@
         ["torso_fat"],
         ["skeletal_muscle"]
     ],
+    "tissue_4_types": [
+        ["subcutaneous_fat"],
+        ["torso_fat"],
+        ["skeletal_muscle"],
+        ["intermuscular_fat"]
+    ],
     "face": [
         ["face"]
     ],

diff --git a/totalsegmentator/python_api.py b/totalsegmentator/python_api.py
@@ -408,6 +408,15 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         folds = [0]
         if fast: raise ValueError("task tissue_types_mr does not work with option --fast")
         show_license_info()
+    elif task == "tissue_4_types":
+        task_id = 485
+        resample = 1.5
+        trainer = "nnUNetTrainer"
+        crop = None
+        model = "3d_fullres"
+        folds = [0]
+        if fast: raise ValueError("task tissue_4_types does not work with option --fast")
+        show_license_info()
     elif task == "face":
         task_id = 303
         resample = 1.5