From d657c52136b9af9f6f10f54e7e3d56b00a9a0fe2 Mon Sep 17 00:00:00 2001
From: wasserth <j.wasserthal@gmx.de>
Date: Thu, 15 Feb 2024 12:37:53 +0100
Subject: [PATCH] change step_size to 0.8 for faster runtime

---
 CHANGELOG.md               | 1 +
 tests/test_locally.py      | 8 +++++---
 totalsegmentator/nnunet.py | 7 +++++--
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 83162618d..f2769b676 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@
 * input to python api can be a Nifti1Image object or a file path
 * upgrade to `nnunetv2>=2.2.1`
 * minor edits and bugfixes
+* use nnU-Net `step_size=0.8` instead of `0.5` for faster runtime while only decreasing dice by 0.001
 
 
 ## Release 2.0.5
diff --git a/tests/test_locally.py b/tests/test_locally.py
index 8865267fe..0bb0be561 100755
--- a/tests/test_locally.py
+++ b/tests/test_locally.py
@@ -181,6 +181,8 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
     cpu_utilization = {}
     gpu_utilization = {}
 
+    device = "gpu"  # "cpu" or "gpu"
+
     for resolution in ["15mm", "3mm"]:
     # for resolution in ["3mm"]:
         img_dir = base_dir / resolution / "ct"
@@ -193,7 +195,7 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
         for img_fn in tqdm(img_dir.glob("*.nii.gz")):
             fast = resolution == "3mm"
             st = time.time()
-            totalsegmentator(img_fn, pred_dir / img_fn.name, fast=fast, ml=True, device="gpu")
+            totalsegmentator(img_fn, pred_dir / img_fn.name, fast=fast, ml=True, device=device)
             times[resolution].append(time.time()-st)
 
         print("Logging...")
@@ -229,7 +231,7 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
             "gpu_utilization_15mm", "gpu_utilization_3mm",
             "python_version", "torch_version", "nnunet_version",
             "cuda_version", "cudnn_version",
-            "gpu_name"]
+            "gpu_name", "comment"]
     overview_file = Path(f"{base_dir}/overview.xlsx")
     if overview_file.exists():
         overview = pd.read_excel(overview_file)
@@ -248,7 +250,7 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
                platform.python_version(), torch.__version__,
                importlib.metadata.version("nnunetv2"),
                float(torch.version.cuda), int(torch.backends.cudnn.version()),
-               torch.cuda.get_device_name(0)]
+               torch.cuda.get_device_name(0), ""]
 
     print("Comparing NEW to PREVIOUS log:")
     if are_logs_similar(last_log, new_log, cols):
diff --git a/totalsegmentator/nnunet.py b/totalsegmentator/nnunet.py
index e9b1a378d..16247057b 100644
--- a/totalsegmentator/nnunet.py
+++ b/totalsegmentator/nnunet.py
@@ -179,8 +179,11 @@ def nnUNetv2_predict(dir_in, dir_out, task_id, model="3d_fullres", folds=None,
         device = torch.device('cuda')
     else:
         device = torch.device('mps')
-    step_size = 0.5
-    # step_size = 0.8  # overall speedup roughly 11%; for fast model no speedup; dice 0.001 worse
+    # step_size = 0.5
+    # overall speedup for 15mm model roughly 11% (GPU) and 100% (CPU) 
+    # overall speedup for  3mm model roughly  0% (GPU) and  10% (CPU)
+    # (dice 0.001 worse on test set -> ok)
+    step_size = 0.8
     disable_tta = not tta
     verbose = False
     save_probabilities = False