bugfix in types in test_locally

wasserth · Mar 18, 2024 · 5140cfd · 5140cfd
1 parent bd9e293
commit 5140cfd
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/tests/test_locally.py b/tests/test_locally.py
@@ -149,7 +149,8 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
         elif isinstance(old_value, pd.Timestamp) and isinstance(new_value, pd.Timestamp):
             continue
         # Check numeric values for similarity within a tolerance
-        elif isinstance(old_value, (int, float)) and isinstance(new_value, (int, float)):
+        elif isinstance(old_value, (int, float, np.integer, np.floating)) and \
+             isinstance(new_value, (int, float, np.integer, np.floating)):
             if old_value == 0 and new_value == 0:
                 continue
             elif old_value == 0 or new_value == 0:
@@ -182,6 +183,8 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
 
     device = "gpu"  # "cpu" or "gpu"
 
+    debug = False
+
     for resolution in ["15mm", "3mm"]:
     # for resolution in ["3mm"]:
         print(f"----- resolution {resolution} ------")
@@ -192,8 +195,8 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
 
         print("Run totalsegmentator...")
         reset_monitors()
-        for img_fn in tqdm(img_dir.glob("*.nii.gz")):
-        # for img_fn in tqdm(list(img_dir.glob("*.nii.gz"))[:1]):
+        subjects = 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"
             st = time.time()
             totalsegmentator(img_fn, pred_dir / img_fn.name, fast=fast, ml=True, device=device)
@@ -207,7 +210,7 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
         gpu_utilization[resolution] = float(np.mean(gpu_utilizations).round(1))
 
         print("Calc metrics...")
-        subjects = [s.name.split(".")[0] for s in img_dir.glob("*.nii.gz")]
+        subjects = [s.name.split(".")[0] for s in subjects]
         res = [calc_metrics(s, gt_dir, pred_dir, class_map["total"]) for s in subjects]
         res = pd.DataFrame(res)
 
@@ -259,9 +262,9 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
 
     print("Comparing PREVIOUS to NEW log:")
     if are_logs_similar(last_log, new_log, cols):
-        print("SUCCESS: no differences")
+        print("\nSUCCESS: no differences\n")
     else:
-        print("ERROR: major differences found")
+        print("\nERROR: major differences found\n")
 
     print(f"Saving to {overview_file}...")
     overview.loc[len(overview)] = new_log

diff --git a/totalsegmentator/bin/totalseg_get_phase.py b/totalsegmentator/bin/totalseg_get_phase.py
@@ -89,9 +89,9 @@ def get_ct_contrast_phase(ct_img: nib.Nifti1Image):
 
 def main():
     """
-    The the contrast phase of a CT scan. Specifically this script will predict the
+    Predicts the contrast phase of a CT scan. Specifically this script will predict the
     pi (post injection) time (in seconds) of a CT scan based on the intensity of different regions
-    in the image
+    in the image. 
     """
     parser = argparse.ArgumentParser(description="Get CT contrast phase.",
                                      epilog="Written by Jakob Wasserthal. If you use this tool please cite https://pubs.rsna.org/doi/10.1148/ryai.230024")