Repository: kornia/kornia

Issue:
AugmentationSequential does not return lists of boxes
### Describe the bug

When passing boxes as lists of tensors with any BBOX datakey to `AugmentationSequential` the resulting augmentation is not returned. Thus, the returned outputs have one fewer element than the inputs that were passed to it. This happens for both forward and inverse.

### Reproduction steps

```bash
import kornia.augmentation as K
import torch

img = torch.rand((3,3,10,10))
bbox = [torch.tensor([[1,5,2,7],[0,3,9,9]]),torch.tensor([[1,5,2,7],[0,3,9,9],[0,5,8,7]]), torch.empty((0,4))]
inputs = [img, bbox]

aug = K.AugmentationSequential(K.Resize((300,300)), data_keys=['input', 'bbox_xyxy'])

transformed = aug(*inputs)

assert len(transformed) == len(inputs)
```
```
Traceback (most recent call last):             
  File "/home/miquelmr/kornia/test.py", line 12, in <module>
    assert len(transformed) == len(inputs)
AssertionError  
```


### Expected behavior

List of boxes is returned.

### 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:
@miquelmarti @shijianjian do you believe is it an easy fix ? Besides, i would like start iterating on how to integrate the new apis from geometry like the [bounding box](https://github.com/kornia/kornia/blob/master/kornia/geometry/boxes.py#L167) class. I'll soon push a new api i've been iterating the background for images too.

In particular, i believe the bbox class should be adjusted a bit to support autocast and stuff like this
```python

class BoundingBoxFormat(Enum):
   XYXY = 0
   ....

# TODO: discuss about if bbox makes sense to be a tensor, eg.g to optimize/refine coords
class BoundingBox(TensorWrapper):
  def __init__(self, data: Union[Tensor, List[Tensor]], data_format: BoundingBoxFormat) -> None:
     self._data = data
     self._data_format = data_format

  # adapt the rest: https://github.com/kornia/kornia/blob/master/kornia/geometry/boxes.py#L167

  def convert(self, new_format: BoundingBoxFormat) -> "BoundingBox":
   # check data_format and apply the corresponding logic to convert from one format to the other
    ...

  def normalized(self) -> "BoundingBox":
     # return [-1,1] e.g to use directly in grid_sample
     ...
```
For the fix, I tried to remove this line where the list of boxes is removed (since not a tensor):
https://github.com/kornia/kornia/blob/370099180c97167d066d94491eda46e7b9c8b939/kornia/augmentation/container/augment.py#L279
But it triggered multiple test errors. I don't understand what's the purpose of this or what the thought was. Was it to handle cases where not all keys are passed and thus there are Nones in the outputs list? If someone can guide me with this I will fix it.
@shijianjian can you guide here?
IIRC, we support the list of tensors in the following lines:

https://github.com/kornia/kornia/blob/370099180c97167d066d94491eda46e7b9c8b939/kornia/augmentation/container/augment.py#L230-L233

for which it was designed to use the `Boxes` wrapper to wrap those tensors, while it might be more intuitive to use the list of tensors directly. You may test and add such supports. @miquelmarti 
It seems that the Boxes wrapper works just fine when passing lists of tensors up to when it is converted back to a tensor or a list of tensors here, depending on what the input was: https://github.com/kornia/kornia/blob/370099180c97167d066d94491eda46e7b9c8b939/kornia/augmentation/container/augment.py#L275

This seems to work fine too but now `outputs` contains a list of tensors (if that was the format passed to the input), which gets removed from the list of outputs in L279. It is this last removal that I don't understand the purpose of and the implications that it would have to just not do that, as it would seem to be required to support returning lists of tensors.
> It seems that the Boxes wrapper works just fine when passing lists of tensors up to when it is converted back to a tensor or a list of tensors here, depending on what the input was:
> 
> https://github.com/kornia/kornia/blob/370099180c97167d066d94491eda46e7b9c8b939/kornia/augmentation/container/augment.py#L275
> 
> This seems to work fine too but now `outputs` contains a list of tensors (if that was the format passed to the input), which gets removed from the list of outputs in L279. It is this last removal that I don't understand the purpose of and the implications that it would have to just not do that, as it would seem to be required to support returning lists of tensors.

I agree here. I think we shall not use any obvious wrapper at all. Say, the user input a list of tensors directly and get the corresponding list of tensors. If they input a wrapped `Box` object, they should get the corresponding `Box` object.

Additionally, we should use our own `Box` data type as the default communication type to pass around different methods. 
Is this somehow been fixed? I cannot reproduce it.
> Is this somehow been fixed? I cannot reproduce it.

Just tried on master, still the same.
It worked on previous versions before #2030 I believe.

> For the fix, I tried to remove this line where the list of boxes is removed (since not a tensor):
> 
> https://github.com/kornia/kornia/blob/370099180c97167d066d94491eda46e7b9c8b939/kornia/augmentation/container/augment.py#L279
> 
> 
> But it triggered multiple test errors. I don't understand what's the purpose of this or what the thought was. Was it to handle cases where not all keys are passed and thus there are Nones in the outputs list? If someone can guide me with this I will fix it.

https://github.com/kornia/kornia/blob/370099180c97167d066d94491eda46e7b9c8b939/kornia/augmentation/container/augment.py#L428

@johnnv1 Is there any reasons that we need these lines?
just for typing, as it didn't fail any tests I thought there was no problem, should be a easy fix and need to fix the typing in some places -- I thought this was a problem when I was doing #2094 (https://github.com/kornia/kornia/blob/9514ea767e0f0ce2eebcd369c99f123236bae1a4/kornia/augmentation/container/augment.py#L430)



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
@@ -18,6 +18,7 @@
 from kornia.constants import DataKey, Resample
 from kornia.core import Tensor
 from kornia.geometry.boxes import Boxes
+from kornia.testing import KORNIA_CHECK_IS_LIST_OF_TENSOR
 from kornia.utils import eye_like
 
 __all__ = ["AugmentationSequential"]
@@ -193,7 +194,7 @@ def inverse(  # type: ignore[override]
         *args: Tensor,
         params: Optional[List[ParamItem]] = None,
         data_keys: Optional[List[Union[str, int, DataKey]]] = None,
-    ) -> Union[Tensor, List[Tensor]]:
+    ) -> Union[Tensor, List[Tensor], List[Union[Tensor, List[Tensor]]]]:
         """Reverse the transformation applied.
 
         Number of input tensors must align with the number of``data_keys``. If ``data_keys`` is not set, use
@@ -276,7 +277,7 @@ def inverse(  # type: ignore[override]
             else:
                 outputs[idx] = input
 
-        _outputs = [i for i in outputs if isinstance(i, Tensor)]
+        _outputs = [i for i in outputs if isinstance(i, Tensor) or KORNIA_CHECK_IS_LIST_OF_TENSOR(i)]
 
         if len(_outputs) == 1 and isinstance(_outputs, list):
             return _outputs[0]
@@ -284,10 +285,15 @@ def inverse(  # type: ignore[override]
         return _outputs
 
     def __packup_output__(  # type: ignore[override]
-        self, output: List[Tensor], label: Optional[Tensor] = None
-    ) -> Union[Tensor, List[Tensor], Tuple[Union[Tensor, List[Tensor]], Optional[Tensor]]]:
+        self, output: List[Union[Tensor, List[Tensor]]], label: Optional[Tensor] = None
+    ) -> Union[
+        Tensor,
+        List[Tensor],
+        List[Union[Tensor, List[Tensor]]],
+        Tuple[Union[Tensor, List[Tensor], List[Union[Tensor, List[Tensor]]]], Optional[Tensor]],
+    ]:
 
-        _out: Union[Tensor, List[Tensor]]
+        _out: Union[Tensor, List[Tensor], List[Union[Tensor, List[Tensor]]]]
 
         if len(output) == 1 and isinstance(output, list):
             _out = output[0]
@@ -331,7 +337,12 @@ def forward(  # type: ignore[override]
         label: Optional[Tensor] = None,
         params: Optional[List[ParamItem]] = None,
         data_keys: Optional[List[Union[str, int, DataKey]]] = None,
-    ) -> Union[Tensor, List[Tensor], Tuple[Union[Tensor, List[Tensor]], Optional[Tensor]]]:
+    ) -> Union[
+        Tensor,
+        List[Tensor],
+        List[Union[Tensor, List[Tensor]]],
+        Tuple[Union[Tensor, List[Tensor], List[Union[Tensor, List[Tensor]]]], Optional[Tensor]],
+    ]:
         """Compute multiple tensors simultaneously according to ``self.data_keys``."""
         if data_keys is None:
             _data_keys = self.data_keys
@@ -425,5 +436,6 @@ def forward(  # type: ignore[override]
                 outputs[idx] = arg.to_tensor()
             else:
                 outputs[idx] = input
-        _outputs = [i for i in outputs if isinstance(i, Tensor)]
+        _outputs = [i for i in outputs if isinstance(i, Tensor) or KORNIA_CHECK_IS_LIST_OF_TENSOR(i)]
+
         return self.__packup_output__(_outputs, label)


Test patch:
diff --git a/kornia/testing/__init__.py b/kornia/testing/__init__.py
--- a/kornia/testing/__init__.py
+++ b/kornia/testing/__init__.py
@@ -4,10 +4,11 @@
 from abc import ABC, abstractmethod
 from copy import deepcopy
 from itertools import product
-from typing import Any, Dict, List, Optional, Tuple, TypeVar, cast
+from typing import Any, Dict, List, Optional, Sequence, Tuple, TypeVar, cast
 
 import torch
 from torch.testing import assert_close as _assert_close
+from typing_extensions import TypeGuard
 
 from kornia.core import Tensor, eye, tensor
 
@@ -21,6 +22,7 @@
     "KORNIA_UNWRAP",
     "KORNIA_CHECK_TYPE",
     "KORNIA_CHECK_IS_TENSOR",
+    "KORNIA_CHECK_IS_LIST_OF_TENSOR",
     "KORNIA_CHECK_SAME_DEVICE",
     "KORNIA_CHECK_SAME_DEVICES",
     "KORNIA_CHECK_IS_COLOR",
@@ -396,6 +398,25 @@ def KORNIA_CHECK_IS_TENSOR(x, msg: Optional[str] = None):
         raise TypeError(f"Not a Tensor type. Got: {type(x)}.\n{msg}")
 
 
+def KORNIA_CHECK_IS_LIST_OF_TENSOR(x: Optional[Sequence[object]]) -> TypeGuard[List[Tensor]]:
+    """Check the input variable is a List of Tensors.
+
+    Args:
+        x: Any sequence of objects
+
+    Return:
+        True if the input is a list of Tensors, otherwise return False.
+
+    Example:
+        >>> x = torch.rand(2, 3, 3)
+        >>> KORNIA_CHECK_IS_LIST_OF_TENSOR(x)
+        False
+        >>> KORNIA_CHECK_IS_LIST_OF_TENSOR([x])
+        True
+    """
+    return isinstance(x, list) and all(isinstance(d, Tensor) for d in x)
+
+
 def KORNIA_CHECK_SAME_DEVICE(x: Tensor, y: Tensor):
     """Check whether two tensor in the same device.
 
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
@@ -426,6 +426,27 @@ def test_random_crops(self, device, dtype):
         assert out_inv[3].shape == points.shape
         assert_close(out_inv[3], points, atol=1e-4, rtol=1e-4)
 
