Repository: kornia/kornia

Issue:
ValueError when applying flip augmentations to boxes
### Describe the bug

Came across a strange bug while applying flip augmentations with bboxes. Running the below code snippet produces the following error some of the time. 

For reference my image is of shape `(B, C, H, W)` and my boxes are a list of length `B` where each element is of shape `(N,4)`

```
ValueError: Input batch size must be the same for both tensors or 1.Got torch.Size([2, 3, 3]) and torch.Size([1, 8, 2])
```

The randomness of the message could be stemming from the randomness in applying the augmentation but setting `p=1.0` or `p=0.0` leads to no error being reported.

### Reproduction steps

```bash
import kornia.augmentation as K
import torch

augs = K.AugmentationSequential(
    K.RandomHorizontalFlip(p=0.5),
    K.RandomVerticalFlip(p=0.5),
    data_keys=["input", "bbox_xyxy"],
)

image = torch.randn((2,3,200,200))
boxes = [
    torch.tensor([
        [1, 2, 2, 4],
        [2, 4, 3, 6],
    ]),
    torch.tensor([
        [1, 2, 2, 4],
    ]),
]

tfm_img, tfm_bbox = augs(image, boxes)
```


### Expected behavior

Script should run successfully all the time

### Environment

```shell
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.5
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.23.2
Libc version: glibc-2.31

Python version: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:56:21)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-117-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.3.58
GPU models and configuration: GPU 0: NVIDIA TITAN Xp
Nvidia driver version: 470.129.06
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] efficientnet-pytorch==0.6.3
[pip3] mypy==0.961
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.4
[pip3] pytorch-lightning==1.6.4
[pip3] pytorch-sphinx-theme==0.0.24
[pip3] segmentation-models-pytorch==0.2.1
[pip3] torch==1.11.0
[pip3] torchmetrics==0.9.0
[pip3] torchvision==0.12.0
[conda] blas                      2.115                       mkl    conda-forge
[conda] blas-devel                3.9.0            15_linux64_mkl    conda-forge
[conda] cudatoolkit               11.5.1              h59c8dcf_10    conda-forge
[conda] efficientnet-pytorch      0.6.3                    pypi_0    pypi
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libblas                   3.9.0            15_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            15_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            15_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            15_linux64_mkl    conda-forge
[conda] mkl                       2022.1.0           h84fe81f_915    conda-forge
[conda] mkl-devel                 2022.1.0           ha770c72_916    conda-forge
[conda] mkl-include               2022.1.0           h84fe81f_915    conda-forge
[conda] numpy                     1.22.4           py39hc58783e_0    conda-forge
[conda] pytorch                   1.11.0          py3.9_cuda11.5_cudnn8.3.2_0    pytorch
[conda] pytorch-lightning         1.6.4                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch-sphinx-theme      0.0.24                   pypi_0    pypi
[conda] segmentation-models-pytorch 0.2.1                    pypi_0    pypi
[conda] torchmetrics              0.9.0                    pypi_0    pypi
[conda] torchvision               0.12.0               py39_cu115    pytorch
```



Discussion:
Tried to narrow down the error and it seems it comes down to the shape of `boxes` that are passed on to `transform_bbox`. Using the above snippet
- If the shape is `(2,2,4,2)`, the script passes 
- If the shape is `(1,2,4,2)`, the script fails
@ashnair1 for now we assume dense tensors. to make it work correctly you should pad your boxes tensors to match with the maximum size (here 2). Ideally, we should implement this as part of the `Boxes` class 
https://github.com/kornia/kornia/blob/master/kornia/geometry/boxes.py#L171

we are open to contributions -- maybe for the next iteration based on https://github.com/kornia/kornia/issues/1754

/cc @shijianjian thoughts ?
@edgarriba The error occurs even if I ensure the number of boxes to be equal. For example, changing the boxes tensor as shown below

```
boxes = [
    torch.tensor([
        [1, 2, 2, 4],
        [2, 4, 3, 6],
    ]),
    torch.tensor([
        [1, 2, 2, 4],
        [2, 4, 3, 6],
    ]),
]
```

still throws

```
ValueError: Input batch size must be the same for both tensors or 1.Got torch.Size([2, 3, 3]) and torch.Size([1, 8, 2])
```




Ok, I've found the cause. This error occurs when elements in your batch have **different augmentations**. So in my above example, if I had batch size of 2 where one element was flipped and the other wasn't this error is triggered.

The 'fix' is to set the `same_on_batch` arg in `AugmentationSequential` to `True`. Perhaps this should be documented somewhere.


Technically this is still a problem and applying distinct augmentations on batch elements should be possible.

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
@@ -100,7 +100,7 @@ def apply_trans(
 
         # If any inputs need to be transformed.
         if mat is not None and to_apply is not None and to_apply.sum() != 0 and input.numel() > 0:
-            input[to_apply] = cls.apply_func(mat, input[to_apply])
+            input[to_apply] = cls.apply_func(mat[to_apply], input[to_apply])
 
         return input, label
 


Test patch: