update local tests for mr

wasserth · May 29, 2024 · 83c4242 · 83c4242
1 parent 4d6fa99
commit 83c4242
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/test_locally.py b/tests/test_locally.py
@@ -195,11 +195,12 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
 
         print("Run totalsegmentator...")
         reset_monitors()
-        subjects = list(img_dir.glob("*.nii.gz"))[:1] if debug else list(img_dir.glob("*.nii.gz"))
+        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"
             st = time.time()
-            totalsegmentator(img_fn, pred_dir / img_fn.name, fast=fast, ml=True, device=device)
+            totalsegmentator(img_fn, pred_dir / img_fn.name, fast=fast, ml=True, task=task, device=device)
             times[resolution].append(time.time()-st)
 
         print("Logging...")