Repository: Project-MONAI/MONAI

Issue:
Sliding window inference with overlap and weighting
**Is your feature request related to a problem? Please describe.**
Image-scale predictions (e.g. segmentations) tend to become worse towards the outside of the image [citation needed, I know], resulting in non-smooth predictions when a sliding window approach is used.

**Describe the solution you'd like**
An option to perform sliding window inference with overlapping patches (e.g. by half patch size) and weighting (e.g. a Gaussian over the image) so we put more weight on predictions from the center of a patch.



Discussion:
Hi @jenspetersen ,

Thanks for your feedback!
That's a very good point. Actually, we already discussed internally about overlapping in sliding window, and want to have it in a future version.
I think that's a very important and valuable improvement to our evaluation method.
Do you have any interest to submit a PR that enhances our existing sliding window method?
Thanks.
@FabianIsensee, you already implemented this, is that something you would/could contribute?
Hi,
yes I already did that a while ago and I actually rewrote the whole thing a couple of days ago because it was not easy to read and quite convoluted. At the moment I am very busy and, as much as I would like to, I cannot contribute in the near future. In case someone is interested, here is my implementation:

https://github.com/MIC-DKFZ/nnUNet/blob/master/nnunet/network_architecture/neural_network.py

The code for sliding window prediction in 3d is here:

https://github.com/MIC-DKFZ/nnUNet/blob/dba763e1abbb311f8f9fbed842ff9bde7dce87e7/nnunet/network_architecture/neural_network.py#L269

2d is here:

https://github.com/MIC-DKFZ/nnUNet/blob/dba763e1abbb311f8f9fbed842ff9bde7dce87e7/nnunet/network_architecture/neural_network.py#L589

Best,
Fabian


Hi @madil90 ,

Could you please help develop this feature request? It will obviously improve the evaluation performance.
Thanks.

@Nic-Ma Yes definitely. I'm looking at @FabianIsensee implementation and will start porting it to MONAI.
@Nic-Ma @madil90 
Getting the patch weights with Gaussians is only one way to do this, but we should probably support other approaches. For example, the next two immediately come to mind. 

- Option 1: Get the window stride (overlap) to be equivalent to the maximum receptive field of a network. This one will be very harsh at rejecting any prediction close to the border of a patch, but will guarantee that the predictions are always valid by construction.
- Option 2: Get the mixing weight to be something else other than a Gaussian (e.g. distance from the centre of the patch). Given that kernels are 3x3, information propagation along the network does not "travel" isotropically, and thus should not be Gaussian (it is actually the Chebyshev Distance). 

Also, we need to be aware that not everything is a semantic segmentation model. E.g. weighted average of 2 predictions is not necessarily ideal for semantic regression models (image synthesis). 
@mjorgecardoso Thanks, that's a very good point. I'll separate the weighting from overlapping in a way that we can implement other weighting algos later.
> @mjorgecardoso Thanks, that's a very good point. I'll separate the weighting from overlapping in a way that we can implement other weighting algos later.

Hi @madil90 ,

Thanks for your help on this PR! Looking forward to your update of our sliding window.

Resolving patch:
diff --git a/monai/data/sliding_window_inference.py b/monai/data/sliding_window_inference.py
--- a/monai/data/sliding_window_inference.py
+++ b/monai/data/sliding_window_inference.py
@@ -11,10 +11,10 @@
 
 import torch
 import torch.nn.functional as F
-from monai.data.utils import dense_patch_slices
+from monai.data.utils import dense_patch_slices, compute_importance_map
 
 
-def sliding_window_inference(inputs, roi_size, sw_batch_size, predictor):
+def sliding_window_inference(inputs, roi_size, sw_batch_size, predictor, overlap=0.25, blend_mode="constant"):
     """Use SlidingWindow method to execute inference.
 
     Args:
@@ -24,6 +24,9 @@ def sliding_window_inference(inputs, roi_size, sw_batch_size, predictor):
         predictor (Callable): given input tensor `patch_data` in shape NCHW[D], `predictor(patch_data)`
             should return a prediction with the same spatial shape and batch_size, i.e. NMHW[D];
             where HW[D] represents the patch spatial size, M is the number of output channels, N is `sw_batch_size`.
+        overlap (float): Amount of overlap between scans.
+        blend_mode (str): How to blend output of overlapping windows. Options are 'constant', 'guassian'. 'constant'
+            gives equal weight to all predictions while gaussian gives less weight to predictions on edges of windows.
 
     Note:
         must be channel first, support both 2D and 3D.
@@ -32,6 +35,7 @@ def sliding_window_inference(inputs, roi_size, sw_batch_size, predictor):
     """
     num_spatial_dims = len(inputs.shape) - 2
     assert len(roi_size) == num_spatial_dims, f"roi_size {roi_size} does not match input dims."
