Repository: kornia/kornia

Issue:
Nearest resampling mode not used for masks in nested AugmentationSequential containers
### Describe the bug

The resampling mode for masks must always be Nearest, otherways values outside of the set of valid values for the mask can appear. When having an `AugmentationSequential` container with affine transformations inside another `AugmentationSequential` container, the resampling mode used for the masks is incorrect, leading to invalid masks.

The nearest neighbor interpolation for masks was added in https://github.com/kornia/kornia/pull/1630. In the original fork for the PR the behavior was correct (I have been using it since then with this exact use case) so something must have gone wrong before merging it. However, I cannot find the problem on master. If someone can take a look and point me in the right direction I will be happy to submit a new PR.

### Reproduction steps

```bash
import kornia
import kornia.augmentation as K
import numpy as np
import torch
from kornia.constants import DataKey
from torchvision.transforms import functional as F

img = torch.zeros((3,5,5))
img[0,2,2] = 1.
mask = torch.zeros((1,5,5))
mask[...,2,2] = 1.

data_keys= [DataKey.INPUT, DataKey.MASK]
tf_pool = {
        'ShearX': K.RandomAffine(degrees=0., shear=17, p=1.),  # shear factor .3
        'ShearY': K.RandomAffine(
            degrees=0., shear=(0, 0, -17, 17), p=1.),
        'TranslateX': K.RandomAffine(degrees=0., translate=(.15, 0), p=1.),
        'TranslateY': K.RandomAffine(degrees=0., translate=(0, .15), p=1.),
        'Rotate': K.RandomAffine(degrees=30, p=1.),
}
aug = kornia.augmentation.AugmentationSequential(
    K.AugmentationSequential(*list(tf_pool.values()), random_apply=2, data_keys=data_keys),
    data_keys=data_keys
)

transformed = aug(img, mask)
img_transformed_inv = aug.inverse(*transformed)

F.to_pil_image(
    (img.squeeze(dim=0).transpose(0, 2).transpose(0, 1).numpy()*255).astype(dtype=np.uint8)
).save('in.png')
F.to_pil_image(
    (mask.transpose(0, 2).transpose(0, 1).numpy()*255).astype(dtype=np.uint8)
).save('mask.png')
F.to_pil_image(
    (transformed[0].squeeze(dim=0).transpose(0, 2).transpose(0, 1).numpy()*255).astype(dtype=np.uint8)
).save('in_t.png')
F.to_pil_image(
    (transformed[1].squeeze(dim=0).transpose(0, 2).transpose(0, 1).numpy()*255).astype(dtype=np.uint8)
).save('mask_t.png')

print(img.unique())
print(transformed[0].unique())

print(mask.unique())
print(transformed[1].unique())
```
Output:
```
tensor([0., 1.])
tensor([0.0000, 0.0533, 0.0533, 0.0533, 0.0533, 1.0000])
tensor([0., 1.])
tensor([0.0000, 0.0533, 0.0533, 0.0533, 0.0533, 1.0000])
```

### Expected behavior

The values that the augmented mask can take are only those in the original mask.

### Environment

```shell
Collecting environment information...
PyTorch version: 1.13.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
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.10.8 (main, Nov 24 2022, 14:13:03) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-4.15.0-196-generic-x86_64-with-glibc2.27
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: 520.61.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==1.13.0
[pip3] torchvision==0.14.0
[conda] numpy                     1.23.5                   pypi_0    pypi
[conda] torch                     1.13.0                   pypi_0    pypi
[conda] torchvision               0.14.0                   pypi_0    pypi
```


### Additional context

_No response_


Discussion:
@shijianjian 
I think I got a working solution after all... Fighting with some corner cases now. I will create a PR soon.
@miquelmarti thanks so much. Feel free also to propose anything that can help with the API.

Resolving patch:
diff --git a/kornia/augmentation/container/augment.py b/kornia/augmentation/container/augment.py
--- a/kornia/augmentation/container/augment.py
+++ b/kornia/augmentation/container/augment.py
@@ -156,7 +156,7 @@ def __init__(
             random_apply_weights=random_apply_weights,
         )
 
-        self.data_keys = [DataKey.get(inp) for inp in data_keys]
+        self.data_keys: List[DataKey] = [DataKey.get(inp) for inp in data_keys]
 
         if not all(in_type in DataKey for in_type in self.data_keys):
             raise AssertionError(f"`data_keys` must be in {DataKey}. Got {data_keys}.")
