update download tool

wasserth · Aug 29, 2023 · d19b604 · d19b604
1 parent 0a38b18
commit d19b604
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 15 deletions.
diff --git a/bin/totalseg_download_weights b/bin/totalseg_download_weights
@@ -10,27 +10,38 @@ from totalsegmentator.config import setup_totalseg
 if __name__ == "__main__":
     """
     Download totalsegmentator weights
+
+    Info: If want to download models with require a license you have to run `totalseg_set_license` first.
     """
     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",
                                                  "hip_implant", "coronary_arteries", "body", "pleural_pericard_effusion",
-                                                 "liver_vessels", "bones_extremities"],
+                                                 "body", "body_fast", "vertebrae_body", 
+                                                 "heartchambers_highres", "appendicular_bones", "tissue_types", "face"],
                         help="Task for which to download the weights", default="total")
 
     args = parser.parse_args()
 
     task_to_id = {
-        "total": [251, 252, 253, 254, 255],
-        "total_fast": [256],
+        "total": [291, 292, 293, 294, 295],
+        "total_fast": [297],
         "lung_vessels": [258],
-        "covid": [201],
         "cerebral_bleed": [150],
         "hip_implant": [260],
         "coronary_arteries": [503],
         "pleural_pericard_effusion": [315],
-        "liver_vessels": [8]
+        "body": [299],
+        "body_fast": [300],
+        "vertebrae_body": [302],
+
+        "heartchambers_highres": [301],
+        "appendicular_bones": [296],
+        "tissue_types": [481],
+        "face": [303],
+
+        # "liver_vessels": [8]
     }
 
     setup_totalseg()

diff --git a/totalsegmentator/python_api.py b/totalsegmentator/python_api.py
@@ -96,16 +96,16 @@ def totalsegmentator(input, output, ml=False, nr_thr_resamp=1, nr_thr_saving=6,
         if fast: raise ValueError("task lung_vessels does not work with option --fast")
         model = "3d_fullres"
         folds = [0]
-    elif task == "covid":
-        task_id = 201
-        resample = None
-        trainer = "nnUNetTrainer"
-        crop = ["lung_upper_lobe_left", "lung_lower_lobe_left", "lung_upper_lobe_right",
-                "lung_middle_lobe_right", "lung_lower_lobe_right"]
-        model = "3d_fullres"
-        folds = [0]
-        print("WARNING: The COVID model finds many types of lung opacity not only COVID. Use with care!")
-        if fast: raise ValueError("task covid does not work with option --fast")
+    # elif task == "covid":
+    #     task_id = 201
+    #     resample = None
+    #     trainer = "nnUNetTrainer"
+    #     crop = ["lung_upper_lobe_left", "lung_lower_lobe_left", "lung_upper_lobe_right",
+    #             "lung_middle_lobe_right", "lung_lower_lobe_right"]
+    #     model = "3d_fullres"
+    #     folds = [0]
+    #     print("WARNING: The COVID model finds many types of lung opacity not only COVID. Use with care!")
+    #     if fast: raise ValueError("task covid does not work with option --fast")
     elif task == "cerebral_bleed":
         task_id = 150
         resample = None