Repository: kornia/kornia

Issue:
`AugmentationSequential` does not support instance masks shape (N, H, W)
### Describe the bug

Currently `AugmentationSequential` does not seem to support instance masks i.e. an image (C, H, W) having multiple masks (N, H, W). 



### Reproduction steps


The following code snippet errors out with a `RuntimeError: Expect (B, C, H, W). Got torch.Size([3, 256, 256]).` error.

```python
import torch
import kornia.augmentation as K

img = torch.randn(3, 256, 256)
masks = torch.ones( 2, 256, 256)

aug_dict = K.AugmentationSequential(
    K.RandomHorizontalFlip(p=1.0),
    data_keys=None,
    keepdim=True,
)
aug_dict.keepdim = True

data = {'image': img, 'mask': masks}
out = aug_dict(data)
```



### Expected behavior

Error should not be raised

### Environment

```shell
wget https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py

kornia==0.7.2
```
- PyTorch Version (e.g., 1.0):
- OS (e.g., Linux):
- How you installed PyTorch (`conda`, `pip`, source):
- Build command you used (if compiling from source):
- Python version:
- CUDA/cuDNN version:
- GPU models and configuration:
- Any other relevant information:




Discussion:

Resolving patch:
diff --git a/kornia/augmentation/container/ops.py b/kornia/augmentation/container/ops.py
--- a/kornia/augmentation/container/ops.py
+++ b/kornia/augmentation/container/ops.py
@@ -1,3 +1,4 @@
+import copy
 from abc import ABCMeta, abstractmethod
 from typing import Any, Callable, Dict, Generic, List, Optional, Type, TypeVar, Union
 
@@ -286,10 +287,11 @@ def transform_list(
         if isinstance(module, (K.GeometricAugmentationBase2D,)):
             tfm_input = []
             params = cls.get_instance_module_param(param)
-            for inp in input:
-                inp = module.transform_tensor(inp)
+            params_i = copy.deepcopy(params)
+            for i, inp in enumerate(input):
+                params_i["batch_prob"] = params["batch_prob"][i]
                 tfm_inp = module.transform_masks(
-                    inp, params=params, flags=module.flags, transform=module.transform_matrix, **extra_args
+                    inp, params=params_i, flags=module.flags, transform=module.transform_matrix, **extra_args
                 )
                 tfm_input.append(tfm_inp)
             input = tfm_input
@@ -302,9 +304,10 @@ def transform_list(
         elif isinstance(module, (_AugmentationBase)):
             tfm_input = []
             params = cls.get_instance_module_param(param)
-            for inp in input:
-                inp = module.transform_tensor(inp)
-                tfm_inp = module.transform_masks(inp, params=params, flags=module.flags, **extra_args)
+            params_i = copy.deepcopy(params)
+            for i, inp in enumerate(input):
+                params_i["batch_prob"] = params["batch_prob"][i]
+                tfm_inp = module.transform_masks(inp, params=params_i, flags=module.flags, **extra_args)
                 tfm_input.append(tfm_inp)
             input = tfm_input
 
@@ -314,7 +317,7 @@ def transform_list(
             for inp in input:
                 tfm_inp = module.transform_masks(inp, params=seq_params, extra_args=extra_args)
                 tfm_input.append(tfm_inp)
-                input = tfm_input
+            input = tfm_input
 
         elif isinstance(module, K.container.ImageSequentialBase):
             tfm_input = []
@@ -329,9 +332,6 @@ def transform_list(
                 "The support for list of masks under auto operations are not yet supported. You are welcome to file a"
                 " PR in our repo."
             )
-            # return [
-            # MaskSequentialOps.transform(inp, module=module.op, param=param, extra_args=extra_args) for inp in input
-            # ]
         return input
 
     @classmethod
diff --git a/kornia/augmentation/utils/helpers.py b/kornia/augmentation/utils/helpers.py
--- a/kornia/augmentation/utils/helpers.py
+++ b/kornia/augmentation/utils/helpers.py
@@ -96,10 +96,10 @@ def _transform_input_by_shape(input: Tensor, reference_shape: Tensor, match_chan
     if len(input.shape) == 2:
         input = input.unsqueeze(0)
 
-    if len(input.shape) == 3 and B == input.shape[-3]:
+    if len(input.shape) == 3:
         # If the first dim matches within the batch_size, add a `C` dim
         # Useful to handler Masks without `C` dimensions
-        input = input.unsqueeze(1)
+        input = input.unsqueeze(1) if B == input.shape[-3] else input.unsqueeze(0)
 
     if match_channel and C:
         if not input.shape[-3] == C:
@@ -211,7 +211,7 @@ def _transform_output_shape(
 
     for dim in range(len(out_tensor.shape) - len(shape)):
         idx = 0
-        if reference_shape is not None and out_tensor.shape[0] == reference_shape[0] and len(shape) > 2:
+        if reference_shape is not None and out_tensor.shape[0] == reference_shape[0] != 1 and len(shape) > 2:
             idx = 1
         if out_tensor.shape[idx] != 1:
             raise AssertionError(f"Dimension {dim} of input is expected to be 1, got {out_tensor.shape[idx]}")


Test patch:
diff --git a/tests/augmentation/test_container.py b/tests/augmentation/test_container.py
--- a/tests/augmentation/test_container.py
+++ b/tests/augmentation/test_container.py
@@ -524,7 +524,7 @@ def test_class(self, device, dtype):
 
     @pytest.mark.parametrize("B", [None, 1, 3, 5])
     @pytest.mark.parametrize("C_i", [1, 3])
-    @pytest.mark.parametrize("C_m", [None, 1])
+    @pytest.mark.parametrize("C_m", [None, 1, 3, 4])
     @pytest.mark.parametrize("keepdim", [True, False])
     def test_masks_without_channel_dim(self, device, dtype, B, C_i, C_m, keepdim):
         img_shape = (B, C_i, 1000, 500) if B else (C_i, 1000, 500)