Repository: kornia/kornia

Issue:
AugmentationSequential: accept sample dict as input
## 🚀 Feature

I would like `AugmentationSequential` to support a dictionary as input.

## Motivation

I'm a [TorchGeo](https://github.com/microsoft/torchgeo) developer. In TorchGeo, every dataset returns a sample dictionary like so:
```python
sample = {
    "input": torch.tensor(...),
    "mask": torch.tensor(...),
    "bbox": torch.tensor(...),
    ...
}
```
(the exact key names don't match at the moment, but we're working on standardizing those)

## Pitch

With the feature I'm envisioning, the following would work:
```python
augs = AugmentationSequential(...)
sample = augs(sample)
```

The exact implementation details would still need to be worked out, but `*args` would go from `Tensor` to `Union[Tensor, Dict[str, Tensor]]`. The dictionary may contain keys that Kornia doesn't know how to support, and these should be ignored. If a sample dictionary contains a known key that the user doesn't want to transform, they can simply pass `data_keys` to override the default detection. If the input is a dict, the output should also be a dict. If implemented correctly, this feature will be backwards compatible with the old behavior so people can still pass these inputs in manually if they want to.

## Alternatives

At the moment, to use Kornia augmentations, we have to use:
```python
augs = AugmentationSequential(..., data_keys=["input", "mask", "bbox", ...])
sample["input"], sample["mask"], sample["bbox"], ... = augs(sample["input"], sample["mask"], sample["bbox"], ...)
```
As you can see, this is much more verbose than necessary. There's no reason we need to duplicate the list of keys so many times.

## Additional context

If this is something you would be interested in, I would be happy to submit a PR to support this. Just wanted to gauge interest first before working on it.

@isaaccorley 


Discussion:
It looks acceptable if you have decent solution, like using `override`. If we accept dict, I would like to support something like:

```python
augs = AugmentationSequential(..., data_keys=NONE)  # so that we do not require to input datakeys in the first place
sample = augs(sample...)
```

Another question is that is we have multiple masks/bboxes for the input, how should we organize the input dict?
> It looks acceptable if you have decent solution, like using override.

Do you mean "overload" instead of "override"? Python doesn't support overloading functions, but you can overload their type hints. It should be possible for the function to be backwards compatible and still support dicts.

> Another question is that is we have multiple masks/bboxes for the input, how should we organize the input dict?

I would not bother supporting this for dict input. If a user needs multiple input/mask/bbox, they can simply use the old syntax.
We can proceed the `dict` support with `overload` as our first attempt. But preferably, I would like to see something like `a dict with duplicated keys`, or `the same key with different values`. XD

It looks like the latter is much doable, but I do not want to use list values (since a user can pass a list of tensors).
looking at the fact that we possibly wouldn't have many duplicate values, I don't think it would be a performance problem to match any case with, for example: mask-1, mask-2, maskAnything, ... -- match any `mask*` case 
The match any strategy looks reasonable to me. Then a user can mark their inputs as `mask-taskA`, `mask-taskB`, etc. We can be more flexible comparing to other libs.

I will try finish the first refactor #2117 this week. Then start implementing this feature on top of the new base. What do you think? @adamjstewart 
i also envision something like this in the midterm

```python
aug = AugmentationPipeline(...)
gen = DataGenerator(...)
net = MyModel(...)
loss = MyGeometricLoss(...)

gen.output_dict >> aug.input
(gen.image | gen.mask) >> aug.input  # or this to be more selective
aug.output >> net.image
(net.output | aug.tranform_matrix) >>  loss.input
```
this is kinda the strategy behind https://github.com/kornia/limbus
which is now under heavy refactor to fully support `asyncio`

@adamjstewart idk if something like that would be interesting for you guys
/cc @lferraz
I tried to take a stab at implementing this but supporting `data_keys = None` and supporting dicts as input/output of inverse/forward is actually quite a large refactor. I can try to force it in there, but it might be better for the person who originally designed it to redesign it with dicts in mind. Do any core developers have any interest in this, or should I try to redesign it myself and minimize any unrequired API changes?
> I tried to take a stab at implementing this but supporting `data_keys = None` and supporting dicts as input/output of inverse/forward is actually quite a large refactor. I can try to force it in there, but it might be better for the person who originally designed it to redesign it with dicts in mind. Do any core developers have any interest in this, or should I try to redesign it myself and minimize any unrequired API changes?

