4747 compatible dtypes for padding (#4749)

* fixes #4747 padding

Signed-off-by: Wenqi Li <[email protected]>

Author: wyli

Diff for: monai/transforms/croppad/array.py

@@ -85,14 +85,15 @@ class Pad(InvertibleTransform):
     in which case `np.pad` will be used.
 
     Args:
-        to_pad: the amount to be padded in each dimension [(low_H, high_H), (low_W, high_W), ...].
+        to_pad: the amount to pad in each dimension (including the channel) [(low_H, high_H), (low_W, high_W), ...].
             if None, must provide in the `__call__` at runtime.
-        mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``,
+        mode: available modes: (Numpy) {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``,
             ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
-            available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}.
+            (PyTorch) {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}.
             One of the listed string values or a user supplied function. Defaults to ``"constant"``.
             See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
             https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html
+            requires pytorch >= 1.10 for best compatibility.
         kwargs: other arguments for the `np.pad` or `torch.pad` function.
             note that `np.pad` treats channel dimension as the first dimension.
 
@@ -122,6 +123,9 @@ def compute_pad_width(self, spatial_shape: Sequence[int]) -> List[Tuple[int, int
     def _np_pad(img: torch.Tensor, pad_width, mode, **kwargs) -> torch.Tensor:
         img_np = img.detach().cpu().numpy() if isinstance(img, torch.Tensor) else img
         mode = convert_pad_mode(dst=img_np, mode=mode).value
+        if mode == "constant" and "value" in kwargs:
+            val = kwargs.pop("value")
+            kwargs["constant_values"] = val
         out = torch.as_tensor(np.pad(img, pad_width, mode=mode, **kwargs))
         if isinstance(img, MetaTensor):
             out = convert_to_dst_type(out, dst=img)[0]
@@ -141,9 +145,9 @@ def __call__(  # type: ignore
             img: data to be transformed, assuming `img` is channel-first and padding doesn't apply to the channel dim.
             to_pad: the amount to be padded in each dimension [(low_H, high_H), (low_W, high_W), ...].
                 default to `self.to_pad`.
-            mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``,
+            mode: available modes: (Numpy) {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``,
                 ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
-                available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}.
+                (PyTorch) {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}.
                 One of the listed string values or a user supplied function. Defaults to ``"constant"``.
                 See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
                 https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html
@@ -163,16 +167,26 @@ def __call__(  # type: ignore
 
         # all zeros, skip padding
         if np.asarray(to_pad_).any():
-            if mode in ["linear_ramp", "maximum", "mean", "median", "minimum", "symmetric", "empty"]:
+            to_pad_ = list(to_pad_)
+            if len(to_pad_) < len(img_t.shape):
+                to_pad_ = list(to_pad_) + [(0, 0)] * (len(img_t.shape) - len(to_pad_))
+            if mode_ in {"linear_ramp", "maximum", "mean", "median", "minimum", "symmetric", "empty"}:
                 out = self._np_pad(img_t, pad_width=to_pad_, mode=mode_, **kwargs_)
             else:
+                mode_ = convert_pad_mode(dst=img_t, mode=mode_).value
                 try:
-                    mode_ = convert_pad_mode(dst=img_t, mode=mode_).value
-                    out = self._pt_pad(img_t, pad_width=to_pad_, mode=mode_, **kwargs_)
-                # but if mode or args don't exist in pytorch, use numpy instead
-                except (ValueError, TypeError) as err:
-                    if "Unsupported option" in str(err) or "unexpected keyword" in str(err):
+                    _pad = (
+                        self._pt_pad
+                        if mode_ in {"reflect", "replicate"}
+                        and img_t.dtype not in {torch.int16, torch.int64, torch.bool, torch.uint8}
+                        else self._np_pad
+                    )
+                    out = _pad(img_t, pad_width=to_pad_, mode=mode_, **kwargs_)
+                except (ValueError, TypeError, RuntimeError) as err:
+                    if "supported" in str(err) or "unexpected keyword" in str(err) or "implemented" in str(err):
                         out = self._np_pad(img_t, pad_width=to_pad_, mode=mode_, **kwargs_)
+                    else:
+                        raise ValueError(f"{mode_}, {kwargs_}, {img_t.dtype}, {img_t.device}") from err
         else:
             out = img_t
         if get_track_meta():

Diff for: monai/transforms/spatial/array.py

@@ -1075,7 +1075,7 @@ class Zoom(InvertibleTransform):
         padding_mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``,
             ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
             available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}.
-            One of the listed string values or a user supplied function. Defaults to ``"constant"``.
+            One of the listed string values or a user supplied function. Defaults to ``"edge"``.
             The mode to pad data after zooming.
             See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
             https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html
@@ -1123,7 +1123,7 @@ def __call__(
             padding_mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``,
                 ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
                 available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}.
-                One of the listed string values or a user supplied function. Defaults to ``"constant"``.
+                One of the listed string values or a user supplied function. Defaults to ``"edge"``.
                 The mode to pad data after zooming.
                 See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
                 https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html

Diff for: monai/transforms/spatial/dictionary.py

@@ -1448,7 +1448,7 @@ class Zoomd(MapTransform, InvertibleTransform):
         padding_mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``,
             ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
             available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}.
-            One of the listed string values or a user supplied function. Defaults to ``"constant"``.
+            One of the listed string values or a user supplied function. Defaults to ``"edge"``.
             The mode to pad data after zooming.
             See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
             https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html
@@ -1521,7 +1521,7 @@ class RandZoomd(RandomizableTransform, MapTransform, InvertibleTransform):
         padding_mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``,
             ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
             available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}.
-            One of the listed string values or a user supplied function. Defaults to ``"constant"``.
+            One of the listed string values or a user supplied function. Defaults to ``"edge"``.
             The mode to pad data after zooming.
             See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
             https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html

Diff for: tests/test_pad_mode.py

@@ -0,0 +1,37 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import torch
+
+from monai.transforms import CastToType, Pad
+from monai.utils import NumpyPadMode, PytorchPadMode
+from tests.utils import SkipIfBeforePyTorchVersion
+
+
+@SkipIfBeforePyTorchVersion((1, 10, 1))
+class TestPadMode(unittest.TestCase):
+    def test_pad(self):
+        expected_shapes = {3: (1, 15, 10), 4: (1, 10, 6, 7)}
+        for t in (float, int, np.uint8, np.int16, np.float32, bool):
+            for d in ("cuda:0", "cpu") if torch.cuda.is_available() else ("cpu",):
+                for s in ((1, 10, 10), (1, 5, 6, 7)):
+                    for m in list(PytorchPadMode) + list(NumpyPadMode):
+                        a = torch.rand(s)
+                        to_pad = [(0, 0), (2, 3)] if len(s) == 3 else [(0, 0), (2, 3), (0, 0), (0, 0)]
+                        out = Pad(to_pad=to_pad, mode=m)(CastToType(dtype=t)(a).to(d))
+                        self.assertEqual(out.shape, expected_shapes[len(s)])
+
+
+if __name__ == "__main__":
+    unittest.main()