Repository: kornia/kornia

Issue:
Augmentations in nested AugmentationSequential container not applied to bboxes or keypoints
### Describe the bug

Augmentations in nested AugmentationSequential container not applied to bboxes or keypoints.

Since the nested AugmentationSequential module is not an instance of `_AugmentationBase` nothing is done for boxes or points in https://github.com/kornia/kornia/blob/4768008f137620e1d79f58b72e5f84311d81ccab/kornia/augmentation/container/utils.py#L83-L92


### Reproduction steps

```bash
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).expand(3, -1, -1)
points = torch.tensor([[[0.0, 0.0], [1.0, 1.0]]], device=device, dtype=dtype).expand(3, -1, -1)
aug = K.AugmentationSequential(K.AugmentationSequential(K.RandomAffine(degrees=90., p=1.)), data_keys=["input", "mask", "bbox_xyxy", "keypoints"])
```


### Expected behavior

Bboxes and points should have been modified according to rotation in affine transformation rotated, but they remain unchanged instead.

### Environment

```shell
PyTorch version: 1.10.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.27

Python version: 3.8.5 (default, Sep  4 2020, 07:30:14)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-4.15.0-162-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: NVIDIA TITAN X (Pascal)
GPU 1: NVIDIA TITAN X (Pascal)
GPU 2: NVIDIA TITAN Xp

Nvidia driver version: 470.82.00
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.1
[pip3] torch==1.10.0
[pip3] torchvision==0.11.1
[conda] blas                      1.0                         mkl    anaconda
[conda] cudatoolkit               11.3.1               ha36c431_9    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2019.4                      243    anaconda
[conda] mkl-service               2.3.0            py38he904b0f_0    anaconda
[conda] mkl_fft                   1.2.0            py38h23d657b_0    anaconda
[conda] mkl_random                1.1.0            py38h962f231_0    anaconda
[conda] numpy                     1.19.1           py38hbc911f0_0    anaconda
[conda] numpy-base                1.19.1           py38hfa32c7d_0    anaconda
[conda] pytorch                   1.10.0          py3.8_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchvision               0.11.1               py38_cu113    pytorch
```


### Additional context

_No response_


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
@@ -17,7 +17,7 @@
 
 
 def _get_geometric_only_param(
-    module: "kornia.augmentation.container.ImageSequential", param: List[ParamItem]
+    module: "kornia.augmentation.ImageSequential", param: List[ParamItem]
 ) -> List[ParamItem]:
     named_modules: Iterator[Tuple[str, nn.Module]] = module.get_forward_sequence(param)
 