+    assert overlap >= 0 and overlap < 1, "overlap must be >= 0 and < 1."
 
     # determine image spatial size and batch size
     # Note: all input images must have the same image size and batch size
@@ -48,8 +52,7 @@ def sliding_window_inference(inputs, roi_size, sw_batch_size, predictor):
     pad_size = [i for k in range(len(inputs.shape) - 1, 1, -1) for i in (0, max(roi_size[k - 2] - inputs.shape[k], 0))]
     inputs = F.pad(inputs, pad=pad_size, mode="constant", value=0)
 
-    # TODO: interval from user's specification
-    scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims)
+    scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims, overlap)
 
     # Store all slices in list
     slices = dense_patch_slices(image_size, roi_size, scan_interval)
@@ -77,22 +80,30 @@ def sliding_window_inference(inputs, roi_size, sw_batch_size, predictor):
     output_classes = output_rois[0].shape[1]
     output_shape = [batch_size, output_classes] + list(image_size)
 
+    # Create importance map
+    importance_map = compute_importance_map(roi_size, mode=blend_mode, device=inputs.device)
+
     # allocate memory to store the full output and the count for overlapping parts
     output_image = torch.zeros(output_shape, dtype=torch.float32, device=inputs.device)
     count_map = torch.zeros(output_shape, dtype=torch.float32, device=inputs.device)
 
     for window_id, slice_index in enumerate(range(0, len(slices), sw_batch_size)):
         slice_index_range = range(slice_index, min(slice_index + sw_batch_size, len(slices)))
-        # store the result in the proper location of the full output
+
+        # store the result in the proper location of the full output. Apply weights from importance map.
         for curr_index in slice_index_range:
             if num_spatial_dims == 3:
                 slice_i, slice_j, slice_k = slices[curr_index]
-                output_image[0, :, slice_i, slice_j, slice_k] += output_rois[window_id][curr_index - slice_index, :]
-                count_map[0, :, slice_i, slice_j, slice_k] += 1.0
+                output_image[0, :, slice_i, slice_j, slice_k] += (
+                    importance_map * output_rois[window_id][curr_index - slice_index, :]
+                )
+                count_map[0, :, slice_i, slice_j, slice_k] += importance_map
             else:
                 slice_i, slice_j = slices[curr_index]
-                output_image[0, :, slice_i, slice_j] += output_rois[window_id][curr_index - slice_index, :]
-                count_map[0, :, slice_i, slice_j] += 1.0
+                output_image[0, :, slice_i, slice_j] += (
+                    importance_map * output_rois[window_id][curr_index - slice_index, :]
+                )
+                count_map[0, :, slice_i, slice_j] += importance_map
 
     # account for any overlapping sections
     output_image /= count_map
@@ -102,15 +113,15 @@ def sliding_window_inference(inputs, roi_size, sw_batch_size, predictor):
     return output_image[..., : original_image_size[0], : original_image_size[1]]  # 2D
 
 
-def _get_scan_interval(image_size, roi_size, num_spatial_dims):
+def _get_scan_interval(image_size, roi_size, num_spatial_dims, overlap):
     assert len(image_size) == num_spatial_dims, "image coord different from spatial dims."
     assert len(roi_size) == num_spatial_dims, "roi coord different from spatial dims."
 
-    scan_interval = [1 for _ in range(num_spatial_dims)]
+    scan_interval = []
     for i in range(num_spatial_dims):
         if roi_size[i] == image_size[i]:
-            scan_interval[i] = int(roi_size[i])
+            scan_interval.append(int(roi_size[i]))
         else:
-            # this means that it's r-16 (if r>=64) and r*0.75 (if r<=64)
-            scan_interval[i] = int(max(roi_size[i] - 16, roi_size[i] * 0.75))
+            # scan interval is (1-overlap)*roi_size
+            scan_interval.append(int(roi_size[i] * (1 - overlap)))
     return tuple(scan_interval)
diff --git a/monai/data/utils.py b/monai/data/utils.py
--- a/monai/data/utils.py
+++ b/monai/data/utils.py
@@ -14,9 +14,11 @@
 import math
 import nibabel as nib
 from itertools import starmap, product
+import torch
 from torch.utils.data._utils.collate import default_collate
 import numpy as np
 from monai.utils import ensure_tuple_size
+from monai.networks.layers.simplelayers import GaussianFilter
 
 
 def get_random_patch(dims, patch_size, rand_state=None):
