Repository: kornia/kornia

Issue:
Padding error with `RandomCrop`
### Describe the bug

Tests are failing because of the following test `TestAugmentationSequential::test_random_crops` possibly introduced by the combination of be104a5 and 3f8e7a3.

Output error:

```python
 if mode.startswith("vertices"):
            batched = boxes.ndim == 4
            if not (3 <= boxes.ndim <= 4 and boxes.shape[-2:] == torch.Size([4, 2])):
>               raise ValueError(f"Boxes shape must be (N, 4, 2) or (B, N, 4, 2) when {mode} mode. Got {boxes.shape}.")
E               ValueError: Boxes shape must be (N, 4, 2) or (B, N, 4, 2) when vertices_plus mode. Got torch.Size([3, 3, 4]).
```

NOTE: This is a major bug, for now test is disabled until this is fixed.

### Reproduction steps

```bash
pytest test/augmentation/test_container.py::TestAugmentationSequential::test_random_crops
```


### Expected behavior

Pass the tests and enable the tests.

### Environment

```shell
PyTorch version: 1.10.1
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.20.5
Libc version: glibc-2.17

Python version: 3.7.11 (default, Jul 27 2021, 14:32:16)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.11.0-43-generic-x86_64-with-debian-bullseye-sid
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1650 with Max-Q Design
Nvidia driver version: 470.86
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.930
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.2
[pip3] pytest-mypy==0.8.1
[pip3] torch==1.10.1
[pip3] torchvision==0.11.2
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py37h7f8727e_0  
[conda] mkl_fft                   1.3.1            py37hd3c417c_0  
[conda] mkl_random                1.2.2            py37h51133e4_0  
[conda] mypy                      0.930                    pypi_0    pypi
[conda] mypy-extensions           0.4.3                    pypi_0    pypi
[conda] numpy                     1.21.2           py37h20f2e39_0  
[conda] numpy-base                1.21.2           py37h79a1101_0  
[conda] pytest-mypy               0.8.1                    pypi_0    pypi
[conda] pytorch                   1.10.1          py3.7_cuda10.2_cudnn7.6.5_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchvision               0.11.2               py37_cu102    pytorch
```


### Additional context

_No response_


