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:
wyli: > Thanks very much for your quick update.
> It looks good to me now, just need some discussion.
> Thanks.

sure I'm working on an enhancement for the gaussian filter, may submit a PR soon
Nic-Ma: > > Thanks very much for your quick update.
> > It looks good to me now, just need some discussion.
> > Thanks.
> 
> sure I'm working on an enhancement for the gaussian filter, may submit a PR soon

Cool, then let's keep this PR open for a while and replace with our own gaussian filter when your PR merged soon.
Nic-Ma: Hi @madil90 ,

As our Gaussian filter had been updated to support multi sigma could you please take a look and update to use that?
Thanks.
Nic-Ma: /black
madil90: @Nic-Ma Yup, working on it.
Nic-Ma: /black

Patch:
diff --git a/monai/data/sliding_window_inference.py b/monai/data/sliding_window_inference.py
index 4f441bb7..84faaffc 100644
--- 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
index e3c47ea8..aebf0bfa 100644
--- a/monai/data/utils.py
+++ b/monai/data/utils.py
@@ -14,9 +14,11 @@ import warnings
 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
index e487eb41..f2a0fcf4 100644
--- a/tests/test_sliding_window_inference.py
+++ b/tests/test_sliding_window_inference.py
@@ -16,25 +16,33 @@ from parameterized import parameterized
 
 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))