Use replication padding in Median filter (Project-MONAI#5329)

Ref Project-MONAI#5264. Follow-up to Project-MONAI#5307. This is a more normal, intuitive variant of median filtering. Signed-off-by: Dženan Zukić <dzenan.zukic@kitware.com> Signed-off-by: Wenqi Li <wenqil@nvidia.com> Co-authored-by: Wenqi Li <wenqil@nvidia.com> Co-authored-by: Wenqi Li <831580+wyli@users.noreply.github.com>

Author: dzenanz

monai/networks/layers/simplelayers.py

@@ -498,17 +498,12 @@ def median_filter(
  kernel = kernel.to(in_tensor)
  # map the local window to single vector
  conv = [F.conv1d, F.conv2d, F.conv3d][spatial_dims - 1]
+ reshaped_input: torch.Tensor = in_tensor.reshape(oprod, 1, *sshape) # type: ignore
 
- if "padding" not in kwargs:
- if pytorch_after(1, 10):
- kwargs["padding"] = "same"
- else:
- # even-sized kernels are not supported
- kwargs["padding"] = [(k - 1) // 2 for k in kernel.shape[2:]]
- elif kwargs["padding"] == "same" and not pytorch_after(1, 10):
- # even-sized kernels are not supported
- kwargs["padding"] = [(k - 1) // 2 for k in kernel.shape[2:]]
- features: torch.Tensor = conv(in_tensor.reshape(oprod, 1, *sshape), kernel, stride=1, **kwargs) # type: ignore
+ # even-sized kernels are not supported
+ padding = [(k - 1) // 2 for k in reversed(kernel.shape[2:]) for _ in range(2)]
+ padded_input: torch.Tensor = F.pad(reshaped_input, pad=padding, mode="replicate")
+ features: torch.Tensor = conv(padded_input, kernel, padding=0, stride=1, **kwargs) # type: ignore
  features = features.view(oprod, -1, *sshape) # type: ignore
 
  # compute the median along the feature axis

tests/test_median_filter.py

@@ -18,41 +18,27 @@
 
 
 class MedianFilterTestCase(unittest.TestCase):
+ def test_3d_big(self):
+ a = torch.ones(1, 1, 2, 3, 5)
+ g = MedianFilter([1, 2, 4]).to(torch.device("cpu:0"))
+
+ expected = a.numpy()
+ out = g(a).cpu().numpy()
+ np.testing.assert_allclose(out, expected, rtol=1e-5)
+
  def test_3d(self):
  a = torch.ones(1, 1, 4, 3, 4)
  g = MedianFilter(1).to(torch.device("cpu:0"))
 
- expected = np.array(
- [
- [
- [
- [[0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0]],
- [[0.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 0.0]],
- [[0.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 0.0]],
- [[0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0]],
- ]
- ]
- ]
- )
+ expected = a.numpy()
  out = g(a).cpu().numpy()
  np.testing.assert_allclose(out, expected, rtol=1e-5)
 
  def test_3d_radii(self):
  a = torch.ones(1, 1, 4, 3, 2)
  g = MedianFilter([3, 2, 1]).to(torch.device("cpu:0"))
 
- expected = np.array(
- [
- [
- [
- [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
- [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
- [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
- [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
- ]
- ]
- ]
- )
+ expected = a.numpy()
  out = g(a).cpu().numpy()
  np.testing.assert_allclose(out, expected, rtol=1e-5)
  if torch.cuda.is_available():

tests/test_median_smooth.py

@@ -23,7 +23,7 @@
  [
  {"radius": 1},
  p([[[1, 1, 1], [2, 2, 2], [3, 3, 3]], [[4, 4, 4], [5, 5, 5], [6, 6, 6]]]),
- p([[[0.0, 1.0, 0.0], [1.0, 2, 1.0], [0.0, 2.0, 0.0]], [[0.0, 4.0, 0.0], [4.0, 5, 4.0], [0.0, 5.0, 0.0]]]),
+ p([[[1, 1, 1], [2, 2, 2], [3, 3, 3]], [[4, 4, 4], [5, 5, 5], [6, 6, 6]]]),
  ]
  )
 

tests/test_median_smoothd.py

@@ -31,39 +31,23 @@
  [
  {"keys": "img", "radius": 1},
  {"img": p(np.array([[[1, 1, 1], [2, 2, 2], [3, 3, 3]], [[4, 4, 4], [5, 5, 5], [6, 6, 6]]]))},
- np.array(
- [[[0.0, 1.0, 0.0], [1.0, 2, 1.0], [0.0, 2.0, 0.0]], [[0.0, 4.0, 0.0], [4.0, 5, 4.0], [0.0, 5.0, 0.0]]]
- ),
+ np.array([[[1, 1, 1], [2, 2, 2], [3, 3, 3]], [[4, 4, 4], [5, 5, 5], [6, 6, 6]]]),
  ]
  )
 
  TESTS.append(
  [
  {"keys": "img", "radius": [1, 1]},
  {"img": p(np.array([[[1, 1, 1], [2, 2, 2], [3, 3, 3]], [[4, 4, 4], [5, 5, 5], [6, 6, 6]]]))},
- np.array(
- [[[0.0, 1.0, 0.0], [1.0, 2, 1.0], [0.0, 2.0, 0.0]], [[0.0, 4.0, 0.0], [4.0, 5, 4.0], [0.0, 5.0, 0.0]]]
- ),
- ]
- )
-
- TESTS.append(
- [
- {"keys": "img", "radius": [1, 1]},
- {"img": p(np.array([[[1, 1, 1], [2, 2, 2], [3, 3, 3]], [[4, 4, 4], [5, 0, 5], [6, 6, 6]]]))},
- np.array(
- [[[0.0, 1.0, 0.0], [1.0, 2, 1.0], [0.0, 2.0, 0.0]], [[0.0, 4.0, 0.0], [4.0, 5, 4.0], [0.0, 5.0, 0.0]]]
- ),
+ np.array([[[1, 1, 1], [2, 2, 2], [3, 3, 3]], [[4, 4, 4], [5, 5, 5], [6, 6, 6]]]),
  ]
  )
 
  TESTS.append(
  [
  {"keys": "img", "radius": [1, 1, 1]},
  {"img": p(np.array([[[[1, 1, 1], [2, 2, 2], [3, 3, 3]], [[4, 4, 4], [5, 5, 5], [6, 6, 6]]]]))},
- np.array(
- [[[[0.0, 0.0, 0.0], [0.0, 2, 0.0], [0.0, 0, 0.0]], [[0.0, 0, 0.0], [0.0, 2, 0.0], [0.0, 0, 0.0]]]]
- ),
+ np.array([[[[2, 2, 2], [3, 3, 3], [3, 3, 3]], [[4, 4, 4], [4, 4, 4], [5, 5, 5]]]]),
  ]
  )