diff --git a/monai/config/type_definitions.py b/monai/config/type_definitions.py index 478830437e..5e294e0687 100644 --- a/monai/config/type_definitions.py +++ b/monai/config/type_definitions.py @@ -58,13 +58,16 @@ DtypeLike = Union[np.dtype, type, None] """Type of datatypes -adapted from https://github.com/numpy/numpy/blob/master/numpy/typing/_dtype_like.py + +Adapted from https://github.com/numpy/numpy/blob/master/numpy/typing/_dtype_like.py """ -# Generic type which can represent either a numpy.ndarray or a torch.Tensor -# Unlike Union can create a dependence between parameter(s) / return(s) NdarrayTensor = TypeVar("NdarrayTensor", np.ndarray, torch.Tensor) +"""NdarrayTensor +Generic type which can represent either a numpy.ndarray or a torch.Tensor +Unlike Union can create a dependence between parameter(s) / return(s) +""" TensorOrList = Union[torch.Tensor, Sequence[torch.Tensor]] """TensorOrList diff --git a/monai/transforms/croppad/array.py b/monai/transforms/croppad/array.py index 0eca034950..06d6badf80 100644 --- a/monai/transforms/croppad/array.py +++ b/monai/transforms/croppad/array.py @@ -19,8 +19,10 @@ import numpy as np import torch +from torch.nn.functional import pad as pad_pt from monai.config import IndexSelection +from monai.config.type_definitions import NdarrayTensor from monai.data.utils import get_random_patch, get_valid_patch_size from monai.transforms.transform import Randomizable, Transform from monai.transforms.utils import ( @@ -34,6 +36,8 @@ weighted_patch_samples, ) from monai.utils import Method, NumpyPadMode, ensure_tuple, ensure_tuple_rep, fall_back_tuple, look_up_option +from monai.utils.enums import TransformBackends +from monai.utils.type_conversion import convert_data_type __all__ = [ "SpatialPad", @@ -54,9 +58,72 @@ ] +class Pad(Transform): + """ + Perform padding for a given an amount of padding in each dimension. + If input is `torch.Tensor` and mode is `constant`, `torch.nn.functional.pad` will be used. + Otherwise, `np.pad` will be used (input converted to `np.ndarray` if necessary). + Uses np.pad so in practice, a mode needs to be provided. See numpy.lib.arraypad.pad + for additional details. + Args: + to_pad: the amount to be padded in each dimension [(low_H, high_H), (low_W, high_W), ...]. + mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``, + ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} + 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 + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__( + self, + to_pad: List[Tuple[int, int]], + mode: Union[NumpyPadMode, str, None] = NumpyPadMode.CONSTANT, + **np_kwargs, + ) -> None: + self.to_pad = to_pad + self.mode = mode or NumpyPadMode.CONSTANT + self.np_kwargs = np_kwargs + + @staticmethod + def _np_pad(img: np.ndarray, all_pad_width, mode, **np_kwargs) -> np.ndarray: + img_np, *_ = convert_data_type(img, np.ndarray) + return np.pad(img_np, all_pad_width, mode=mode, **np_kwargs) # type: ignore + + @staticmethod + def _pt_pad(img: torch.Tensor, all_pad_width, mode, **np_kwargs) -> torch.Tensor: + pt_pad_width = [val for sublist in all_pad_width for val in sublist[::-1]][::-1] + return pad_pt(img, pt_pad_width, mode=mode, **np_kwargs) + + def __call__(self, img: NdarrayTensor, mode: Optional[Union[NumpyPadMode, str]] = None) -> NdarrayTensor: + """ + Args: + img: data to be transformed, assuming `img` is channel-first and + padding doesn't apply to the channel dim. + mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``, + ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} + One of the listed string values or a user supplied function. Defaults to ``self.mode``. + See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html + """ + if not np.asarray(self.to_pad).any(): + # all zeros, skip padding + return img + mode = mode or self.mode + mode = mode.value if isinstance(mode, NumpyPadMode) else mode + if isinstance(img, torch.Tensor) and mode == "constant" and not self.np_kwargs: + pad = self._pt_pad + else: + pad = self._np_pad # type: ignore + return pad(img, self.to_pad, mode, **self.np_kwargs) + + class SpatialPad(Transform): """ Performs padding to the data, symmetric for all sides or all on one side for each dimension. + + If input is `torch.Tensor` and mode is `constant`, `torch.nn.functional.pad` will be used. + Otherwise, `np.pad` will be used (input converted to `np.ndarray` if necessary). + Uses np.pad so in practice, a mode needs to be provided. See numpy.lib.arraypad.pad for additional details. @@ -77,6 +144,8 @@ class SpatialPad(Transform): """ + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + def __init__( self, spatial_size: Union[Sequence[int], int], @@ -99,7 +168,7 @@ def _determine_data_pad_width(self, data_shape: Sequence[int]) -> List[Tuple[int return pad_width return [(0, max(sp_i - data_shape[i], 0)) for i, sp_i in enumerate(spatial_size)] - def __call__(self, img: np.ndarray, mode: Optional[Union[NumpyPadMode, str]] = None) -> np.ndarray: + def __call__(self, img: NdarrayTensor, mode: Optional[Union[NumpyPadMode, str]] = None) -> NdarrayTensor: """ Args: img: data to be transformed, assuming `img` is channel-first and @@ -115,9 +184,9 @@ def __call__(self, img: np.ndarray, mode: Optional[Union[NumpyPadMode, str]] = N # all zeros, skip padding return img - mode = look_up_option(self.mode if mode is None else mode, NumpyPadMode).value - img = np.pad(img, all_pad_width, mode=mode, **self.np_kwargs) - return img + mode = look_up_option(mode or self.mode, NumpyPadMode) + padder = Pad(all_pad_width, mode, **self.np_kwargs) + return padder(img) class BorderPad(Transform): @@ -145,6 +214,8 @@ class BorderPad(Transform): """ + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + def __init__( self, spatial_border: Union[Sequence[int], int], @@ -155,7 +226,7 @@ def __init__( self.mode: NumpyPadMode = look_up_option(mode, NumpyPadMode) self.np_kwargs = np_kwargs - def __call__(self, img: np.ndarray, mode: Optional[Union[NumpyPadMode, str]] = None): + def __call__(self, img: NdarrayTensor, mode: Optional[Union[NumpyPadMode, str]] = None) -> NdarrayTensor: """ Args: img: data to be transformed, assuming `img` is channel-first and @@ -189,8 +260,10 @@ def __call__(self, img: np.ndarray, mode: Optional[Union[NumpyPadMode, str]] = N f"[1, len(spatial_shape)={len(spatial_shape)}, 2*len(spatial_shape)={2*len(spatial_shape)}]." ) - mode = look_up_option(self.mode if mode is None else mode, NumpyPadMode).value - return np.pad(img, [(0, 0)] + data_pad_width, mode=mode, **self.np_kwargs) + all_pad_width = [(0, 0)] + data_pad_width + mode = look_up_option(mode or self.mode, NumpyPadMode) + padder = Pad(all_pad_width, mode, **self.np_kwargs) + return padder(img) class DivisiblePad(Transform): @@ -198,6 +271,8 @@ class DivisiblePad(Transform): Pad the input data, so that the spatial sizes are divisible by `k`. """ + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + def __init__( self, k: Union[Sequence[int], int], @@ -226,7 +301,7 @@ def __init__( self.method: Method = Method(method) self.np_kwargs = np_kwargs - def __call__(self, img: np.ndarray, mode: Optional[Union[NumpyPadMode, str]] = None) -> np.ndarray: + def __call__(self, img: NdarrayTensor, mode: Optional[Union[NumpyPadMode, str]] = None) -> NdarrayTensor: """ Args: img: data to be transformed, assuming `img` is channel-first diff --git a/monai/transforms/croppad/dictionary.py b/monai/transforms/croppad/dictionary.py index a642bf406b..7c517cae96 100644 --- a/monai/transforms/croppad/dictionary.py +++ b/monai/transforms/croppad/dictionary.py @@ -25,6 +25,7 @@ import numpy as np from monai.config import IndexSelection, KeysCollection +from monai.config.type_definitions import NdarrayTensor from monai.data.utils import get_random_patch, get_valid_patch_size from monai.transforms.croppad.array import ( BorderPad, @@ -49,7 +50,7 @@ ) from monai.utils import ImageMetaKey as Key from monai.utils import Method, NumpyPadMode, ensure_tuple, ensure_tuple_rep, fall_back_tuple -from monai.utils.enums import InverseKeys +from monai.utils.enums import InverseKeys, TransformBackends __all__ = [ "NumpyPadModeSequence", @@ -106,6 +107,8 @@ class SpatialPadd(MapTransform, InvertibleTransform): Performs padding to the data, symmetric for all sides or all on one side for each dimension. """ + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + def __init__( self, keys: KeysCollection, @@ -140,7 +143,7 @@ def __init__( self.mode = ensure_tuple_rep(mode, len(self.keys)) self.padder = SpatialPad(spatial_size, method, **np_kwargs) - def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + def __call__(self, data: Mapping[Hashable, NdarrayTensor]) -> Dict[Hashable, NdarrayTensor]: d = dict(data) for key, m in self.key_iterator(d, self.mode): self.push_transform(d, key, extra_info={"mode": m.value if isinstance(m, Enum) else m}) @@ -174,6 +177,8 @@ class BorderPadd(MapTransform, InvertibleTransform): Dictionary-based wrapper of :py:class:`monai.transforms.BorderPad`. """ + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + def __init__( self, keys: KeysCollection, @@ -211,7 +216,7 @@ def __init__( self.mode = ensure_tuple_rep(mode, len(self.keys)) self.padder = BorderPad(spatial_border=spatial_border, **np_kwargs) - def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + def __call__(self, data: Mapping[Hashable, NdarrayTensor]) -> Dict[Hashable, NdarrayTensor]: d = dict(data) for key, m in self.key_iterator(d, self.mode): self.push_transform(d, key, extra_info={"mode": m.value if isinstance(m, Enum) else m}) @@ -249,6 +254,8 @@ class DivisiblePadd(MapTransform, InvertibleTransform): Dictionary-based wrapper of :py:class:`monai.transforms.DivisiblePad`. """ + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + def __init__( self, keys: KeysCollection, @@ -283,7 +290,7 @@ def __init__( self.mode = ensure_tuple_rep(mode, len(self.keys)) self.padder = DivisiblePad(k=k, method=method, **np_kwargs) - def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]: + def __call__(self, data: Mapping[Hashable, NdarrayTensor]) -> Dict[Hashable, NdarrayTensor]: d = dict(data) for key, m in self.key_iterator(d, self.mode): self.push_transform(d, key, extra_info={"mode": m.value if isinstance(m, Enum) else m}) diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py index 1871eedb6f..59b4d64174 100644 --- a/monai/transforms/utility/array.py +++ b/monai/transforms/utility/array.py @@ -39,6 +39,8 @@ min_version, optional_import, ) +from monai.utils.enums import TransformBackends +from monai.utils.type_conversion import convert_data_type PILImageImage, has_pil = optional_import("PIL.Image", name="Image") pil_image_fromarray, _ = optional_import("PIL.Image", name="fromarray") @@ -288,6 +290,8 @@ class CastToType(Transform): specified PyTorch data type. """ + backends = [TransformBackends.TORCH, TransformBackends.NUMPY] + def __init__(self, dtype=np.float32) -> None: """ Args: @@ -295,9 +299,7 @@ def __init__(self, dtype=np.float32) -> None: """ self.dtype = dtype - def __call__( - self, img: Union[np.ndarray, torch.Tensor], dtype: Optional[Union[DtypeLike, torch.dtype]] = None - ) -> Union[np.ndarray, torch.Tensor]: + def __call__(self, img: NdarrayTensor, dtype: Optional[Union[DtypeLike, torch.dtype]] = None) -> NdarrayTensor: """ Apply the transform to `img`, assuming `img` is a numpy array or PyTorch Tensor. @@ -308,11 +310,10 @@ def __call__( TypeError: When ``img`` type is not in ``Union[numpy.ndarray, torch.Tensor]``. """ - if isinstance(img, np.ndarray): - return img.astype(self.dtype if dtype is None else dtype) # type: ignore - if isinstance(img, torch.Tensor): - return torch.as_tensor(img, dtype=self.dtype if dtype is None else dtype) - raise TypeError(f"img must be one of (numpy.ndarray, torch.Tensor) but is {type(img).__name__}.") + if not isinstance(img, (torch.Tensor, np.ndarray)): + raise TypeError(f"img must be one of (numpy.ndarray, torch.Tensor) but is {type(img).__name__}.") + img_out, *_ = convert_data_type(img, output_type=type(img), dtype=dtype or self.dtype) + return img_out class ToTensor(Transform): diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py index 9c0a709bbf..9f6680fdf5 100644 --- a/monai/transforms/utility/dictionary.py +++ b/monai/transforms/utility/dictionary.py @@ -58,7 +58,7 @@ ) from monai.transforms.utils import extreme_points_to_image, get_extreme_points from monai.utils import convert_to_numpy, ensure_tuple, ensure_tuple_rep -from monai.utils.enums import InverseKeys +from monai.utils.enums import InverseKeys, TransformBackends __all__ = [ "AddChannelD", @@ -393,6 +393,8 @@ class CastToTyped(MapTransform): Dictionary-based wrapper of :py:class:`monai.transforms.CastToType`. """ + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + def __init__( self, keys: KeysCollection, @@ -413,9 +415,7 @@ def __init__( self.dtype = ensure_tuple_rep(dtype, len(self.keys)) self.converter = CastToType() - def __call__( - self, data: Mapping[Hashable, Union[np.ndarray, torch.Tensor]] - ) -> Dict[Hashable, Union[np.ndarray, torch.Tensor]]: + def __call__(self, data: Mapping[Hashable, NdarrayTensor]) -> Dict[Hashable, NdarrayTensor]: d = dict(data) for key, dtype in self.key_iterator(d, self.dtype): d[key] = self.converter(d[key], dtype=dtype) diff --git a/monai/utils/type_conversion.py b/monai/utils/type_conversion.py index d4d911c5b2..54c92c46ef 100644 --- a/monai/utils/type_conversion.py +++ b/monai/utils/type_conversion.py @@ -1,5 +1,5 @@ import re -from typing import Any, Optional, Sequence, Tuple, Union +from typing import Any, Optional, Sequence, Tuple, Type, Union, cast import numpy as np import torch @@ -147,7 +147,7 @@ def convert_to_numpy(data): def convert_data_type( data: Any, - output_type: Optional[type] = None, + output_type: Optional[Type[NdarrayTensor]] = None, device: Optional[torch.device] = None, dtype: Optional[Union[DtypeLike, torch.dtype]] = None, ) -> Tuple[NdarrayTensor, type, Optional[torch.device]]: @@ -176,16 +176,16 @@ def convert_data_type( data = convert_to_tensor(data) if dtype != data.dtype: data = data.to(dtype) # type: ignore - elif output_type is np.ndarray: + if device is not None: + data = data.to(device) + return cast(NdarrayTensor, data), orig_type, orig_device # pytype: disable=invalid-annotation + if output_type is np.ndarray: if orig_type is not np.ndarray: data = convert_to_numpy(data) if data is not None and dtype != data.dtype: data = data.astype(dtype) # type: ignore - - if isinstance(data, torch.Tensor) and device is not None: - data = data.to(device) - - return data, orig_type, orig_device + return cast(NdarrayTensor, data), orig_type, orig_device # pytype: disable=invalid-annotation + raise ValueError(f"Unsupported output type: {output_type}") def convert_to_dst_type(src: Any, dst: NdarrayTensor) -> Tuple[NdarrayTensor, type, Optional[torch.device]]: diff --git a/tests/test_border_pad.py b/tests/test_border_pad.py index b011601694..9e6a8a6a08 100644 --- a/tests/test_border_pad.py +++ b/tests/test_border_pad.py @@ -16,6 +16,7 @@ from monai.transforms import BorderPad from monai.utils import NumpyPadMode +from tests.utils import TEST_NDARRAYS TEST_CASE_1 = [ {"spatial_border": 2, "mode": "constant"}, @@ -45,11 +46,12 @@ class TestBorderPad(unittest.TestCase): @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4]) def test_pad_shape(self, input_param, input_data, expected_val): - padder = BorderPad(**input_param) - result = padder(input_data) - self.assertAlmostEqual(result.shape, expected_val.shape) - result = padder(input_data, mode=input_param["mode"]) - self.assertAlmostEqual(result.shape, expected_val.shape) + for p in TEST_NDARRAYS: + padder = BorderPad(**input_param) + r1 = padder(p(input_data)) + r2 = padder(input_data, mode=input_param["mode"]) + self.assertAlmostEqual(r1.shape, expected_val.shape) + self.assertAlmostEqual(r2.shape, expected_val.shape) def test_pad_kwargs(self): padder = BorderPad(spatial_border=2, mode="constant", constant_values=((0, 0), (1, 1), (2, 2))) diff --git a/tests/test_cast_to_type.py b/tests/test_cast_to_type.py index 5e81b41650..0ef25cbafa 100644 --- a/tests/test_cast_to_type.py +++ b/tests/test_cast_to_type.py @@ -16,17 +16,24 @@ from parameterized import parameterized from monai.transforms import CastToType +from monai.utils.type_conversion import get_equivalent_dtype +from tests.utils import TEST_NDARRAYS -TEST_CASE_1 = [{"dtype": np.float64}, np.array([[0, 1], [1, 2]], dtype=np.float32), np.float64] - -TEST_CASE_2 = [{"dtype": torch.float64}, torch.tensor([[0, 1], [1, 2]], dtype=torch.float32), torch.float64] +TESTS = [] +for p in TEST_NDARRAYS: + for out_dtype in (np.float64, torch.float64): + TESTS.append([out_dtype, p(np.array([[0, 1], [1, 2]], dtype=np.float32)), out_dtype]) class TestCastToType(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2]) - def test_type(self, input_param, input_data, expected_type): - result = CastToType(**input_param)(input_data) - self.assertEqual(result.dtype, expected_type) + @parameterized.expand(TESTS) + def test_type(self, out_dtype, input_data, expected_type): + + result = CastToType(dtype=out_dtype)(input_data) + self.assertEqual(result.dtype, get_equivalent_dtype(expected_type, type(result))) + + result = CastToType()(input_data, out_dtype) + self.assertEqual(result.dtype, get_equivalent_dtype(expected_type, type(result))) if __name__ == "__main__": diff --git a/tests/test_divisible_pad.py b/tests/test_divisible_pad.py index e4415a2f22..ca15b4b347 100644 --- a/tests/test_divisible_pad.py +++ b/tests/test_divisible_pad.py @@ -12,27 +12,36 @@ import unittest import numpy as np +import torch from parameterized import parameterized from monai.transforms import DivisiblePad - -# pad first dim to be divisible by 7, the second unchanged. -TEST_CASE_1 = [ - {"k": (7, -1), "mode": "constant"}, - np.zeros((3, 8, 7)), - np.zeros((3, 14, 7)), -] - -# pad all dimensions to be divisible by 5 -TEST_CASE_2 = [ - {"k": 5, "mode": "constant", "method": "end"}, - np.zeros((3, 10, 5, 17)), - np.zeros((3, 10, 5, 20)), -] +from tests.utils import TEST_NDARRAYS + +TESTS = [] + +for p in TEST_NDARRAYS: + # pad first dim to be divisible by 7, the second unchanged. + TESTS.append( + [ + {"k": (7, -1), "mode": "constant"}, + p(np.zeros((3, 8, 7))), + p(np.zeros((3, 14, 7))), + ] + ) + + # pad all dimensions to be divisible by 5 + TESTS.append( + [ + {"k": 5, "mode": "constant", "method": "end"}, + p(np.zeros((3, 10, 5, 17))), + p(np.zeros((3, 10, 5, 20))), + ] + ) class TestDivisiblePad(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2]) + @parameterized.expand(TESTS) def test_pad_shape(self, input_param, input_data, expected_val): padder = DivisiblePad(**input_param) result = padder(input_data) @@ -42,9 +51,11 @@ def test_pad_shape(self, input_param, input_data, expected_val): def test_pad_kwargs(self): padder = DivisiblePad(k=5, mode="constant", constant_values=((0, 0), (1, 1), (2, 2))) - result = padder(np.zeros((3, 8, 4))) - np.testing.assert_allclose(result[:, :1, :4], np.ones((3, 1, 4))) - np.testing.assert_allclose(result[:, :, 4:5], np.ones((3, 10, 1)) + 1) + for p in TEST_NDARRAYS: + result = padder(p(np.zeros((3, 8, 4)))) + result = result.cpu() if isinstance(result, torch.Tensor) else result + torch.testing.assert_allclose(result[:, :1, :4], np.ones((3, 1, 4)), rtol=1e-7, atol=0) + torch.testing.assert_allclose(result[:, :, 4:5], np.ones((3, 10, 1)) + 1, rtol=1e-7, atol=0) if __name__ == "__main__": diff --git a/tests/test_get_equivalent_dtype.py b/tests/test_get_equivalent_dtype.py index 96f4a4d720..04ba5ae5fb 100644 --- a/tests/test_get_equivalent_dtype.py +++ b/tests/test_get_equivalent_dtype.py @@ -26,9 +26,9 @@ TESTS.append((p(np.array(1.0, dtype=np.float32)), im_dtype)) -class TestDtypeConvert(unittest.TestCase): +class TestGetEquivalentDtype(unittest.TestCase): @parameterized.expand(TESTS) - def test_dtype_convert(self, im, input_dtype): + def test_get_equivalent_dtype(self, im, input_dtype): out_dtype = get_equivalent_dtype(input_dtype, type(im)) self.assertEqual(out_dtype, im.dtype) diff --git a/tests/test_spatial_pad.py b/tests/test_spatial_pad.py index 93241610de..86d010bbad 100644 --- a/tests/test_spatial_pad.py +++ b/tests/test_spatial_pad.py @@ -10,47 +10,91 @@ # limitations under the License. import unittest +from typing import List import numpy as np +import torch from parameterized import parameterized from monai.transforms import SpatialPad +from monai.utils.enums import NumpyPadMode +from monai.utils.misc import set_determinism +from tests.utils import TEST_NDARRAYS -TEST_CASE_1 = [ - {"spatial_size": [15, 8, 8], "method": "symmetric", "mode": "constant"}, - np.zeros((3, 8, 8, 4)), - np.zeros((3, 15, 8, 8)), -] +TESTS = [] -TEST_CASE_2 = [ - {"spatial_size": [15, 8, 8], "method": "end", "mode": "constant"}, - np.zeros((3, 8, 8, 4)), - np.zeros((3, 15, 8, 8)), +# Numpy modes +MODES: List = [ + "constant", + "edge", + "linear_ramp", + "maximum", + "mean", + "median", + "minimum", + "reflect", + "symmetric", + "wrap", + "empty", ] +MODES += [NumpyPadMode(i) for i in MODES] -TEST_CASE_3 = [ - {"spatial_size": [15, 4, -1], "method": "symmetric", "mode": "constant"}, - np.zeros((3, 8, 8, 4)), - np.zeros((3, 15, 8, 4)), -] +for mode in MODES: + TESTS.append( + [ + {"spatial_size": [50, 50], "method": "end", "mode": mode}, + (1, 2, 2), + (1, 50, 50), + ] + ) + + TESTS.append( + [ + {"spatial_size": [15, 4, -1], "method": "symmetric", "mode": mode}, + (3, 8, 8, 4), + (3, 15, 8, 4), + ] + ) class TestSpatialPad(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) - def test_pad_shape(self, input_param, input_data, expected_val): - padder = SpatialPad(**input_param) - result = padder(input_data) - np.testing.assert_allclose(result.shape, expected_val.shape) - result = padder(input_data, mode=input_param["mode"]) - np.testing.assert_allclose(result.shape, expected_val.shape) + def setUp(self) -> None: + set_determinism(seed=0) + + def tearDown(self) -> None: + set_determinism(None) + + @staticmethod + def get_arr(shape): + return np.random.randint(100, size=shape).astype(float) + + @parameterized.expand(TESTS) + def test_pad_shape(self, input_param, input_shape, expected_shape): + results_1 = [] + results_2 = [] + input_data = self.get_arr(input_shape) + # check result is the same regardless of input type + for p in TEST_NDARRAYS: + padder = SpatialPad(**input_param) + r1 = padder(p(input_data)) + r2 = padder(p(input_data), mode=input_param["mode"]) + results_1.append(r1.cpu() if isinstance(r1, torch.Tensor) else r1) + results_2.append(r2.cpu() if isinstance(r2, torch.Tensor) else r2) + for results in (results_1, results_2): + np.testing.assert_allclose(results[-1].shape, expected_shape) + if input_param["mode"] not in ("empty", NumpyPadMode.EMPTY): + torch.testing.assert_allclose(results[0], results[-1], atol=0, rtol=1e-5) def test_pad_kwargs(self): padder = SpatialPad( spatial_size=[15, 8], method="end", mode="constant", constant_values=((0, 0), (1, 1), (2, 2)) ) - result = padder(np.zeros((3, 8, 4))) - np.testing.assert_allclose(result[:, 8:, :4], np.ones((3, 7, 4))) - np.testing.assert_allclose(result[:, :, 4:], np.ones((3, 15, 4)) + 1) + for p in TEST_NDARRAYS: + result = padder(p(np.zeros((3, 8, 4)))) + if isinstance(result, torch.Tensor): + result = result.cpu().numpy() + torch.testing.assert_allclose(result[:, 8:, :4], np.ones((3, 7, 4)), rtol=1e-7, atol=0) + torch.testing.assert_allclose(result[:, :, 4:], np.ones((3, 15, 4)) + 1, rtol=1e-7, atol=0) if __name__ == "__main__":