Fix #8366: Add strict shape validation to sliding_window_inference

Signed-off-by: 林旻佑 <[email protected]>

Author: 林旻佑

monai/inferers/utils.py

@@ -13,7 +13,7 @@
 
 import itertools
 from collections.abc import Callable, Iterable, Mapping, Sequence
-from typing import Any
+from typing import Any, Optional
 
 import numpy as np
 import torch
@@ -33,12 +33,62 @@
     optional_import,
 )
 
+
 tqdm, _ = optional_import("tqdm", name="tqdm")
 _nearest_mode = "nearest-exact"
 
 __all__ = ["sliding_window_inference"]
 
 
+
+def assert_channel_first(
+    t: torch.Tensor,
+    name: str,
+    num_classes: Optional[int] = None,
+    allow_binary_two_channel: bool = False,
+) -> None:
+    """
+    Enforce channel-first layout (NCHW/NCDHW) without guessing.
+
+    Args:
+        t: Input tensor to validate.
+        name: Name of the tensor for error messages.
+        num_classes: Optional expected channel count at dim=1.
+        allow_binary_two_channel: If True and num_classes==2, accept C=2.
+
+    Raises:
+        ValueError: If tensor shape doesn't match channel-first layout or
+            num_classes constraint.
+
+    Note:
+        - Accepts only 4D (NCHW) or 5D (NCDHW) tensors with channel at dim=1.
+        - Non-tensors and tensors with other dimensionalities are silently ignored.
+        - Users must apply EnsureChannelFirst/EnsureChannelFirstd upstream for
+          channel-last data.
+    """
+    if not isinstance(t, torch.Tensor):
+        return
+    if t.ndim not in (4, 5):
+        return
+
+    c = int(t.shape[1])
+    layout = "NCHW" if t.ndim == 4 else "NCDHW"
+    layout_last = "NHWC" if t.ndim == 4 else "NDHWC"
+
+    if num_classes is not None:
+        ok = (c == num_classes) or (num_classes == 1 and c == 1)
+        if allow_binary_two_channel and num_classes == 2:
+            ok = ok or (c == 2)
+        if not ok:
+            raise ValueError(
+                f"{name}: expected {layout} with C(dim=1)==num_classes, "
+                f"but got shape={tuple(t.shape)} (C={c}) and num_classes={num_classes}. "
+                f"If your data is {layout_last}, please apply EnsureChannelFirst/EnsureChannelFirstd upstream."
+            )
+    # No guessing when num_classes is None; we simply require channel at dim=1.
+    # If callers provided NHWC/NDHWC, they must convert upstream.
+
+
 def sliding_window_inference(
     inputs: torch.Tensor | MetaTensor,
     roi_size: Sequence[int] | int,
@@ -131,11 +181,30 @@ def sliding_window_inference(
         kwargs: optional keyword args to be passed to ``predictor``.
 
     Note:
-        - input must be channel-first and have a batch dim, supports N-D sliding window.
+        - Inputs must be channel-first and have a batch dim (NCHW / NCDHW).
+        - If your data is NHWC/NDHWC, please apply `EnsureChannelFirst` / `EnsureChannelFirstd` upstream.
 
     """
-    buffered = buffer_steps is not None and buffer_steps > 0
     num_spatial_dims = len(inputs.shape) - 2
+
+    # Only perform strict shape validation if roi_size is a sequence (explicit dimensions).
+    # If roi_size is an integer, it is broadcast to all dimensions, so we cannot
+    # infer the expected dimensionality to enforce a strict check here.
+    if not isinstance(roi_size, int):
+        roi_dims = len(roi_size)
+        if num_spatial_dims != roi_dims:
+            raise ValueError(
+                f"inputs must have {roi_dims + 2} dimensions for {roi_dims}D roi_size "
+                f"(Batch, Channel, {', '.join(['Spatial'] * roi_dims)}), "
+                f"but got inputs shape {inputs.shape}.\n"
+                "If you have channel-last data (e.g. B, D, H, W, C), please use "
+                "monai.transforms.EnsureChannelFirst or EnsureChannelFirstd upstream."
+            )
+    # -----------------------------------------------------------------
+    # ---- Strict validation: do NOT guess or permute layouts ----
+    if isinstance(inputs, torch.Tensor):
+        assert_channel_first(inputs, "inputs")
+    buffered = buffer_steps is not None and buffer_steps > 0
     if buffered:
         if buffer_dim < -num_spatial_dims or buffer_dim > num_spatial_dims:
             raise ValueError(f"buffer_dim must be in [{-num_spatial_dims}, {num_spatial_dims}], got {buffer_dim}.")

tests/inferers/test_sliding_window_inference.py

@@ -372,6 +372,26 @@ def compute_dict(data):
         for rr, _ in zip(result_dict, expected_dict):
             np.testing.assert_allclose(result_dict[rr].cpu().numpy(), expected_dict[rr], rtol=1e-4)
 
+    def test_strict_shape_validation(self):
+        """Test strict shape validation to ensure inputs match roi_size dimensions."""
+        device = "cpu"
+        roi_size = (16, 16, 16)
+        sw_batch_size = 4
+        
+        def predictor(data):
+            return data
+
+        # Case 1: Input has fewer dimensions than expected (e.g., missing Batch or Channel)
+        # 3D roi_size requires 5D input (B, C, D, H, W), giving 4D here.
+        inputs_4d = torch.randn((1, 16, 16, 16), device=device)
+        with self.assertRaisesRegex(ValueError, "inputs must have 5 dimensions"):
+            sliding_window_inference(inputs_4d, roi_size, sw_batch_size, predictor)
+
+        # Case 2: Input is 3D (missing Batch AND Channel)
+        inputs_3d = torch.randn((16, 16, 16), device=device)
+        with self.assertRaisesRegex(ValueError, "inputs must have 5 dimensions"):
+            sliding_window_inference(inputs_3d, roi_size, sw_batch_size, predictor)
+
 
 class TestSlidingWindowInferenceCond(unittest.TestCase):
     @parameterized.expand(TEST_CASES)