Repository: Project-MONAI/MONAI

Issue:
predictor access of coordinates in `sliding_window_inference`
**Is your feature request related to a problem? Please describe.**
https://github.com/Project-MONAI/MONAI/blob/5feb353030e0bb204e21e1de338cd81b5972bb8a/monai/inferers/utils.py#L219-L223

`predictor` currently takes a batch of image windows as the primary input, it has no access to the relative position information of the windows with respect to the input whole volume. would be great to add an option to send `unravel_slice` to `predictor`.

cc @heyufan1995 


Discussion:

Patch:
diff --git a/monai/apps/pathology/inferers/inferer.py b/monai/apps/pathology/inferers/inferer.py
--- a/monai/apps/pathology/inferers/inferer.py
+++ b/monai/apps/pathology/inferers/inferer.py
@@ -178,6 +178,7 @@ def __call__(
             self.process_output,
             self.buffer_steps,
             self.buffer_dim,
+            False,
             *args,
             **kwargs,
         )
diff --git a/monai/inferers/inferer.py b/monai/inferers/inferer.py
--- a/monai/inferers/inferer.py
+++ b/monai/inferers/inferer.py
@@ -426,6 +426,8 @@ class SlidingWindowInferer(Inferer):
             (i.e. no overlapping among the buffers) non_blocking copy may be automatically enabled for efficiency.
         buffer_dim: the spatial dimension along which the buffers are created.
             0 indicates the first spatial dimension. Default is -1, the last spatial dimension.
+        with_coord: whether to pass the window coordinates to ``network``. Defaults to False.
+            If True, the ``network``'s 2nd input argument should accept the window coordinates.
 
     Note:
         ``sw_batch_size`` denotes the max number of windows per network inference iteration,
@@ -449,6 +451,7 @@ def __init__(
         cpu_thresh: int | None = None,
         buffer_steps: int | None = None,
         buffer_dim: int = -1,
+        with_coord: bool = False,
     ) -> None:
         super().__init__()
         self.roi_size = roi_size
@@ -464,6 +467,7 @@ def __init__(
         self.cpu_thresh = cpu_thresh
         self.buffer_steps = buffer_steps
         self.buffer_dim = buffer_dim
+        self.with_coord = with_coord
 
         # compute_importance_map takes long time when computing on cpu. We thus
         # compute it once if it's static and then save it for future usage
@@ -525,6 +529,7 @@ def __call__(
             None,
             buffer_steps,
             buffer_dim,
+            self.with_coord,
             *args,
             **kwargs,
         )
diff --git a/monai/inferers/utils.py b/monai/inferers/utils.py
--- a/monai/inferers/utils.py
+++ b/monai/inferers/utils.py
@@ -57,6 +57,7 @@ def sliding_window_inference(
     process_fn: Callable | None = None,
     buffer_steps: int | None = None,
     buffer_dim: int = -1,
+    with_coord: bool = False,
     *args: Any,
     **kwargs: Any,
 ) -> torch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]:
@@ -125,6 +126,8 @@ def sliding_window_inference(
             (i.e. no overlapping among the buffers) non_blocking copy may be automatically enabled for efficiency.
         buffer_dim: the spatial dimension along which the buffers are created.
             0 indicates the first spatial dimension. Default is -1, the last spatial dimension.
+        with_coord: whether to pass the window coordinates to ``predictor``. Default is False.
+            If True, the signature of ``predictor`` should be ``predictor(patch_data, patch_coord, ...)``.
         args: optional args to be passed to ``predictor``.
         kwargs: optional keyword args to be passed to ``predictor``.
 
@@ -220,7 +223,10 @@ def sliding_window_inference(
             win_data = torch.cat([inputs[win_slice] for win_slice in unravel_slice]).to(sw_device)
         else:
             win_data = inputs[unravel_slice[0]].to(sw_device)
-        seg_prob_out = predictor(win_data, *args, **kwargs)  # batched patch
+        if with_coord:
+            seg_prob_out = predictor(win_data, unravel_slice, *args, **kwargs)  # batched patch
+        else:
+            seg_prob_out = predictor(win_data, *args, **kwargs)  # batched patch
 
         # convert seg_prob_out to tuple seg_tuple, this does not allocate new memory.
         dict_keys, seg_tuple = _flatten_struct(seg_prob_out)


Test patch:
diff --git a/tests/test_sliding_window_hovernet_inference.py b/tests/test_sliding_window_hovernet_inference.py
--- a/tests/test_sliding_window_hovernet_inference.py
+++ b/tests/test_sliding_window_hovernet_inference.py
@@ -237,6 +237,7 @@ def compute(data, test1, test2):
             None,
             None,
             0,
+            False,
             t1,
             test2=t2,
         )
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
@@ -294,6 +294,7 @@ def compute(data, test1, test2):
             None,
             None,
             0,
+            False,
             t1,
             test2=t2,
         )