add headneck test case to test locally

wasserth · Jul 2, 2024 · 154d5d7 · 154d5d7
1 parent 2415f48
commit 154d5d7
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tests/test_locally.py b/tests/test_locally.py
@@ -198,7 +198,12 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
         subjects = sorted(list(img_dir.glob("*.nii.gz"))[:1] if debug else list(img_dir.glob("*.nii.gz")))
         for img_fn in tqdm(subjects):
             fast = resolution == "3mm"
-            task = "total_mr" if img_fn.name.split(".")[0].endswith("_mr") else "total"
+            if img_fn.name.split(".")[0].endswith("_mr"):
+                task = "total_mr" 
+            elif img_fn.name.split(".")[0].endswith("_headneck"):
+                task = "headneck_bones_vessels" 
+            else:
+                task = "total"
             st = time.time()
             totalsegmentator(img_fn, pred_dir / img_fn.name, fast=fast, ml=True, task=task, device=device)
             times[resolution].append(time.time()-st)