From 2054382e167b1875b7932f09644b13b9f1dfb099 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Mon, 26 Aug 2024 12:52:17 +0800 Subject: [PATCH 1/6] init --- docs/source/augmentation.module.rst | 1 + kornia/augmentation/_2d/geometric/__init__.py | 2 +- kornia/augmentation/_2d/geometric/rotation.py | 98 +++++++++++++++++++ kornia/augmentation/__init__.py | 2 + tests/augmentation/test_augmentation.py | 78 +++++++++++++++ 5 files changed, 180 insertions(+), 1 deletion(-) diff --git a/docs/source/augmentation.module.rst b/docs/source/augmentation.module.rst index 99b0d7e233f..2279e997cd6 100644 --- a/docs/source/augmentation.module.rst +++ b/docs/source/augmentation.module.rst @@ -61,6 +61,7 @@ Geometric .. autoclass:: RandomHorizontalFlip .. autoclass:: RandomPerspective .. autoclass:: RandomResizedCrop +.. autoclass:: RandomRot90 .. autoclass:: RandomRotation .. autoclass:: RandomShear .. autoclass:: RandomThinPlateSpline diff --git a/kornia/augmentation/_2d/geometric/__init__.py b/kornia/augmentation/_2d/geometric/__init__.py index bfea3dac0a9..da02b025ccd 100644 --- a/kornia/augmentation/_2d/geometric/__init__.py +++ b/kornia/augmentation/_2d/geometric/__init__.py @@ -8,7 +8,7 @@ from kornia.augmentation._2d.geometric.perspective import RandomPerspective from kornia.augmentation._2d.geometric.resize import LongestMaxSize, Resize, SmallestMaxSize from kornia.augmentation._2d.geometric.resized_crop import RandomResizedCrop -from kornia.augmentation._2d.geometric.rotation import RandomRotation +from kornia.augmentation._2d.geometric.rotation import RandomRotation, RandomRot90 from kornia.augmentation._2d.geometric.shear import RandomShear from kornia.augmentation._2d.geometric.thin_plate_spline import RandomThinPlateSpline from kornia.augmentation._2d.geometric.translate import RandomTranslate diff --git a/kornia/augmentation/_2d/geometric/rotation.py b/kornia/augmentation/_2d/geometric/rotation.py index f2372b84873..f907e0edead 100644 --- a/kornia/augmentation/_2d/geometric/rotation.py +++ b/kornia/augmentation/_2d/geometric/rotation.py @@ -110,3 +110,101 @@ def inverse_transform( transform=as_tensor(transform, device=input.device, dtype=input.dtype), flags=flags, ) + + +class RandomRot90(GeometricAugmentationBase2D): + r"""Apply a random 90 * n degree rotation to a tensor image or a batch of tensor images. + + Args: + times: the range of n times 90 degree rotation needs to be applied. + resample: Default: the interpolation mode. + same_on_batch: apply the same transformation across the batch. + align_corners: interpolation flag. + p: probability of applying the transformation. + keepdim: whether to keep the output shape the same as input (True) or broadcast it + to the batch form (False). + + Shape: + - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)` + - Output: :math:`(B, C, H, W)` + + .. note:: + This function internally uses :func:`kornia.geometry.transform.affine`. + + Examples: + >>> rng = torch.manual_seed(1) + >>> input = torch.tensor([[1., 0., 0., 2.], + ... [0., 0., 0., 0.], + ... [0., 1., 2., 0.], + ... [0., 0., 1., 2.]]) + >>> aug = RandomRot90(times=(1, 1), p=1.) + >>> out = aug(input) + >>> out + tensor([[[[2.0000e+00, 0.0000e+00, 0.0000e+00, 2.0000e+00], + [0.0000e+00, 2.3842e-07, 2.0000e+00, 1.0000e+00], + [5.9605e-08, 3.5527e-15, 1.0000e+00, 0.0000e+00], + [1.0000e+00, 5.9605e-08, 0.0000e+00, 0.0000e+00]]]]) + >>> aug.transform_matrix + tensor([[[-4.3711e-08, 1.0000e+00, 1.1921e-07], + [-1.0000e+00, -4.3711e-08, 3.0000e+00], + [ 0.0000e+00, 0.0000e+00, 1.0000e+00]]]) + >>> inv = aug.inverse(out) + + To apply the exact augmenation again, you may take the advantage of the previous parameter state: + >>> input = torch.randn(1, 3, 32, 32) + >>> aug = RandomRot90(times=(-1, 1), p=1.) + >>> (aug(input) == aug(input, params=aug._params)).all() + tensor(True) + """ + + def __init__( + self, + times: Tuple[int, int], + resample: Union[str, int, Resample] = Resample.BILINEAR.name, + same_on_batch: bool = False, + align_corners: bool = True, + p: float = 0.5, + keepdim: bool = False, + ) -> None: + super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) + self._param_generator = rg.PlainUniformGenerator((times, "times", 0.0, (-3, 3))) + + self.flags = {"resample": Resample.get(resample), "align_corners": align_corners} + + def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor: + # TODO: Update to use `get_rotation_matrix2d` + angles: Tensor = 90. * params["times"].long().to(input) + + center: Tensor = _compute_tensor_center(input) + rotation_mat: Tensor = _compute_rotation_matrix(angles, center.expand(angles.shape[0], -1)) + + # rotation_mat is B x 2 x 3 and we need a B x 3 x 3 matrix + trans_mat: Tensor = eye_like(3, input, shared_memory=False) + trans_mat[:, 0] = rotation_mat[:, 0] + trans_mat[:, 1] = rotation_mat[:, 1] + + return trans_mat + + def apply_transform( + self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + ) -> Tensor: + if not isinstance(transform, Tensor): + raise TypeError(f"Expected the `transform` be a Tensor. Got {type(transform)}.") + + return affine(input, transform[..., :2, :3], flags["resample"].name.lower(), "zeros", flags["align_corners"]) + + def inverse_transform( + self, + input: Tensor, + flags: Dict[str, Any], + transform: Optional[Tensor] = None, + size: Optional[Tuple[int, int]] = None, + ) -> Tensor: + if not isinstance(transform, Tensor): + raise TypeError(f"Expected the `transform` be a Tensor. Got {type(transform)}.") + return self.apply_transform( + input, + params=self._params, + transform=as_tensor(transform, device=input.device, dtype=input.dtype), + flags=flags, + ) diff --git a/kornia/augmentation/__init__.py b/kornia/augmentation/__init__.py index f0338975b94..20f9f4d1446 100644 --- a/kornia/augmentation/__init__.py +++ b/kornia/augmentation/__init__.py @@ -48,6 +48,7 @@ RandomRain, RandomResizedCrop, RandomRGBShift, + RandomRot90, RandomRotation, RandomSaltAndPepperNoise, RandomSaturation, @@ -135,6 +136,7 @@ "RandomPlasmaBrightness", "RandomPlasmaContrast", "RandomResizedCrop", + "RandomRot90", "RandomRotation", "RandomRGBShift", "RandomSaltAndPepperNoise", diff --git a/tests/augmentation/test_augmentation.py b/tests/augmentation/test_augmentation.py index 939cdc6e3bb..c5f36cc323d 100644 --- a/tests/augmentation/test_augmentation.py +++ b/tests/augmentation/test_augmentation.py @@ -48,6 +48,7 @@ RandomRain, RandomResizedCrop, RandomRGBShift, + RandomRot90, RandomRotation, RandomSaltAndPepperNoise, RandomSaturation, @@ -697,6 +698,83 @@ def test_exception(self): self._create_augmentation_from_params(degrees=(360.0, -360.0)) +class TestRandomRot90(CommonTests): + possible_params: Dict["str", Tuple] = { + "times": ((-3, 3), (1, 1)), + "resample": (0, Resample.BILINEAR.name, Resample.BILINEAR), + "align_corners": (False, True), + } + _augmentation_cls = RandomRot90 + _default_param_set: Dict["str", Any] = { + "times": (-3, 3), + "align_corners": True, + } + + @pytest.fixture(params=[_default_param_set], scope="class") + def param_set(self, request): + return request.param + + @pytest.mark.parametrize( + "input_shape,expected_output_shape", + [((3, 4, 5), (1, 3, 4, 5)), ((2, 3, 4, 5), (2, 3, 4, 5))], + ) + def test_cardinality(self, input_shape, expected_output_shape): + self._test_cardinality_implementation( + input_shape=input_shape, + expected_output_shape=expected_output_shape, + params=self._default_param_set, + ) + + def test_random_p_1(self): + torch.manual_seed(42) + + input_tensor = torch.tensor( + [[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]], + device=self.device, + dtype=self.dtype, + ) + expected_output = torch.tensor( + [[[[0.3, 0.6, 0.9], [0.2, 0.5, 0.8], [0.1, 0.4, 0.7]]]], + device=self.device, + dtype=self.dtype, + ) + + parameters = {"times": (1, 2), "align_corners": True} + self._test_random_p_1_implementation( + input_tensor=input_tensor, + expected_output=expected_output, + params=parameters, + ) + + def test_batch(self): + if self.dtype == torch.float16: + pytest.skip("not work for half-precision") + + torch.manual_seed(12) + + input_tensor = torch.tensor( + [[[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]]], + device=self.device, + dtype=self.dtype, + ).repeat((2, 1, 1, 1)) + expected_output = input_tensor + expected_transformation = kornia.eye_like(3, input_tensor) + parameters = {"times": (0, 0), "align_corners": True} + self._test_random_p_1_implementation( + input_tensor=input_tensor, + expected_output=expected_output, + expected_transformation=expected_transformation, + params=parameters, + ) + + def test_exception(self): + # Wrong type + with pytest.raises(TypeError): + self._create_augmentation_from_params(times="") + with pytest.raises(ValueError): + self._create_augmentation_from_params(times=(30, 60), align_corners=0) + + class TestRandomGrayscaleAlternative(CommonTests): possible_params: Dict["str", Tuple] = {} From 0ae40e4731ad374d77c3c2302406ef602e28a163 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 26 Aug 2024 04:55:11 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/augmentation/_2d/geometric/__init__.py | 2 +- kornia/augmentation/_2d/geometric/rotation.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/kornia/augmentation/_2d/geometric/__init__.py b/kornia/augmentation/_2d/geometric/__init__.py index da02b025ccd..078f2e0b1b3 100644 --- a/kornia/augmentation/_2d/geometric/__init__.py +++ b/kornia/augmentation/_2d/geometric/__init__.py @@ -8,7 +8,7 @@ from kornia.augmentation._2d.geometric.perspective import RandomPerspective from kornia.augmentation._2d.geometric.resize import LongestMaxSize, Resize, SmallestMaxSize from kornia.augmentation._2d.geometric.resized_crop import RandomResizedCrop -from kornia.augmentation._2d.geometric.rotation import RandomRotation, RandomRot90 +from kornia.augmentation._2d.geometric.rotation import RandomRot90, RandomRotation from kornia.augmentation._2d.geometric.shear import RandomShear from kornia.augmentation._2d.geometric.thin_plate_spline import RandomThinPlateSpline from kornia.augmentation._2d.geometric.translate import RandomTranslate diff --git a/kornia/augmentation/_2d/geometric/rotation.py b/kornia/augmentation/_2d/geometric/rotation.py index f907e0edead..4b3b5cbe703 100644 --- a/kornia/augmentation/_2d/geometric/rotation.py +++ b/kornia/augmentation/_2d/geometric/rotation.py @@ -173,7 +173,7 @@ def __init__( def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor: # TODO: Update to use `get_rotation_matrix2d` - angles: Tensor = 90. * params["times"].long().to(input) + angles: Tensor = 90.0 * params["times"].long().to(input) center: Tensor = _compute_tensor_center(input) rotation_mat: Tensor = _compute_rotation_matrix(angles, center.expand(angles.shape[0], -1)) From 3ece44d206b9dae4ddb2ded339045cdb144feb0d Mon Sep 17 00:00:00 2001 From: shijianjian Date: Tue, 27 Aug 2024 11:14:48 +0800 Subject: [PATCH 3/6] update --- docs/source/augmentation.module.rst | 2 +- kornia/augmentation/_2d/geometric/__init__.py | 2 +- kornia/augmentation/_2d/geometric/rotation.py | 26 ++++++++++--------- kornia/augmentation/__init__.py | 4 +-- tests/augmentation/test_augmentation.py | 6 ++--- 5 files changed, 21 insertions(+), 19 deletions(-) diff --git a/docs/source/augmentation.module.rst b/docs/source/augmentation.module.rst index 2279e997cd6..e87fef7c253 100644 --- a/docs/source/augmentation.module.rst +++ b/docs/source/augmentation.module.rst @@ -61,7 +61,7 @@ Geometric .. autoclass:: RandomHorizontalFlip .. autoclass:: RandomPerspective .. autoclass:: RandomResizedCrop -.. autoclass:: RandomRot90 +.. autoclass:: RandomRotation90 .. autoclass:: RandomRotation .. autoclass:: RandomShear .. autoclass:: RandomThinPlateSpline diff --git a/kornia/augmentation/_2d/geometric/__init__.py b/kornia/augmentation/_2d/geometric/__init__.py index 078f2e0b1b3..d51a7524eee 100644 --- a/kornia/augmentation/_2d/geometric/__init__.py +++ b/kornia/augmentation/_2d/geometric/__init__.py @@ -8,7 +8,7 @@ from kornia.augmentation._2d.geometric.perspective import RandomPerspective from kornia.augmentation._2d.geometric.resize import LongestMaxSize, Resize, SmallestMaxSize from kornia.augmentation._2d.geometric.resized_crop import RandomResizedCrop -from kornia.augmentation._2d.geometric.rotation import RandomRot90, RandomRotation +from kornia.augmentation._2d.geometric.rotation import RandomRotation, RandomRotation90 from kornia.augmentation._2d.geometric.shear import RandomShear from kornia.augmentation._2d.geometric.thin_plate_spline import RandomThinPlateSpline from kornia.augmentation._2d.geometric.translate import RandomTranslate diff --git a/kornia/augmentation/_2d/geometric/rotation.py b/kornia/augmentation/_2d/geometric/rotation.py index 4b3b5cbe703..582d9a17f4b 100644 --- a/kornia/augmentation/_2d/geometric/rotation.py +++ b/kornia/augmentation/_2d/geometric/rotation.py @@ -112,7 +112,7 @@ def inverse_transform( ) -class RandomRot90(GeometricAugmentationBase2D): +class RandomRotation90(GeometricAugmentationBase2D): r"""Apply a random 90 * n degree rotation to a tensor image or a batch of tensor images. Args: @@ -129,30 +129,32 @@ class RandomRot90(GeometricAugmentationBase2D): - Output: :math:`(B, C, H, W)` .. note:: - This function internally uses :func:`kornia.geometry.transform.affine`. + This function internally uses :func:`kornia.geometry.transform.affine`. This version is relatively + slow as it operates based on affine transformations. Examples: >>> rng = torch.manual_seed(1) + >>> torch.set_printoptions(sci_mode=False) >>> input = torch.tensor([[1., 0., 0., 2.], ... [0., 0., 0., 0.], ... [0., 1., 2., 0.], ... [0., 0., 1., 2.]]) - >>> aug = RandomRot90(times=(1, 1), p=1.) + >>> aug = RandomRotation90(times=(1, 1), p=1.) >>> out = aug(input) >>> out - tensor([[[[2.0000e+00, 0.0000e+00, 0.0000e+00, 2.0000e+00], - [0.0000e+00, 2.3842e-07, 2.0000e+00, 1.0000e+00], - [5.9605e-08, 3.5527e-15, 1.0000e+00, 0.0000e+00], - [1.0000e+00, 5.9605e-08, 0.0000e+00, 0.0000e+00]]]]) + tensor([[[[ 2.0000, 0.0000, 0.0000, 2.0000], + [ 0.0000, 0.0000, 2.0000, 1.0000], + [ 0.0000, 0.0000, 1.0000, 0.0000], + [ 1.0000, 0.0000, 0.0000, 0.0000]]]]) >>> aug.transform_matrix - tensor([[[-4.3711e-08, 1.0000e+00, 1.1921e-07], - [-1.0000e+00, -4.3711e-08, 3.0000e+00], - [ 0.0000e+00, 0.0000e+00, 1.0000e+00]]]) + tensor([[[ -0.0000, 1.0000, 0.0000], + [ -1.0000, -0.0000, 3.0000], + [ 0.0000, 0.0000, 1.0000]]]) >>> inv = aug.inverse(out) To apply the exact augmenation again, you may take the advantage of the previous parameter state: >>> input = torch.randn(1, 3, 32, 32) - >>> aug = RandomRot90(times=(-1, 1), p=1.) + >>> aug = RandomRotation90(times=(-1, 1), p=1.) >>> (aug(input) == aug(input, params=aug._params)).all() tensor(True) """ @@ -173,7 +175,7 @@ def __init__( def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor: # TODO: Update to use `get_rotation_matrix2d` - angles: Tensor = 90.0 * params["times"].long().to(input) + angles: Tensor = 90. * params["times"].round().to(input) center: Tensor = _compute_tensor_center(input) rotation_mat: Tensor = _compute_rotation_matrix(angles, center.expand(angles.shape[0], -1)) diff --git a/kornia/augmentation/__init__.py b/kornia/augmentation/__init__.py index 20f9f4d1446..ee5b5ddca06 100644 --- a/kornia/augmentation/__init__.py +++ b/kornia/augmentation/__init__.py @@ -48,7 +48,7 @@ RandomRain, RandomResizedCrop, RandomRGBShift, - RandomRot90, + RandomRotation90, RandomRotation, RandomSaltAndPepperNoise, RandomSaturation, @@ -136,7 +136,7 @@ "RandomPlasmaBrightness", "RandomPlasmaContrast", "RandomResizedCrop", - "RandomRot90", + "RandomRotation90", "RandomRotation", "RandomRGBShift", "RandomSaltAndPepperNoise", diff --git a/tests/augmentation/test_augmentation.py b/tests/augmentation/test_augmentation.py index c5f36cc323d..b762b063bc5 100644 --- a/tests/augmentation/test_augmentation.py +++ b/tests/augmentation/test_augmentation.py @@ -48,7 +48,7 @@ RandomRain, RandomResizedCrop, RandomRGBShift, - RandomRot90, + RandomRotation90, RandomRotation, RandomSaltAndPepperNoise, RandomSaturation, @@ -698,13 +698,13 @@ def test_exception(self): self._create_augmentation_from_params(degrees=(360.0, -360.0)) -class TestRandomRot90(CommonTests): +class TestRandomRotation90(CommonTests): possible_params: Dict["str", Tuple] = { "times": ((-3, 3), (1, 1)), "resample": (0, Resample.BILINEAR.name, Resample.BILINEAR), "align_corners": (False, True), } - _augmentation_cls = RandomRot90 + _augmentation_cls = RandomRotation90 _default_param_set: Dict["str", Any] = { "times": (-3, 3), "align_corners": True, From c70fbf568d4d4b2726a025329eb34e5c549d6832 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 27 Aug 2024 03:15:18 +0000 Subject: [PATCH 4/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/augmentation/_2d/geometric/rotation.py | 2 +- kornia/augmentation/__init__.py | 2 +- tests/augmentation/test_augmentation.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/kornia/augmentation/_2d/geometric/rotation.py b/kornia/augmentation/_2d/geometric/rotation.py index 582d9a17f4b..f348e774cda 100644 --- a/kornia/augmentation/_2d/geometric/rotation.py +++ b/kornia/augmentation/_2d/geometric/rotation.py @@ -175,7 +175,7 @@ def __init__( def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor: # TODO: Update to use `get_rotation_matrix2d` - angles: Tensor = 90. * params["times"].round().to(input) + angles: Tensor = 90.0 * params["times"].round().to(input) center: Tensor = _compute_tensor_center(input) rotation_mat: Tensor = _compute_rotation_matrix(angles, center.expand(angles.shape[0], -1)) diff --git a/kornia/augmentation/__init__.py b/kornia/augmentation/__init__.py index ee5b5ddca06..466dd6ffe4e 100644 --- a/kornia/augmentation/__init__.py +++ b/kornia/augmentation/__init__.py @@ -48,8 +48,8 @@ RandomRain, RandomResizedCrop, RandomRGBShift, - RandomRotation90, RandomRotation, + RandomRotation90, RandomSaltAndPepperNoise, RandomSaturation, RandomSharpness, diff --git a/tests/augmentation/test_augmentation.py b/tests/augmentation/test_augmentation.py index b762b063bc5..72da82a1a77 100644 --- a/tests/augmentation/test_augmentation.py +++ b/tests/augmentation/test_augmentation.py @@ -48,8 +48,8 @@ RandomRain, RandomResizedCrop, RandomRGBShift, - RandomRotation90, RandomRotation, + RandomRotation90, RandomSaltAndPepperNoise, RandomSaturation, RandomSnow, From f204de962e845ea45c880132da8c92c7263c7185 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Tue, 27 Aug 2024 12:49:25 +0800 Subject: [PATCH 5/6] update --- tests/augmentation/test_augmentation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/augmentation/test_augmentation.py b/tests/augmentation/test_augmentation.py index 72da82a1a77..9873d9bf1ef 100644 --- a/tests/augmentation/test_augmentation.py +++ b/tests/augmentation/test_augmentation.py @@ -739,7 +739,7 @@ def test_random_p_1(self): dtype=self.dtype, ) - parameters = {"times": (1, 2), "align_corners": True} + parameters = {"times": (1, 1), "align_corners": True} self._test_random_p_1_implementation( input_tensor=input_tensor, expected_output=expected_output, From bdae0ba9c617a915911d20bc22487b636328cca4 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Tue, 27 Aug 2024 15:01:29 +0800 Subject: [PATCH 6/6] update --- kornia/augmentation/_2d/geometric/rotation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/kornia/augmentation/_2d/geometric/rotation.py b/kornia/augmentation/_2d/geometric/rotation.py index f348e774cda..c109884bbc9 100644 --- a/kornia/augmentation/_2d/geometric/rotation.py +++ b/kornia/augmentation/_2d/geometric/rotation.py @@ -151,6 +151,7 @@ class RandomRotation90(GeometricAugmentationBase2D): [ -1.0000, -0.0000, 3.0000], [ 0.0000, 0.0000, 1.0000]]]) >>> inv = aug.inverse(out) + >>> torch.set_printoptions(profile='default') To apply the exact augmenation again, you may take the advantage of the previous parameter state: >>> input = torch.randn(1, 3, 32, 32)