[nnUNet/TF2] Fix sliding window evaluation for non cubic volumes

Author: michal2409

TensorFlow2/Segmentation/nnUNet/main.py

@@ -37,7 +37,8 @@ def main(args):
     assert pValue.contents.value == 128
 
     hvd_init()
-    set_seed(args.seed)
+    if args.seed is not None:
+        set_seed(args.seed)
     set_tf_flags(args)
     data = DataModule(args)
     data.setup()

TensorFlow2/Segmentation/nnUNet/models/sliding_window.py

@@ -27,20 +27,8 @@ def batch_window_slices(slices, image_batch_size, batch_size):
         )
     return batched_window_slices
 
-@tf.function
-def gaussian_kernel_tf_v2(roi_size, sigma):
-        """
-        adapted from: https://gist.github.com/blzq
-        """
-        kernel_size = roi_size[0]
-        sigma = sigma * kernel_size
-        gauss = tf.range(start = 0, limit = kernel_size, dtype = tf.float32) - (kernel_size - 1.0) / 2.0
-        xx, yy, zz = tf.meshgrid(gauss, gauss, gauss)
-        kernel = tf.exp(-(xx ** 2 + yy ** 2 + zz ** 2) / (2.0 * sigma ** 2))
-        kernel = tf.math.pow(kernel, 1/len(roi_size))
-        kernel = kernel / tf.reduce_max(kernel)
-        return kernel
 
+@tf.function
 def gaussian_kernel(roi_size, sigma):
     gauss = signal.windows.gaussian(roi_size[0], std=sigma * roi_size[0])
     for s in roi_size[1:]:
@@ -57,7 +45,7 @@ def get_importance_kernel(roi_size, blend_mode, sigma):
     if blend_mode == "constant":
         return tf.ones(roi_size, dtype=tf.float32)
     elif blend_mode == "gaussian":
-        return gaussian_kernel_tf_v2(roi_size, sigma=sigma)
+        return gaussian_kernel(roi_size, sigma)
     else:
         raise ValueError(f'Invalid blend mode: {blend_mode}. Use either "constant" or "gaussian".')
 
@@ -133,7 +121,6 @@ def sliding_window_inference(
     image_size = list(input_padded.shape[1:-1])
 
     importance_kernel = get_importance_kernel(roi_size, blend_mode, sigma=sigma)
-
     output_shape = (batch_size,) + tuple(image_size) + (n_class,)
     importance_map = tf.tile(
         tf.reshape(importance_kernel, shape=[1, *roi_size, 1]),

TensorFlow2/Segmentation/nnUNet/runtime/args.py

@@ -62,7 +62,7 @@ def get_main_args():
     p.arg("--data", type=Path, default=Path("/data"), help="Path to data directory")
     p.arg("--task", type=str, default="01", help="Task number, MSD uses numbers 01-10")
     p.arg("--dim", type=int, choices=[2, 3], default=3, help="UNet dimension")
-    p.arg("--seed", type=non_negative_int, help="Random seed")
+    p.arg("--seed", type=non_negative_int, default=None, help="Random seed")
     p.flag("--benchmark", help="Run model benchmarking")
     p.boolean_flag("--tta", default=False, help="Enable test time augmentation")
     p.boolean_flag("--save-preds", "--save_preds", default=False, help="Save predictions")