+    def test_bbox(self, device, dtype):
+        img = torch.rand((3, 3, 10, 10), device=device, dtype=dtype)
+        bbox = [
+            torch.tensor([[1, 5, 2, 7], [0, 3, 9, 9]], device=device, dtype=dtype),
+            torch.tensor([[1, 5, 2, 7], [0, 3, 9, 9], [0, 5, 8, 7]], device=device, dtype=dtype),
+            torch.empty((0, 4), device=device, dtype=dtype),
+        ]
+
+        inputs = [img, bbox]
+
+        aug = K.AugmentationSequential(K.Resize((30, 30)), data_keys=['input', 'bbox_xyxy'])
+
+        transformed = aug(*inputs)
+
+        assert len(transformed) == len(inputs)
+        bboxes_transformed = transformed[-1]
+        assert len(bboxes_transformed) == len(bbox)
+        assert len(bboxes_transformed[0]) == 2
+        assert len(bboxes_transformed[1]) == 3
+        assert len(bboxes_transformed[2]) == 0
+
     @pytest.mark.parametrize('random_apply', [1, (2, 2), (1, 2), (2,), 10, True, False])
     def test_forward_and_inverse(self, random_apply, device, dtype):
         inp = torch.randn(1, 3, 1000, 500, device=device, dtype=dtype)