This person is @shijianjian , we just did a huge refactor to the augmentations module, so makes sense to keep improving to support more features 
I think we may totally ignore the `data_keys` argument here, just to overwrite the runtime `data_keys`.

```python
augs = AugmentationSequential(...)
sample = augs({"image-a": imagea, "image-b": imageb})
```

In the implementation, the forward signature changed to:
```python
def forward(self, *input: Union[PREVIOUS_TYPE, Dict[str, Tensor]], data_keys: Optional=None) -> Union[PREVIOUS_TYPE, Dict[str, Tensor]]:
      if len(input) == 1 and isinstance(input, dict):
           if data_keys is not None: raise Error
           data_keys = read_datakeys_from_dict(input[0])

      input = self._preproc_dict(...)
      RUN_ITERATIONS_HERE
      output = self._postproc_dict(...)

     return output
```

This shall be straight-forward to implement.
@shijianjian do you want to submit a PR for this or should I try to hack on it? I have something with 10x as many lines of code changed and it still doesn't work because I'm mapping dicts to lists but I'm not yet mapping them back. It's not just data_keys that needs to change, it's also transform_op (which uses data_keys) that will need to be set dynamically. I can open a draft PR if you want to see my current solution but it's pretty ugly.
@adamjstewart I may not be working on this since it is not a critical feature. The only thing you need to implement is to map the data keys. You do not need to handle the `transform_op`, it shall read the input `data_keys` and perform the augmentations automatically.

> ```python
> def forward(self, *input: Union[PREVIOUS_TYPE, Dict[str, Tensor]], data_keys: Optional=None) -> Union[PREVIOUS_TYPE, Dict[str, Tensor]]:
>       if len(input) == 1 and isinstance(input, dict):
>            if data_keys is not None: raise Error
>            data_keys = read_datakeys_from_dict(input[0])
> 
>       input = self._preproc_dict(...)
>       RUN_ITERATIONS_HERE
>       output = self._postproc_dict(...)
> 
>      return output
> ```

As shown here, having the implementation of `read_datakeys_from_dict`, `_preproc_dict`, `_postproc_dict` shall be enough to make it work.


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
@@ -1,5 +1,5 @@
 import warnings
-from typing import Any, Dict, List, Optional, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
 
 from kornia.augmentation._2d.base import RigidAffineAugmentationBase2D
 from kornia.augmentation._3d.base import AugmentationBase3D, RigidAffineAugmentationBase3D
@@ -168,12 +168,37 @@ class AugmentationSequential(TransformMatrixMinIn, ImageSequential):
         ...    random_apply=10,
         ... )
         >>> out = aug_list(input, mask, bbox)