@@ -399,3 +401,40 @@ def create_file_basename(postfix, input_file_name, folder_path, data_root_dir=""
 
     # add the sub-folder plus the postfix name to become the file basename in the output path
     return os.path.join(subfolder_path, filename + "_" + postfix)
+
+
+def compute_importance_map(patch_size, mode="constant", sigma_scale=0.125, device=None):
+    """Get importance map for different weight modes.
+
+    Args:
+        patch_size (tuple): Size of the required importance map. This should be either H, W [,D].
+        mode (str): Importance map type. Options are 'constant' (Each weight has value 1.0)
+            or 'gaussian' (Importance becomes lower away from center).
+        sigma_scale (float): Sigma_scale to calculate sigma for each dimension 
+            (sigma = sigma_scale * dim_size). Used for gaussian mode only.
+        device (str of pytorch device): Device to put importance map on.
+
+    Returns:
+        Tensor of size patch_size.
+    """
+    importance_map = None
+    if mode == "constant":
+        importance_map = torch.ones(patch_size, device=device).float()
+    elif mode == "gaussian":
+        center_coords = [i // 2 for i in patch_size]
+        sigmas = [i * sigma_scale for i in patch_size]
+
+        importance_map = torch.zeros(patch_size, device=device)
+        importance_map[tuple(center_coords)] = 1
+        pt_gaussian = GaussianFilter(len(patch_size), sigmas).to(device)
+        importance_map = pt_gaussian(importance_map.unsqueeze(0).unsqueeze(0))
+        importance_map = importance_map.squeeze(0).squeeze(0)
+        importance_map = importance_map / torch.max(importance_map)
+        importance_map = importance_map.float()
+
+        # importance_map cannot be 0, otherwise we may end up with nans!
+        importance_map[importance_map == 0] = torch.min(importance_map[importance_map != 0])
+    else:
+        raise ValueError('mode must be "constant" or "gaussian".')
+
+    return importance_map


Test patch:
diff --git a/tests/test_sliding_window_inference.py b/tests/test_sliding_window_inference.py
--- a/tests/test_sliding_window_inference.py
+++ b/tests/test_sliding_window_inference.py
@@ -16,25 +16,33 @@
 
 from monai.data import sliding_window_inference
 
-TEST_CASE_1 = [(1, 3, 16, 15, 7), (4, 10, 7), 3]  # 3D small roi
+TEST_CASE_1 = [(1, 3, 16, 15, 7), (4, 10, 7), 3, 0.25, "constant"]  # 3D small roi
 
-TEST_CASE_2 = [(1, 3, 16, 15, 7), (20, 22, 23), 10]  # 3D large roi
+TEST_CASE_2 = [(1, 3, 16, 15, 7), (20, 22, 23), 10, 0.25, "constant"]  # 3D large roi
 
-TEST_CASE_3 = [(1, 3, 15, 7), (2, 6), 1000]  # 2D small roi, large batch
+TEST_CASE_3 = [(1, 3, 15, 7), (2, 6), 1000, 0.25, "constant"]  # 2D small roi, large batch
 
-TEST_CASE_4 = [(1, 3, 16, 7), (80, 50), 7]  # 2D large roi
+TEST_CASE_4 = [(1, 3, 16, 7), (80, 50), 7, 0.25, "constant"]  # 2D large roi
+
+TEST_CASE_5 = [(1, 3, 16, 15, 7), (20, 22, 23), 10, 0.5, "constant"]  # 3D large overlap
+
+TEST_CASE_6 = [(1, 3, 16, 7), (80, 50), 7, 0.5, "gaussian"]  # 2D large overlap, gaussian
+
+TEST_CASE_7 = [(1, 3, 16, 15, 7), (4, 10, 7), 3, 0.25, "gaussian"]  # 3D small roi, guassian
 
 
 class TestSlidingWindowInference(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
-    def test_sliding_window_default(self, image_shape, roi_shape, sw_batch_size):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7])
+    def test_sliding_window_default(self, image_shape, roi_shape, sw_batch_size, overlap, mode):
         inputs = torch.ones(*image_shape)
         device = torch.device("cpu:0")
 
         def compute(data):
             return data + 1
 
-        result = sliding_window_inference(inputs.to(device), roi_shape, sw_batch_size, compute)
+        result = sliding_window_inference(
+            inputs.to(device), roi_shape, sw_batch_size, compute, overlap, blend_mode=mode
+        )
         expected_val = np.ones(image_shape, dtype=np.float32) + 1
         self.assertTrue(np.allclose(result.numpy(), expected_val))