Repository: kornia/kornia

Issue:
[BUG] Pipelining multiple augmentations return wrong transformation matrix
## 🐛 Bug

Transformation matrices get calculated wrongly when fed though more than one augmentation.  

## To Reproduce

```
>>> import kornia as K
>>> import torch
>>> from kornia.augmentation import random_generator as rg
>>> input = torch.tensor([[1., 0., 0.],
                  [0., 1., 2.,],
                  [0., 0., 1.,]]).view((1,1,3,3))

>>> params = rg.random_rotation_generator(batch_size=1,degrees=torch.tensor([90.,90.]))
>>> hflip = K.augmentation.RandomHorizontalFlip(1.0,return_transform=True)
>>> rot = K.augmentation.RandomRotation(degrees=90.0, return_transform=True)
>>> rotated_img,T_rot = rot(input,params)
>>> final_img_true, T_hflip = hflip(rotated_img)
>>> T_true = T_hflip @ T_rot
>>> final_img,T = hflip(rot(input,params))

>>> assert torch.allclose(final_img_true,final_img), "Result images don't match."
>>> assert torch.allclose(T_true,T), "Transformations don't match."
AssertionError: Transformations don't match.

```
## Expected behavior

A pipeline of multiple the augmentation moduls should return the correct transformation matrix.

## Environment

PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.4 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.18.1

Python version: 3.6 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: GPU 0: GeForce RTX 2080 Ti
Nvidia driver version: 440.44
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.19.1
[pip3] torch==1.6.0
[pip3] torchvision==0.7.0



Discussion:

Resolving patch:
diff --git a/kornia/augmentation/augmentation.py b/kornia/augmentation/augmentation.py
--- a/kornia/augmentation/augmentation.py
+++ b/kornia/augmentation/augmentation.py
@@ -58,7 +58,7 @@ def forward(self, input: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
             output = self.apply_transform(input[0], self._params)
             transformation_matrix = self.compute_transformation(input[0], self._params)
             if return_transform:
-                return output, input[1] @ transformation_matrix
+                return output, transformation_matrix @ input[1]
             else:
                 return output, input[1]
 


Test patch:
diff --git a/test/augmentation/test_augmentation.py b/test/augmentation/test_augmentation.py
--- a/test/augmentation/test_augmentation.py
+++ b/test/augmentation/test_augmentation.py
@@ -1,3 +1,4 @@
+from unittest.mock import patch
 from typing import Union, Tuple
 
 import pytest
@@ -10,10 +11,81 @@
 import kornia
 import kornia.testing as utils  # test utils
 from kornia.constants import pi
-from kornia.augmentation import RandomHorizontalFlip, RandomVerticalFlip, ColorJitter, \
+from kornia.augmentation import AugmentationBase, RandomHorizontalFlip, RandomVerticalFlip, ColorJitter, \
     RandomErasing, RandomGrayscale, RandomRotation, RandomCrop, RandomResizedCrop, RandomMotionBlur
 
 
+class TestAugmentationBase:
+
+    def test_forward(self, device, dtype):
+        torch.manual_seed(42)
+        input = torch.rand((2, 3, 4, 5), device=device, dtype=dtype)
+        input_transform = torch.rand((2, 3, 3), device=device, dtype=dtype)
+        expected_output = torch.rand((2, 3, 4, 5), device=device, dtype=dtype)
+        expected_transform = torch.rand((2, 3, 3), device=device, dtype=dtype)
+        augmentation = AugmentationBase(return_transform=False)
+
+        with patch.object(augmentation, "apply_transform", autospec=True) as apply_transform, \
+                patch.object(augmentation, "generate_parameters", autospec=True) as generate_parameters, \
+                patch.object(augmentation, "compute_transformation", autospec=True) as compute_transformation:
+
+            # Calling the augmentation with a single tensor shall return the expected tensor using the generated params.
+            params = {"foo": 0}
+            generate_parameters.return_value = params
+            apply_transform.return_value = expected_output
+            compute_transformation.return_value = expected_transform
+            output = augmentation(input)
+            apply_transform.assert_called_once_with(input, params)
+            assert output is expected_output
+
+            # Calling the augmentation with a tensor and set return_transform shall
+            # return the expected tensor and transformation.
+            output, transformation = augmentation(input, return_transform=True)
+            assert output is expected_output
+            assert transformation is expected_transform
+
+            # Calling the augmentation with a tensor and params shall return the expected tensor using the given params.
+            params = {"bar": 1}
+            apply_transform.reset_mock()
+            generate_parameters.return_value = None
+            output = augmentation(input, params=params)
+            apply_transform.assert_called_once_with(input, params)
+            assert output is expected_output
+
+            # Calling the augmentation with a tensor,a transformation and set
+            # return_transform shall return the expected tensor and the proper
+            # transformation matrix.
+            expected_final_transformation = expected_transform @ input_transform
+            output, transformation = augmentation((input, input_transform), return_transform=True)
+            assert output is expected_output
+            assert torch.allclose(expected_final_transformation, transformation)
+            assert transformation.shape[0] == input.shape[0]
+
+    def test_gradcheck(self, device, dtype):
+        torch.manual_seed(42)
+
+        input = torch.rand((1, 1, 3, 3), device=device, dtype=dtype)
+        output = torch.rand((1, 1, 3, 3), device=device, dtype=dtype)
+        input_transform = torch.rand((1, 3, 3), device=device, dtype=dtype)
+        other_transform = torch.rand((1, 3, 3), device=device, dtype=dtype)
+
+        input = utils.tensor_to_gradcheck_var(input)  # to var
+        input_transform = utils.tensor_to_gradcheck_var(input_transform)  # to var
+        output = utils.tensor_to_gradcheck_var(output)  # to var
+        other_transform = utils.tensor_to_gradcheck_var(other_transform)  # to var
+
+        augmentation = AugmentationBase(return_transform=True)
+
+        with patch.object(augmentation, "apply_transform", autospec=True) as apply_transform, \
+                patch.object(augmentation, "generate_parameters", autospec=True) as generate_parameters, \
+                patch.object(augmentation, "compute_transformation", autospec=True) as compute_transformation:
+
+            apply_transform.return_value = output
+            compute_transformation.return_value = other_transform
+
+            assert gradcheck(augmentation, ((input, input_transform)), raise_exception=True)
+
+
 class TestRandomHorizontalFlip:
 
     def smoke_test(self, device):