Discussion:

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
@@ -181,11 +181,17 @@ def inverse(  # type: ignore
         """
         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]
+
         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)}."
             )
+
+        args = self._arguments_preproc(*args, data_keys=_data_keys)
+
         if params is None:
             if self._params is None:
                 raise ValueError(
@@ -194,10 +200,15 @@ def inverse(  # type: ignore
                 )
             params = self._params
 
-        outputs = []
-        for input, dcate in zip(args, data_keys):
-            if dcate == DataKey.INPUT and isinstance(input, (tuple, list)):
-                input, _ = input  # ignore the transformation matrix whilst inverse
+        outputs: List[torch.Tensor] = [None] * len(data_keys)  # type: ignore
+        for idx, (arg, dcate) in enumerate(zip(args, data_keys)):
+            if dcate == DataKey.INPUT and isinstance(arg, (tuple, list)):
+                input, _ = arg  # ignore the transformation matrix whilst inverse
+            # Using tensors straight-away
+            if isinstance(arg, (Boxes,)):
+                input = arg.data  # all boxes are in (B, N, 4, 2) format now.
+            else:
+                input = arg
             for (name, module), param in zip_longest(list(self.get_forward_sequence(params))[::-1], params[::-1]):
                 if isinstance(module, (_AugmentationBase, ImageSequential)):
                     param = params[name] if name in params else param
@@ -220,7 +231,11 @@ def inverse(  # type: ignore
                     raise ValueError(f"Unsupported Sequential {module}.")
                 else:
                     raise NotImplementedError(f"data_key {dcate} is not implemented for {module}.")
-            outputs.append(input)
+            if isinstance(arg, (Boxes,)):
+                arg._data = input
+                outputs[idx] = arg.to_tensor()
+            else:
+                outputs[idx] = input
 
         if len(outputs) == 1 and isinstance(outputs, (tuple, list)):
             return outputs[0]
@@ -282,10 +297,12 @@ def forward(  # type: ignore
         Tuple[List[AugmentationSequentialInput], Optional[torch.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 = list(DataKey.get(inp) for inp in data_keys)
+            _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
@@ -1,3 +1,4 @@
+import warnings
 from abc import ABCMeta, abstractmethod
 from functools import partial
 from typing import Callable, Dict, Iterator, List, Optional, Tuple, Type, Union, cast
@@ -6,8 +7,7 @@
 import torch.nn as nn
 
 import kornia  # lazy loading for circular dependencies
-from kornia.augmentation._2d.geometric.base import GeometricAugmentationBase2D
-from kornia.augmentation._2d.mix.base import MixAugmentationBase
+from kornia.augmentation import GeometricAugmentationBase2D, MixAugmentationBase, RandomCrop
 from kornia.augmentation.base import TensorWithTransformMat, _AugmentationBase
 from kornia.augmentation.container.base import ParamItem
 from kornia.constants import DataKey
@@ -269,38 +269,62 @@ class BBoxApplyInverse(ApplyInverseImpl):
     This is for transform boxes in the format (B, N, 4, 2).
     """
 
-    apply_func = partial(transform_bbox, mode="xyxy", restore_coordinates=True)
+    @classmethod
+    def _get_padding_size(cls, module: nn.Module, param: Optional[ParamItem]) -> Optional[torch.Tensor]:
+        if isinstance(module, RandomCrop):
+            _param = cast(Dict[str, torch.Tensor], param.data)  # type: ignore
+            return _param.get("padding_size")
+        return None
 
+    @classmethod
+    def pad(cls, input: torch.Tensor, padding_size: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            input: (B, N, 4, 2)
+            padding_size: (B, 4)
+        """
+        assert len(input.shape) in (3, 4,), input.shape
+        assert len(padding_size.shape) == 2, padding_size.shape
 
-class BBoxXYXYApplyInverse(ApplyInverseImpl):
-    """Apply and inverse transformations for bounding box tensors.
+        _input = input.clone()
 
-    This is for transform boxes in the format [xmin, ymin, xmax, ymax].
-    """
+        if input.dim() == 3:
+            # B,4,2 to B,1,4,2
+            _input = _input[:, None]
 
-    apply_func = partial(transform_bbox, mode="xyxy", restore_coordinates=True)
+        _input[..., 0] += padding_size[..., None, :1]  # left padding
+        _input[..., 1] += padding_size[..., None, 2:3]  # top padding
 
-    @classmethod
-    def _get_padding_size(cls, module, param):
-        if isinstance(module, GeometricAugmentationBase2D):  # TODO: Or only RandomCrop?
-            _param = cast(Dict[str, torch.Tensor], param.data)  # type: ignore
-            return _param.get("padding_size")
-        else:
-            return None
+        if input.dim() == 3:
+            _input = _input[:, 0]  # squeeze back
 
-    @classmethod
-    def pad(cls, input, padding_size):
-        for i in range(len(padding_size)):
-            input[i, :, 0::2] += padding_size[i][0]  # left padding
-            input[i, :, 1::2] += padding_size[i][2]  # top padding
-        return input
+        return _input
 
     @classmethod
-    def unpad(cls, input, padding_size):
-        for i in range(len(padding_size)):
-            input[i, :, 0::2] -= padding_size[i][0]  # left padding
-            input[i, :, 1::2] -= padding_size[i][2]  # top padding
-        return input
+    def unpad(cls, input: torch.Tensor, padding_size: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            input: (B, N, 4, 2)
+            padding_size: (B, 4)
+        """
+        assert len(input.shape) in (3, 4,), input.shape
+        assert len(padding_size.shape) == 2, padding_size.shape
+
+        _input = input.clone()
+
+        if input.dim() == 3:
+            # B,4,2 to B,1,4,2
+            _input = _input[:, None]
+
+        _input[..., 0] -= padding_size[..., None, :1]  # left padding
+        _input[..., 1] -= padding_size[..., None, 2:3]  # top padding
+
+        if input.dim() == 3:
+            _input = _input[:, 0]  # squeeze back
+
+        return _input
+
+    apply_func = partial(transform_bbox, mode="xyxy", restore_coordinates=True)
 
     @classmethod
     def apply_trans(
@@ -309,7 +333,7 @@ def apply_trans(
         """Apply a transformation with respect to the parameters.
 
         Args:
-            input: the input tensor.
+            input: the input tensor, (B, N, 4, 2) or (B, 4, 2).
             label: the optional label tensor.
             module: any torch Module but only kornia augmentation modules will count
                 to apply transformations.
@@ -337,13 +361,50 @@ def inverse(cls, input: torch.Tensor, module: nn.Module, param: Optional[ParamIt
                 to apply transformations.
             param: the corresponding parameters to the module.
         """
-        inverse = super().inverse(input, module, param)
+        _input = input.clone()
+
+        _input = super().inverse(_input, module, param)
 
         padding_size = cls._get_padding_size(module, param)
         if padding_size is not None:
-            inverse = cls.unpad(inverse, padding_size)
+            _input = cls.unpad(_input, padding_size)
 
-        return inverse
+        return _input
+
+
+class BBoxXYXYApplyInverse(BBoxApplyInverse):
+    """Apply and inverse transformations for bounding box tensors.
+
+    This is for transform boxes in the format [xmin, ymin, xmax, ymax].
+    """
+
+    apply_func = partial(transform_bbox, mode="xyxy", restore_coordinates=True)
+
+    @classmethod
+    def pad(cls, input, padding_size):
+        for i in range(len(padding_size)):
+            input[i, :, 0::2] += padding_size[i][0]  # left padding
+            input[i, :, 1::2] += padding_size[i][2]  # top padding
+        return input
+
+    @classmethod
+    def unpad(cls, input, padding_size):
+        for i in range(len(padding_size)):
+            input[i, :, 0::2] -= padding_size[i][0]  # left padding
+            input[i, :, 1::2] -= padding_size[i][2]  # top padding
+        return input
+
+    @classmethod
+    def apply_trans(
+        cls, input: torch.Tensor, label: Optional[torch.Tensor], module: nn.Module, param: ParamItem
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        warnings.warn("BBoxXYXYApplyInverse is no longer maintained. Please use BBoxApplyInverse instead.")
+        return super().apply_trans(input, label=label, module=module, param=param)
+
+    @classmethod
+    def inverse(cls, input: torch.Tensor, module: nn.Module, param: Optional[ParamItem] = None) -> torch.Tensor:
+        warnings.warn("BBoxXYXYApplyInverse is no longer maintained. Please use BBoxApplyInverse instead.")
+        return super().inverse(input, module=module, param=param)
 
 
 class BBoxXYWHApplyInverse(BBoxXYXYApplyInverse):
@@ -371,12 +432,56 @@ def unpad(cls, input, padding_size):
         return input
 
 
-class KeypointsApplyInverse(BBoxXYWHApplyInverse):
-    """Apply and inverse transformations for keypoints tensors."""
+class KeypointsApplyInverse(BBoxApplyInverse):
+    """Apply and inverse transformations for keypoints tensors.
+
+    This is for transform keypoints in the format (B, N, 2).
+    """
 
     # Hot fix for the typing mismatching
     apply_func = partial(transform_points)
 
+    @classmethod
+    def pad(cls, input: torch.Tensor, padding_size: torch.Tensor) -> torch.Tensor:
+
+        assert len(input.shape) in (2, 3,), input.shape
+        assert len(padding_size.shape) == 2, padding_size.shape
+
+        _input = input.clone()
+
+        if input.dim() == 2:
+            # B,2 to B,1,2
+            _input = _input[:, None]
+
+        _input[..., 0] += padding_size[..., :1]  # left padding
+        _input[..., 1] += padding_size[..., 2:3]  # top padding
+
+        if input.dim() == 2:
+            _input = _input[:, 0]  # squeeze back
+
+        return _input
+
+    @classmethod
+    def unpad(cls, input: torch.Tensor, padding_size: torch.Tensor) -> torch.Tensor:
+
+        assert len(input.shape) in (2, 3,), input.shape
+        assert len(padding_size.shape) == 2, padding_size.shape
+
+        _input = input.clone()
+
+        if input.dim() == 2:
+            # B,2 to B,1,2
+            _input = _input[:, None]
+
+        # unpad only xy, not wh
+        _input[..., 0] -= padding_size[..., :1]  # left padding
+        _input[..., 1] -= padding_size[..., 2:3]  # top padding
+
+        if input.dim() == 2:
+            _input = _input[:, 0]  # squeeze back
+
+        return _input
+
 
 class ApplyInverse:
     """Apply and inverse transformations for any tensors (e.g. mask, box, points)."""
@@ -385,7 +490,7 @@ class ApplyInverse:
     def _get_func_by_key(cls, dcate: Union[str, int, DataKey]) -> Type[ApplyInverseInterface]:
         if DataKey.get(dcate) == DataKey.INPUT:
             return InputApplyInverse
-        if DataKey.get(dcate) in [DataKey.MASK]:
+        if DataKey.get(dcate) == DataKey.MASK:
             return MaskApplyInverse
         if DataKey.get(dcate) in [DataKey.BBOX, DataKey.BBOX_XYXY, DataKey.BBOX_XYWH]:
             # We are converting to (B, 4, 2) internally for all formats.
diff --git a/kornia/geometry/boxes.py b/kornia/geometry/boxes.py
--- a/kornia/geometry/boxes.py
+++ b/kornia/geometry/boxes.py
@@ -207,6 +207,8 @@ def __init__(
         if not (3 <= boxes.ndim <= 4 and boxes.shape[-2:] == (4, 2)):
             raise ValueError(f"Boxes shape must be (N, 4, 2) or (B, N, 4, 2). Got {boxes.shape}.")
 
+        self._is_batched = False if boxes.ndim == 3 else True
+
         self._data = boxes
         self._mode = mode
 
@@ -432,10 +434,6 @@ def transform_boxes_(self, M: torch.Tensor) -> "Boxes":
         """Inplace version of :func:`Boxes.transform_boxes`"""
         return self.transform_boxes(M, inplace=True)
 
-    @property
-    def _is_batched(self) -> bool:
-        return self._data.ndim == 4
-
     @property
     def data(self) -> torch.Tensor:
         return self._data
@@ -479,7 +477,10 @@ class Boxes3D:
         ``width = xmax - xmin + 1``, ``height = ymax - ymin + 1`` and ``depth = zmax - zmin + 1``. Examples of
         `hexahedrons <https://en.wikipedia.org/wiki/Hexahedron>`_ are cubes and rhombohedrons.
     """
-    def __init__(self, boxes: torch.Tensor, raise_if_not_floating_point: bool = True):
+    def __init__(
+        self, boxes: torch.Tensor, raise_if_not_floating_point: bool = True,
+        mode: str = "xyzxyz_plus"
+    ) -> None:
         if not isinstance(boxes, torch.Tensor):
             raise TypeError(f"Input boxes is not a Tensor. Got: {type(boxes)}.")
 
@@ -497,7 +498,10 @@ def __init__(self, boxes: torch.Tensor, raise_if_not_floating_point: bool = True
         if not (3 <= boxes.ndim <= 4 and boxes.shape[-2:] == (8, 3)):
             raise ValueError(f"3D bbox shape must be (N, 8, 3) or (B, N, 8, 3). Got {boxes.shape}.")
 
+        self._is_batched = False if boxes.ndim == 3 else True
+
         self._data = boxes
+        self._mode = mode
 
     def get_boxes_shape(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         r"""Compute boxes heights and widths.
@@ -594,7 +598,7 @@ def from_tensor(cls, boxes: torch.Tensor, mode: str = "xyzxyz", validate_boxes:
         hexahedrons = hexahedrons if batched else hexahedrons.squeeze(0)
         # Due to some torch.jit.script bug (at least <= 1.9), you need to pass all arguments to __init__ when
         # constructing the class from inside of a method.
-        return cls(hexahedrons, raise_if_not_floating_point=False)
+        return cls(hexahedrons, raise_if_not_floating_point=False, mode=mode)
 
     def to_tensor(self, mode: str = "xyzxyz") -> torch.Tensor:
         r"""Cast :class:`Boxes3D` to a tensor. ``mode`` controls which 3D boxes format should be use to represent boxes
@@ -776,20 +780,20 @@ def transform_boxes(self, M: torch.Tensor, inplace: bool = False) -> "Boxes3D":
             self._data = transformed_boxes
             return self
 
-        return Boxes3D(transformed_boxes, False)
+        return Boxes3D(transformed_boxes, False, "xyzxyz_plus")
 
     def transform_boxes_(self, M: torch.Tensor) -> "Boxes3D":
         """Inplace version of :func:`Boxes3D.transform_boxes`"""
         return self.transform_boxes(M, inplace=True)
 
-    @property
-    def _is_batched(self) -> bool:
-        return self._data.ndim == 4
-
     @property
     def data(self) -> torch.Tensor:
         return self._data
 
+    @property
+    def mode(self) -> str:
+        return self._mode
+
     @property
     def device(self) -> torch.device:
         """Returns boxes device."""


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
@@ -377,7 +377,7 @@ def test_forward_and_inverse(self, random_apply, return_transform, device, dtype
             return_transform=return_transform,
         )
         out = aug(inp, mask, bbox, keypoints)
-        if return_transform and isinstance(out, (tuple, list)):
+        if return_transform and isinstance(out[0], (tuple, list)):
             assert out[0][0].shape == inp.shape
         else:
             assert out[0].shape == inp.shape
@@ -387,7 +387,10 @@ def test_forward_and_inverse(self, random_apply, return_transform, device, dtype
         reproducibility_test((inp, mask, bbox, keypoints), aug)
 
         out_inv = aug.inverse(*out)
-        assert out_inv[0].shape == inp.shape
+        if return_transform and isinstance(out_inv[0], (tuple, list)):
+            assert out_inv[0][0].shape == inp.shape
+        else:
+            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