@@ -200,15 +200,11 @@ def inverse(  # type: ignore
         ``self.data_keys`` by default.
         """
         if data_keys is None:
-            data_keys = cast(List[Union[str, int, DataKey]], self.data_keys)
-
-        _data_keys: List[DataKey] = [DataKey.get(inp) for inp in data_keys]
+            _data_keys = self.data_keys
+        else:
+            _data_keys = [DataKey.get(inp) for inp in data_keys]
 
-        if len(args) != len(data_keys):
-            raise AssertionError(
-                "The number of inputs must align with the number of data_keys, "
-                f"Got {len(args)} and {len(data_keys)}."
-            )
+        self._validate_args_datakeys(*args, data_keys=_data_keys)
 
         args = self._arguments_preproc(*args, data_keys=_data_keys)
 
@@ -220,11 +216,11 @@ def inverse(  # type: ignore
                 )
             params = self._params
 
-        outputs: List[Tensor] = [None] * len(data_keys)  # type: ignore
-        for idx, (arg, dcate) in enumerate(zip(args, data_keys)):
+        outputs: List[Tensor] = [None] * len(_data_keys)  # type: ignore
+        for idx, (arg, dcate) in enumerate(zip(args, _data_keys)):
 
-            if DataKey.INPUT in self.extra_args:
-                extra_args = self.extra_args[DataKey.INPUT]
+            if dcate in self.extra_args:
+                extra_args = self.extra_args[dcate]
             else:
                 extra_args = {}
 
@@ -324,12 +320,11 @@ def forward(  # type: ignore
         data_keys: Optional[List[Union[str, int, DataKey]]] = None,
     ) -> Union[Tensor, Tuple[Tensor, Optional[Tensor]], List[Tensor], Tuple[List[Tensor], Optional[Tensor]]]:
         """Compute multiple tensors simultaneously according to ``self.data_keys``."""
-        _data_keys: List[DataKey]
         if data_keys is None:
             _data_keys = self.data_keys
         else:
             _data_keys = [DataKey.get(inp) for inp in data_keys]
-            self.data_keys = _data_keys
+
         self._validate_args_datakeys(*args, data_keys=_data_keys)
 
         args = self._arguments_preproc(*args, data_keys=_data_keys)
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
@@ -213,7 +213,12 @@ def inverse(
             temp = module.apply_inverse_func
             module.apply_inverse_func = InputApplyInverse
             if isinstance(module, kornia.augmentation.AugmentationSequential):
-                input = cast(Tensor, module.inverse(input, params=None if param is None else cast(List, param.data)))
+                input = cast(
+                    Tensor,
+                    module.inverse(
+                        input, params=None if param is None else cast(List, param.data), data_keys=[cls.data_key]
+                    ),
+                )
             else:
                 input = module.inverse(
                     input, params=None if param is None else cast(List, param.data), extra_args=extra_args
@@ -275,7 +280,14 @@ def apply_trans(
             temp = module.apply_inverse_func
             module.apply_inverse_func = MaskApplyInverse
             geo_param: List[ParamItem] = _get_geometric_only_param(module, _param)
-            input = cls.make_input_only_sequential(module)(input, label=None, params=geo_param)
+            if isinstance(module, kornia.augmentation.AugmentationSequential):
+                input = cls.make_input_only_sequential(module)(
+                    input, label=None, params=geo_param, data_keys=[cls.data_key]
+                )
+            else:
+                input = cls.make_input_only_sequential(module)(
+                    input, label=None, params=geo_param, extra_args=extra_args
+                )
             module.apply_inverse_func = temp
         else:
             pass  # No need to update anything
@@ -299,7 +311,17 @@ def inverse(
         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))
+            if isinstance(module, kornia.augmentation.AugmentationSequential):
+                input = cast(
+                    Tensor,
+                    module.inverse(
+                        input, params=None if param is None else cast(List, param.data), data_keys=[cls.data_key]
+                    ),
+                )
+            else:
+                input = module.inverse(
+                    input, params=None if param is None else cast(List, param.data), extra_args=extra_args
+                )
             module.apply_inverse_func = temp
         return input
 


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
@@ -440,12 +440,14 @@ def test_forward_and_inverse(self, random_apply, device, dtype):
         assert out[1].shape == mask.shape
         assert out[2].shape == bbox.shape
         assert out[3].shape == keypoints.shape
+        assert set(out[1].unique().tolist()).issubset(set(mask.unique().tolist()))
 
         out_inv = aug.inverse(*out)
         assert out_inv[0].shape == inp.shape
         assert out_inv[1].shape == mask.shape
         assert out_inv[2].shape == bbox.shape
         assert out_inv[3].shape == keypoints.shape
+        assert set(out_inv[1].unique().tolist()).issubset(set(mask.unique().tolist()))
 
         if random_apply is False:
             reproducibility_test((inp, mask, bbox, keypoints), aug)