Repository: kornia/kornia

Issue:
bbox and keypoints won't work with RandomRotation
### Describe the bug

Using the latest version of kornia(0.6.3), RandomRotation cannot product correct bbox and keypoints while the mask is corrrect.

Here is the script(which modified from https://kornia-tutorials.readthedocs.io/en/latest/data_augmentation_sequential.html):
``` python
from matplotlib import pyplot as plt
import numpy as np
import torch
import cv2

from kornia import augmentation as K
from kornia.augmentation import AugmentationSequential
from kornia.geometry import bbox_to_mask
from kornia.utils import image_to_tensor, tensor_to_image
from torchvision.transforms import transforms

to_tensor = transforms.ToTensor()
to_pil = transforms.ToPILImage()

def plot_resulting_image(img, bbox, keypoints, mask):
    img = img * mask
    img_draw = cv2.polylines(np.array(to_pil(img)), bbox.numpy(), isClosed=True, color=(255, 0, 0))
    for k in keypoints[0]:
        img_draw = cv2.circle(img_draw, tuple(k.numpy()[:2]), radius=6, color=(255, 0, 0), thickness=-1)
    return img_draw

img = cv2.imread("panda.jpg", cv2.IMREAD_COLOR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
h, w = img.shape[:2]

img_tensor = image_to_tensor(img)
img_tensor = img_tensor.float() / 255.

aug_list = AugmentationSequential(
    K.RandomRotation(45.15,p=1.0),
    data_keys=["input", "bbox", "keypoints", "mask"],
    return_transform=True,
    same_on_batch=False,
)

bbox = torch.tensor([[[355,10],[660,10],[660,250],[355,250]]])
keypoints = torch.tensor([[[465, 115], [545, 116]]])
mask = bbox_to_mask(torch.tensor([[[155,0],[900,0],[900,400],[155,400]]]), w, h).float()


out_tensor = aug_list(img_tensor, bbox.float(), keypoints.float(), mask)

img_out = plot_resulting_image(
    out_tensor[0][0][0],
    out_tensor[1][0:1].int(),
    out_tensor[2][0:1].int(),
    out_tensor[3][0,0],
)
plt.imshow(img_out); plt.axis('off');
plt.show()
```

On my machine, this produce:
![image](https://user-images.githubusercontent.com/18608159/153137959-960676fe-5d80-4238-9807-9ac5f438a357.png)
, which is obvious not correct.


### Reproduction steps

```bash
1. `wget https://tinypng.com/images/social/website.jpg -O panda.jpg`
2. run the script above
3.
...
```


### Expected behavior

the keypoints and bbox should be on the panda's face.

### Environment

```shell
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```
- PyTorch Version (e.g., 1.0):
- OS (e.g., Linux):
- How you installed PyTorch (`conda`, `pip`, source):
- Build command you used (if compiling from source):
- Python version:
- CUDA/cuDNN version:
- GPU models and configuration:
- Any other relevant information:
```


### Additional context

Collecting environment information...
PyTorch version: 1.10.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.25

Python version: 3.6.9 (default, Dec  8 2021, 21:08:43)  [GCC 8.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-91-generic-x86_64-with-Ubuntu-18.04-bionic
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: GPU 0: GeForce GTX 1050 Ti
Nvidia driver version: 440.33.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.0
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.2
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.19.5
[pip3] numpy-quaternion==2021.11.4.15.26.3
[pip3] pytorch-lightning==1.5.5
[pip3] torch==1.10.0
[pip3] torch-tb-profiler==0.2.1
[pip3] torchmetrics==0.5.1
[pip3] torchvision==0.11.1
[conda] blas                      1.0                         mkl  
[conda] mkl                       2021.2.0           h06a4308_296  
[conda] mkl-service               2.3.0            py38h27cfd23_1  
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.1            py38ha9443f7_2  
[conda] mypy_extensions           0.4.3                    py38_0  
[conda] numpy                     1.20.1           py38h93e21f0_0  
[conda] numpy-base                1.20.1           py38h7d8b39e_0  
[conda] numpydoc                  1.1.0              pyhd3eb1b0_1


Discussion:

Resolving patch:
diff --git a/kornia/augmentation/container/utils.py b/kornia/augmentation/container/utils.py
--- a/kornia/augmentation/container/utils.py
+++ b/kornia/augmentation/container/utils.py
@@ -79,8 +79,12 @@ def apply_trans(
                 to apply transformations.
             param: the corresponding parameters to the module.
         """
-
-        mat: Optional[torch.Tensor] = cls._get_transformation(input, module, param)
+        mat: Optional[torch.Tensor]
+        if hasattr(module, "_transform_matrix") and module._transform_matrix is not None:
+            mat = cast(torch.Tensor, module._transform_matrix)
+        else:
+            mat = cls._get_transformation(input, module, param)
+        mat = torch.as_tensor(mat, device=input.device, dtype=input.dtype)
         to_apply = None
         if isinstance(module, _AugmentationBase):
             to_apply = param.data['batch_prob']  # type: ignore
@@ -103,7 +107,12 @@ def inverse(cls, input: torch.Tensor, module: nn.Module, param: Optional[ParamIt
                 to apply transformations.
             param: the corresponding parameters to the module.
         """
-        mat: Optional[torch.Tensor] = cls._get_transformation(input, module, param)
+        mat: Optional[torch.Tensor]
+        if hasattr(module, "_transform_matrix") and module._transform_matrix is not None:
+            mat = cast(torch.Tensor, module._transform_matrix)
+        else:
+            mat = cls._get_transformation(input, module, param)
+        mat = torch.as_tensor(mat, device=input.device, dtype=input.dtype)
 
         if mat is not None:
             transform: torch.Tensor = cls._get_inverse_transformation(mat)


Test patch:
diff --git a/test/augmentation/test_container.py b/test/augmentation/test_container.py
--- a/test/augmentation/test_container.py
+++ b/test/augmentation/test_container.py
@@ -363,6 +363,7 @@ def test_random_erasing(self, device, dtype):
         assert torch.all(out[1][out[0] == fill_value] == 0.)
 
     def test_random_crops(self, device, dtype):
+        torch.manual_seed(233)
         input = torch.randn(3, 3, 3, 3, device=device, dtype=dtype)
         bbox = torch.tensor(
             [[[1.0, 1.0, 2.0, 2.0], [0.0, 0.0, 1.0, 2.0], [0.0, 0.0, 2.0, 1.0]]], device=device, dtype=dtype
@@ -370,6 +371,7 @@ def test_random_crops(self, device, dtype):
         points = torch.tensor([[[0.0, 0.0], [1.0, 1.0]]], device=device, dtype=dtype).expand(3, -1, -1)
         aug = K.AugmentationSequential(
             K.RandomCrop((3, 3), padding=1, cropping_mode='resample', fill=0),
+            K.RandomAffine((360., 360.), p=1.),
             data_keys=["input", "mask", "bbox_xyxy", "keypoints"],
         )