properly add vertebrae_body model

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## Master
 * Bugfix: add flush to DummyFile
 * Require python >= 3.9 in setup.py
+* properly add `vertebrae_body` model
 
 
 ## Release 2.0.5

README.md

@@ -60,6 +60,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, skeletal_muscle, torso_fat
+* **vertebrae_body**: vertebral body of all vertebrae (without the vertebral arch)
 * **face**: face_region
 
 Usage:

resources/improvements_in_v2.md

@@ -39,7 +39,7 @@ heart_myocardium, heart_atrium_left, heart_ventricle_left, heart_atrium_right, h
 Some of these new classes were available in some preliminary version as additional tasks in v1. Now they are properly added.
 
 The following tasks are freely available only for non-commercial usage (all other tasks can also be used commercially):
-`appendicular_bones`, `tissue_types`, `face`, `heartchambers_highres`
+`appendicular_bones`, `tissue_types`, `face`, `heartchambers_highres`, `vertebrae_body`
 
 
 ## Speed improvements

totalsegmentator/bin/TotalSegmentator.py

@@ -45,22 +45,18 @@ def main():
 
     # cerebral_bleed: Intracerebral hemorrhage 
     # liver_vessels: hepatic vessels
-    parser.add_argument("-ta", "--task", choices=["total", "body", "vertebrae_body",
+    parser.add_argument("-ta", "--task", choices=["total", "body",
                         "lung_vessels", "cerebral_bleed", "hip_implant", "coronary_arteries", 
                         "pleural_pericard_effusion", "test",
                         "appendicular_bones", "tissue_types", "heartchambers_highres", 
-                        "face", 
-                        ],
+                        "face", "vertebrae_body"],
                         # future: liver_vessels, head, 
                         help="Select which model to use. This determines what is predicted.",
                         default="total")
 
     parser.add_argument("-rs", "--roi_subset", type=str, nargs="+",
                         help="Define a subset of classes to save (space separated list of class names). If running 1.5mm model, will only run the appropriate models for these rois.")
 
-    # When this is used together with --roi_subset and --ml, then statistics will be calculated for all classes
-    # in the class_map_part_X with is calculated. Some of the rois will be cropped and therefore the volume will
-    # only be of the cropped region, making it an incorrect volume.
     parser.add_argument("-s", "--statistics", action="store_true", 
                         help="Calc volume (in mm3) and mean intensity. Results will be in statistics.json",
                         default=False)

totalsegmentator/bin/totalseg_download_weights.py

@@ -35,11 +35,11 @@ def main():
         "pleural_pericard_effusion": [315],
         "body": [299],
         "body_fast": [300],
-        "vertebrae_body": [302],
 
         "heartchambers_highres": [301],
         "appendicular_bones": [304],
         "tissue_types": [481],
+        "vertebrae_body": [302],
         "face": [303],
 
         # "liver_vessels": [8]

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, 302]:
+                    315, 299, 300]:
         download_pretrained_weights(task_id)
         sleep(5)

totalsegmentator/libs.py

@@ -215,10 +215,6 @@ 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"
-    elif task_id == 302:
-        weights_path = config_dir / "Dataset302_vertebrae_body_1559subj"
-        # WEIGHTS_URL = url + "/static/totalseg_v2/Dataset302_vertebrae_body_1559subj.zip"
-        WEIGHTS_URL = url + "/v2.0.0-weights/Dataset302_vertebrae_body_1559subj.zip"
 
     # Models from other projects 
     elif task_id == 258:
@@ -269,6 +265,9 @@ def download_pretrained_weights(task_id):
         weights_path = config_dir / "Dataset303_face_1559subj"
     elif task_id == 481:
         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"
 
     else:
         raise ValueError(f"For task_id {task_id} no download path was found.")

totalsegmentator/map_to_binary.py

@@ -308,6 +308,7 @@
     "heartchambers_highres": 301,
     "appendicular_bones": 304, 
     "tissue_types": 481, 
+    "vertebrae_body": 302,
     "face": 303
 }
 # future

totalsegmentator/python_api.py

@@ -173,6 +173,8 @@ def totalsegmentator(input, output, ml=False, nr_thr_resamp=1, nr_thr_saving=6,
         model = "3d_fullres"
         folds = None
         if fast: raise ValueError("task liver_vessels does not work with option --fast")
+
+    # Commercial models
     elif task == "vertebrae_body":
         task_id = 302
         resample = 1.5
@@ -181,8 +183,7 @@ def totalsegmentator(input, output, ml=False, nr_thr_resamp=1, nr_thr_saving=6,
         model = "3d_fullres"
         folds = [0]
         if fast: raise ValueError("task vertebrae_body does not work with option --fast")
-
-    # Commercial models
+        show_license_info()
     elif task == "heartchambers_highres":
         task_id = 301
         resample = None