@@ -84,6 +84,8 @@ def apply_trans(
         to_apply = None
         if isinstance(module, _AugmentationBase):
             to_apply = param.data['batch_prob']  # type: ignore
+        if isinstance(module, kornia.augmentation.ImageSequential):
+            to_apply = torch.ones(input.shape[0], device=input.device, dtype=input.dtype).bool()
 
         # If any inputs need to be transformed.
         if mat is not None and to_apply is not None and to_apply.sum() != 0:
@@ -114,7 +116,7 @@ def _get_transformation(
     ) -> Optional[torch.Tensor]:
 
         if (
-            isinstance(module, (GeometricAugmentationBase2D, kornia.augmentation.container.ImageSequential))
+            isinstance(module, (GeometricAugmentationBase2D, kornia.augmentation.ImageSequential))
             and param is None
         ):
             raise ValueError(f"Parameters of transformation matrix for {module} has not been computed.")
@@ -122,7 +124,7 @@ def _get_transformation(
         if isinstance(module, GeometricAugmentationBase2D):
             _param = cast(Dict[str, torch.Tensor], param.data)  # type: ignore
             mat = module.get_transformation_matrix(input, _param)
-        elif isinstance(module, kornia.augmentation.container.ImageSequential) and not module.is_intensity_only():
+        elif isinstance(module, kornia.augmentation.ImageSequential) and not module.is_intensity_only():
             _param = cast(List[ParamItem], param.data)  # type: ignore
             mat = module.get_transformation_matrix(input, _param)  # type: ignore
         else:
@@ -136,6 +138,7 @@ def _get_inverse_transformation(cls, transform: torch.Tensor) -> torch.Tensor:
 
 class InputApplyInverse(ApplyInverseImpl):
     """Apply and inverse transformations for (image) input tensors."""
+    data_key = DataKey.INPUT
 
     @classmethod
     def apply_trans(  # type: ignore
@@ -154,12 +157,15 @@ def apply_trans(  # type: ignore
             input, label = module(input, label=label, params=param.data)
         elif isinstance(module, (_AugmentationBase,)):
             input = module(input, params=param.data)
-        elif isinstance(module, kornia.augmentation.container.ImageSequential):
+        elif isinstance(module, kornia.augmentation.ImageSequential):
             temp = module.apply_inverse_func
             temp2 = module.return_label
             module.apply_inverse_func = InputApplyInverse
             module.return_label = True
-            input, label = module(input, label=label, params=param.data)
+            if isinstance(module, kornia.augmentation.AugmentationSequential):
+                input, label = module(input, label=label, params=param.data, data_keys=[cls.data_key])
+            else:
+                input, label = module(input, label=label, params=param.data)
             module.apply_inverse_func = temp
             module.return_label = temp2
         else:
@@ -184,7 +190,7 @@ def inverse(cls, input: torch.Tensor, module: nn.Module, param: Optional[ParamIt
         """
         if isinstance(module, GeometricAugmentationBase2D):
             input = module.inverse(input, params=None if param is None else cast(Dict, param.data))
-        elif isinstance(module, kornia.augmentation.container.ImageSequential):
+        elif isinstance(module, kornia.augmentation.ImageSequential):
             temp = module.apply_inverse_func
             module.apply_inverse_func = InputApplyInverse
             input = module.inverse(input, params=None if param is None else cast(List, param.data))
@@ -194,9 +200,10 @@ def inverse(cls, input: torch.Tensor, module: nn.Module, param: Optional[ParamIt
 
 class MaskApplyInverse(ApplyInverseImpl):
     """Apply and inverse transformations for mask tensors."""
+    data_key = DataKey.MASK
 
     @classmethod
-    def make_input_only_sequential(cls, module: "kornia.augmentation.container.ImageSequential") -> Callable:
+    def make_input_only_sequential(cls, module: "kornia.augmentation.ImageSequential") -> Callable:
         """Disable all other additional inputs (e.g. ) for ImageSequential."""
 
         def f(*args, **kwargs):
@@ -232,7 +239,7 @@ def apply_trans(
         if isinstance(module, GeometricAugmentationBase2D):
             _param = cast(Dict[str, torch.Tensor], _param)
             input = module(input, params=_param, return_transform=False)
-        elif isinstance(module, kornia.augmentation.container.ImageSequential) and not module.is_intensity_only():
+        elif isinstance(module, kornia.augmentation.ImageSequential) and not module.is_intensity_only():
             _param = cast(List[ParamItem], _param)
             temp = module.apply_inverse_func
             module.apply_inverse_func = MaskApplyInverse
@@ -255,7 +262,7 @@ def inverse(cls, input: torch.Tensor, module: nn.Module, param: Optional[ParamIt
         """
         if isinstance(module, GeometricAugmentationBase2D):
             input = module.inverse(input, params=None if param is None else cast(Dict, param.data))
-        elif isinstance(module, kornia.augmentation.container.ImageSequential):
+        elif isinstance(module, kornia.augmentation.ImageSequential):
             temp = module.apply_inverse_func
             module.apply_inverse_func = MaskApplyInverse
             input = module.inverse(input, params=None if param is None else cast(List, param.data))


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
@@ -418,7 +418,9 @@ def test_forward_and_inverse(self, random_apply, return_transform, device, dtype
                 K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0, return_transform=True)
             ),
             K.AugmentationSequential(
-                K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0, return_transform=True)
+                K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0, return_transform=True),
+                K.RandomAffine(360, p=1.0),
+                data_keys=["input", "mask", "bbox", "keypoints"]
             ),
             K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0),
             K.RandomAffine(360, p=1.0),