allow different resampling for all axis

wasserth · Jun 25, 2024 · c6bb6db · c6bb6db
1 parent 6753496
commit c6bb6db
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/totalsegmentator/nnunet.py b/totalsegmentator/nnunet.py
@@ -294,7 +294,7 @@ def nnUNet_predict_image(file_in: Union[str, Path, Nifti1Image], file_out, task_
                          v1_order=False):
     """
     crop: string or a nibabel image
-    resample: None or float  (target spacing for all dimensions)
+    resample: None or float (target spacing for all dimensions) or list of floats
     """
     if not isinstance(file_in, Nifti1Image):
         file_in = Path(file_in)
@@ -320,6 +320,9 @@ def nnUNet_predict_image(file_in: Union[str, Path, Nifti1Image], file_out, task_
         class_map_parts = class_map_parts_mr
         map_taskid_to_partname = map_taskid_to_partname_mr
 
+    if type(resample) is float:
+        resample = [resample, resample, resample]
+
     # for debugging
     # tmp_dir = file_in.parent / ("nnunet_tmp_" + ''.join(random.Random().choices(string.ascii_uppercase + string.digits, k=8)))
     # (tmp_dir).mkdir(exist_ok=True)
@@ -380,7 +383,7 @@ def nnUNet_predict_image(file_in: Union[str, Path, Nifti1Image], file_out, task_
             st = time.time()
             img_in_shape = img_in.shape
             img_in_zooms = img_in.header.get_zooms()
-            img_in_rsp = change_spacing(img_in, [resample, resample, resample],
+            img_in_rsp = change_spacing(img_in, resample,
                                         order=3, dtype=np.int32, nr_cpus=nr_threads_resampling)  # 4 cpus instead of 1 makes it a bit slower
             if verbose:
                 print(f"  from shape {img_in.shape} to shape {img_in_rsp.shape}")
@@ -413,7 +416,7 @@ def nnUNet_predict_image(file_in: Union[str, Path, Nifti1Image], file_out, task_
             nib.save(nib.Nifti1Image(img_in_rsp_data[:, :, third*2+1-margin:], img_in_rsp.affine),
                     tmp_dir / "s03_0000.nii.gz")
 
-        if task_name == "total" and resample < 3.0:
+        if task_name == "total" and resample is not None and resample[0] < 3.0:
             # 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)
@@ -549,7 +552,7 @@ def nnUNet_predict_image(file_in: Union[str, Path, Nifti1Image], file_out, task_
             if verbose: print(f"  back to original shape: {img_in_shape}")
             # Use force_affine otherwise output affine sometimes slightly off (which then is even increased
             # by undo_canonical)
-            img_pred = change_spacing(img_pred, [resample, resample, resample], img_in_shape,
+            img_pred = change_spacing(img_pred, resample, img_in_shape,
                                       order=0, dtype=np.uint8, nr_cpus=nr_threads_resampling,
                                       force_affine=img_in.affine)
 

diff --git a/totalsegmentator/python_api.py b/totalsegmentator/python_api.py
@@ -277,8 +277,10 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         folds = None
         if fast: raise ValueError("task head_muscles does not work with option --fast")
     elif task == "headneck_muscles":
-        task_id = 778
-        resample = None
+        # task_id = 778
+        # resample = None
+        task_id = [778, 779]
+        resample = [1., 0.75, 0.75]
         trainer = "nnUNetTrainer_DASegOrd0_NoMirroring"
         # crop = ["skull", "clavicula_left", "clavicula_right", "vertebrae_C5", "vertebrae_T1", "vertebrae_T4"]
         # crop_addon = [10, 10, 10]