+
+    How to use a dictionary as input with AugmentationSequential? The dictionary should starts with
+    one of the datakey availables.
+
+        >>> import kornia.augmentation as K
+        >>> img = torch.randn(1, 3, 256, 256)
+        >>> mask = [torch.ones(1, 3, 256, 256), torch.ones(1, 2, 256, 256)]
+        >>> bbox = [
+        ...    torch.tensor([[28.0, 53.0, 143.0, 164.0], [254.0, 158.0, 364.0, 290.0], [307.0, 204.0, 413.0, 350.0]]),
+        ...    torch.tensor([[254.0, 158.0, 364.0, 290.0], [307.0, 204.0, 413.0, 350.0]])
+        ... ]
+        >>> bbox = [Boxes.from_tensor(i).data for i in bbox]
+        >>> aug_dict = K.AugmentationSequential(
+        ...    K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0),
+        ...    K.RandomHorizontalFlip(p=1.0),
+        ...    K.ImageSequential(K.RandomHorizontalFlip(p=1.0)),
+        ...    K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0)),
+        ...    data_keys=None,
+        ...    same_on_batch=False,
+        ...    random_apply=10,
+        ... )
+        >>> data = {'image': img, 'mask': mask[0], 'mask-b': mask[1], 'bbox': bbox[0], 'bbox-other':bbox[1]}
+        >>> out = aug_dict(data)
+        >>> out.keys()
+        dict_keys(['image', 'mask', 'mask-b', 'bbox', 'bbox-other'])
     """
 
     def __init__(
         self,
         *args: Union[_AugmentationBase, ImageSequential],
-        data_keys: Union[List[str], List[int], List[DataKey]] = [DataKey.INPUT],
+        data_keys: Optional[Union[List[str], List[int], List[DataKey]]] = [DataKey.INPUT],
         same_on_batch: Optional[bool] = None,
         keepdim: Optional[bool] = None,
         random_apply: Union[int, bool, Tuple[int, int]] = False,
@@ -202,13 +227,18 @@ def __init__(
             AugmentationSequential,
         )
 
-        self.data_keys = [DataKey.get(inp) for inp in data_keys]
+        self.data_keys: Optional[List[DataKey]]
+        if data_keys is not None:
+            self.data_keys = [DataKey.get(inp) for inp in data_keys]
+        else:
+            self.data_keys = 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 {self.data_keys}.")
+        if self.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 {self.data_keys}.")
 
-        if self.data_keys[0] != DataKey.INPUT:
-            raise NotImplementedError(f"The first input must be {DataKey.INPUT}.")
+            if self.data_keys[0] != DataKey.INPUT:
+                raise NotImplementedError(f"The first input must be {DataKey.INPUT}.")
 
         self.transform_op = AugmentationSequentialOps(self.data_keys)
 
@@ -241,20 +271,28 @@ def identity_matrix(self, input: Tensor) -> Tensor:
 
     def inverse(  # type: ignore[override]
         self,
-        *args: DataType,
+        *args: Union[DataType, Dict[str, DataType]],
         params: Optional[List[ParamItem]] = None,
         data_keys: Optional[Union[List[str], List[int], List[DataKey]]] = None,
-    ) -> Union[DataType, List[DataType]]:
+    ) -> Union[DataType, List[DataType], Dict[str, DataType]]:
         """Reverse the transformation applied.
 
         Number of input tensors must align with the number of``data_keys``. If ``data_keys`` is not set, use
         ``self.data_keys`` by default.
         """
+        original_keys = None
+        if len(args) == 1 and isinstance(args[0], dict):
+            original_keys, data_keys, args = self._preproc_dict_data(args[0])
+
+        # args here should already be `DataType`
+        # NOTE: how to right type to: unpacked args <-> tuple of args to unpack
+        # issue with `self._preproc_dict_data` return args type
+
         self.transform_op.data_keys = self.transform_op.preproc_datakeys(data_keys)
 
-        self._validate_args_datakeys(*args, data_keys=self.transform_op.data_keys)
+        self._validate_args_datakeys(*args, data_keys=self.transform_op.data_keys)  # type: ignore
 
-        in_args = self._arguments_preproc(*args, data_keys=self.transform_op.data_keys)
+        in_args = self._arguments_preproc(*args, data_keys=self.transform_op.data_keys)  # type: ignore
 
         if params is None:
             if self._params is None:
@@ -276,6 +314,9 @@ def inverse(  # type: ignore[override]
 
         outputs = self._arguments_postproc(args, outputs, data_keys=self.transform_op.data_keys)  # type: ignore
 
+        if isinstance(original_keys, tuple):
+            return {k: v for v, k in zip(outputs, original_keys)}
+
         if len(outputs) == 1 and isinstance(outputs, list):
             return outputs[0]
 
@@ -341,18 +382,23 @@ def _arguments_postproc(
 
     def forward(  # type: ignore[override]
         self,
-        *args: DataType,
+        *args: Union[DataType, Dict[str, DataType]],
         params: Optional[List[ParamItem]] = None,
         data_keys: Optional[Union[List[str], List[int], List[DataKey]]] = None,
-    ) -> Union[DataType, List[DataType]]:
+    ) -> Union[DataType, List[DataType], Dict[str, DataType]]:
         """Compute multiple tensors simultaneously according to ``self.data_keys``."""
         self.clear_state()
 
+        # Unpack/handle dictionary args
+        original_keys = None
+        if len(args) == 1 and isinstance(args[0], dict):
+            original_keys, data_keys, args = self._preproc_dict_data(args[0])
+
         self.transform_op.data_keys = self.transform_op.preproc_datakeys(data_keys)
 
-        self._validate_args_datakeys(*args, data_keys=self.transform_op.data_keys)
+        self._validate_args_datakeys(*args, data_keys=self.transform_op.data_keys)  # type: ignore
 
-        in_args = self._arguments_preproc(*args, data_keys=self.transform_op.data_keys)
+        in_args = self._arguments_preproc(*args, data_keys=self.transform_op.data_keys)  # type: ignore
 
         if params is None:
             # image data must exist if params is not provided.
@@ -386,11 +432,44 @@ def forward(  # type: ignore[override]
 
         self._params = params
 
+        if isinstance(original_keys, tuple):
+            return {k: v for v, k in zip(outputs, original_keys)}
+
         if len(outputs) == 1 and isinstance(outputs, list):
             return outputs[0]
 
         return outputs
 
+    def _preproc_dict_data(
+        self, data: Dict[str, DataType]
+    ) -> Tuple[Tuple[str, ...], List[DataKey], Tuple[DataType, ...]]:
+        if self.data_keys is not None:
+            raise ValueError("If you are using a dictionary as input, the data_keys should be None.")
+
+        data_keys = self._read_datakeys_from_dict(tuple(data.keys()))
+        keys = tuple(data.keys())
+        data_unpacked = tuple(v for v in data.values())
+
+        return keys, data_keys, data_unpacked
+
+    def _read_datakeys_from_dict(self, keys: Sequence[str]) -> List[DataKey]:
+        def retrieve_key(key: str) -> DataKey:
+            """Try to retrieve the datakey value by matching `<datakey>*`"""
+            # Alias cases, like INPUT, will not be get by the enum iterator.
+            if key.upper().startswith("INPUT"):
+                return DataKey.INPUT
+
+            for dk in DataKey:
+                if key.upper().startswith(dk.name):
+                    return DataKey.get(dk.name)
+
+            allowed_dk = " | ".join(f"`{d.name}`" for d in DataKey)
+            raise ValueError(
+                f"You input data dictionary keys should starts with some of datakey values: {allowed_dk}. Got `{key}`"
+            )
+
+        return [DataKey.get(retrieve_key(k)) for k in keys]
+
     def _preproc_boxes(self, arg: DataType, dcate: DataKey) -> Boxes:
         if DataKey.get(dcate) in [DataKey.BBOX]:
             mode = "vertices_plus"
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
@@ -67,20 +67,25 @@ def inverse(cls, input: T, module: Module, param: ParamItem, extra_args: Dict[st
 
 
 class AugmentationSequentialOps:
-    def __init__(self, data_keys: List[DataKey]) -> None:
+    def __init__(self, data_keys: Optional[List[DataKey]]) -> None:
         self._data_keys = data_keys
 
     @property
-    def data_keys(self) -> List[DataKey]:
+    def data_keys(self) -> Optional[List[DataKey]]:
         return self._data_keys
 
     @data_keys.setter
-    def data_keys(self, data_keys: List[DataKey]) -> None:
-        self._data_keys = [DataKey.get(inp) for inp in data_keys]
+    def data_keys(self, data_keys: Optional[Union[List[DataKey], List[str], List[int]]]) -> None:
+        if data_keys:
+            self._data_keys = [DataKey.get(inp) for inp in data_keys]
+        else:
+            self._data_keys = None
 
     def preproc_datakeys(self, data_keys: Optional[Union[List[str], List[int], List[DataKey]]] = None) -> List[DataKey]:
         if data_keys is None:
-            return self.data_keys
+            if isinstance(self.data_keys, list):
+                return self.data_keys
+            raise ValueError("Sequential ops needs data keys to be able to process.")
         else:
             return [DataKey.get(inp) for inp in data_keys]
 


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
@@ -35,6 +35,12 @@ def reproducibility_test(input, seq):
         assert_close(output_1, output_2[0])
     elif isinstance(output_2, (torch.Tensor,)) and isinstance(output_1, (torch.Tensor,)):
         assert_close(output_1, output_2, msg=f"{seq._params}")
+    elif isinstance(output_1, dict) and isinstance(output_2, dict):
+        [
+            assert_close(o1, o2)
+            for o1, o2 in zip(output_1.values(), output_2.values())
+            if isinstance(o1, (torch.Tensor,)) and isinstance(o2, (torch.Tensor,))
+        ]
     else:
         assert False, ("cannot compare", type(output_1), type(output_2))
 
@@ -706,6 +712,47 @@ def test_transform_list_of_masks_and_boxes(self, device, dtype):
         assert len(out[1]) == len(bbox)
         assert len(out[2]) == len(mask)
 
+    @pytest.mark.slow
+    @pytest.mark.parametrize("random_apply", [1, (2, 2), (1, 2), (2,), 10, True, False])
+    def test_dict_as_input_forward_and_inverse(self, random_apply, device, dtype):
+        inp = torch.randn(1, 3, 1000, 500, device=device, dtype=dtype)
+        bbox = torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype)
+        keypoints = torch.tensor([[[465, 115], [545, 116]]], device=device, dtype=dtype)
+        mask = bbox_to_mask(
+            torch.tensor([[[155, 0], [900, 0], [900, 400], [155, 400]]], device=device, dtype=dtype), 1000, 500
+        )[:, None]
+        aug = K.AugmentationSequential(
+            K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)),
+            K.AugmentationSequential(
+                K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0),
+                K.RandomAffine(360, p=1.0),
+                K.RandomAffine(360, p=1.0),
+                data_keys=["input", "mask", "bbox", "keypoints"],
+            ),
+            K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0),
+            K.RandomAffine(360, p=1.0),
+            data_keys=None,
+            random_apply=random_apply,
+        )
+
+        data = {"input": inp, "mask": mask, "bbox": bbox, "keypoints": keypoints}
+        out = aug(data)
+        assert out["input"].shape == inp.shape
+        assert out["mask"].shape == mask.shape
+        assert out["bbox"].shape == bbox.shape
+        assert out["keypoints"].shape == keypoints.shape
+        assert set(out["mask"].unique().tolist()).issubset(set(mask.unique().tolist()))
+
+        out_inv = aug.inverse(out)
+        assert out_inv["input"].shape == inp.shape
+        assert out_inv["mask"].shape == mask.shape
+        assert out_inv["bbox"].shape == bbox.shape
+        assert out_inv["keypoints"].shape == keypoints.shape
+        assert set(out_inv["mask"].unique().tolist()).issubset(set(mask.unique().tolist()))
+
+        if random_apply is False:
+            reproducibility_test(data, aug)
+
     @pytest.mark.jit()
     @pytest.mark.skip(reason="turn off due to Union Type")
     def test_jit(self, device, dtype):
diff --git a/tests/contrib/test_prompter.py b/tests/contrib/test_prompter.py
--- a/tests/contrib/test_prompter.py
+++ b/tests/contrib/test_prompter.py
@@ -88,17 +88,21 @@ def test_module(self):
 
     def test_dynamo(self, device, torch_optimizer):
         dtype = torch.float32
-        img = torch.rand(3, 128, 75, device=device, dtype=dtype)
+        batch_size = 1
+        N = 2
+        inpt = torch.rand(3, 77, 128, device=device, dtype=dtype)
+        keypoints = torch.randint(0, min(inpt.shape[-2:]), (batch_size, N, 2), device=device).to(dtype=dtype)
+        labels = torch.randint(0, 1, (batch_size, N), device=device).to(dtype=dtype)
 
         prompter = VisualPrompter(SamConfig("vit_b"), device, dtype)
-        prompter.set_image(img)
+        prompter.set_image(inpt)
 
-        expected = prompter.predict(img)
+        expected = prompter.predict(keypoints=keypoints, keypoints_labels=labels)
         prompter.reset_image()
 
         prompter.compile()
-        prompter.set_image(img)
-        actual = prompter.predict(img)
+        prompter.set_image(inpt)
+        actual = prompter.predict(keypoints=keypoints, keypoints_labels=labels)
 
         self.assert_close(expected.logits, actual.logits)
         self.assert_close(expected.scores, actual.scores)