Skip to content

Commit 974c1d9

Browse files
林旻佑林旻佑
authored andcommitted
Fix #8366: Add strict shape validation to sliding_window_inference
Signed-off-by: 林旻佑 <[email protected]>
1 parent e267705 commit 974c1d9

File tree

2 files changed

+38
-2
lines changed

2 files changed

+38
-2
lines changed

monai/inferers/utils.py

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -131,11 +131,27 @@ def sliding_window_inference(
131131
kwargs: optional keyword args to be passed to ``predictor``.
132132
133133
Note:
134-
- input must be channel-first and have a batch dim, supports N-D sliding window.
134+
- Inputs must be channel-first and have a batch dim (NCHW / NCDHW).
135+
- If your data is NHWC/NDHWC, please apply `EnsureChannelFirst` / `EnsureChannelFirstd` upstream.
135136
136137
"""
137-
buffered = buffer_steps is not None and buffer_steps > 0
138138
num_spatial_dims = len(inputs.shape) - 2
139+
140+
# Only perform strict shape validation if roi_size is a sequence (explicit dimensions).
141+
# If roi_size is an integer, it is broadcast to all dimensions, so we cannot
142+
# infer the expected dimensionality to enforce a strict check here.
143+
if not isinstance(roi_size, Sequence):
144+
roi_dims = len(roi_size)
145+
if num_spatial_dims != roi_dims:
146+
raise ValueError(
147+
f"inputs must have {roi_dims + 2} dimensions for {roi_dims}D roi_size "
148+
f"(Batch, Channel, {', '.join(['Spatial'] * roi_dims)}), "
149+
f"but got inputs shape {inputs.shape}.\n"
150+
"If you have channel-last data (e.g. B, D, H, W, C), please use "
151+
"monai.transforms.EnsureChannelFirst or EnsureChannelFirstd upstream."
152+
)
153+
# -----------------------------------------------------------------
154+
buffered = buffer_steps is not None and buffer_steps > 0
139155
if buffered:
140156
if buffer_dim < -num_spatial_dims or buffer_dim > num_spatial_dims:
141157
raise ValueError(f"buffer_dim must be in [{-num_spatial_dims}, {num_spatial_dims}], got {buffer_dim}.")

tests/inferers/test_sliding_window_inference.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -372,6 +372,26 @@ def compute_dict(data):
372372
for rr, _ in zip(result_dict, expected_dict):
373373
np.testing.assert_allclose(result_dict[rr].cpu().numpy(), expected_dict[rr], rtol=1e-4)
374374

375+
def test_strict_shape_validation(self):
376+
"""Test strict shape validation to ensure inputs match roi_size dimensions."""
377+
device = "cpu"
378+
roi_size = (16, 16, 16)
379+
sw_batch_size = 4
380+
381+
def predictor(data):
382+
return data
383+
384+
# Case 1: Input has fewer dimensions than expected (e.g., missing Batch or Channel)
385+
# 3D roi_size requires 5D input (B, C, D, H, W), giving 4D here.
386+
inputs_4d = torch.randn((1, 16, 16, 16), device=device)
387+
with self.assertRaisesRegex(ValueError, "inputs must have 5 dimensions"):
388+
sliding_window_inference(inputs_4d, roi_size, sw_batch_size, predictor)
389+
390+
# Case 2: Input is 3D (missing Batch AND Channel)
391+
inputs_3d = torch.randn((16, 16, 16), device=device)
392+
with self.assertRaisesRegex(ValueError, "inputs must have 5 dimensions"):
393+
sliding_window_inference(inputs_3d, roi_size, sw_batch_size, predictor)
394+
375395

376396
class TestSlidingWindowInferenceCond(unittest.TestCase):
377397
@parameterized.expand(TEST_CASES)

0 commit comments

Comments
 (0)