Repository: kornia/kornia

Issue:
DataKey.CLASS not supported by AugmentationSequential
### Describe the bug

Label arguments were removed in https://github.com/kornia/kornia/pull/2117 but the CLASS data key is not actually supported by AugmentationSequential and throws an error:
```
 ---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[18], line 11
      6 labels = torch.ones((1,1))
      7 aug = K.AugmentationSequential(
      8     K.RandomCrop((3,3), pad_if_needed=True),
      9     data_keys=[DataKey.INPUT, DataKey.CLASS])
---> 11 transformed = aug(img, labels)

File ~/miniconda3/envs/ssl/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File ~/miniconda3/envs/ssl/lib/python3.10/site-packages/kornia/augmentation/container/augment.py:329, in AugmentationSequential.forward(self, params, data_keys, *args)
    325 self.transform_op.data_keys = self.transform_op.preproc_datakeys(data_keys)
    327 self._validate_args_datakeys(*args, data_keys=self.transform_op.data_keys)
--> 329 in_args = self._arguments_preproc(*args, data_keys=self.transform_op.data_keys)
    331 if params is None:
    332     # image data must exist if params is not provided.
    333     if DataKey.INPUT in self.transform_op.data_keys:

File ~/miniconda3/envs/ssl/lib/python3.10/site-packages/kornia/augmentation/container/augment.py:280, in AugmentationSequential._arguments_preproc(self, data_keys, *args)
    278         inp.append(self._preproc_boxes(arg, dcate))
    279     else:
--> 280         raise NotImplementedError(f"input type of {dcate} is not implemented.")
    281 return inp

NotImplementedError: input type of DataKey.CLASS is not implemented.
```

### Reproduction steps

```bash
import kornia.augmentation as K
from kornia.constants import DataKey
import torch

img = torch.zeros((1,1,5,5))
labels = torch.ones((1,1))
aug = K.AugmentationSequential(
    K.RandomCrop((3,3), pad_if_needed=True),
    data_keys=[DataKey.INPUT, DataKey.CLASS])

img, labels = aug(img, labels)
```


### Expected behavior

No Error and labels remains unchanged.

### Environment

```shell
PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.90.1-microsoft-standard-WSL2-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   39 bits physical, 48 bits virtual
CPU(s):                          8
On-line CPU(s) list:             0-7
Thread(s) per core:              2
Core(s) per socket:              4
Socket(s):                       1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           126
Model name:                      Intel(R) Core(TM) i7-1065G7 CPU @ 1.30GHz
Stepping:                        5
CPU MHz:                         1497.603
BogoMIPS:                        2995.20
Virtualization:                  VT-x
Hypervisor vendor:               Microsoft
Virtualization type:             full
L1d cache:                       192 KiB
L1i cache:                       128 KiB
L2 cache:                        2 MiB
L3 cache:                        8 MiB
Vulnerability Itlb multihit:     KVM: Mitigation: VMX disabled
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:          Mitigation; Enhanced IBRS
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Unknown: Dependent on hypervisor status
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology tsc_reliable nonstop_tsc cpuid pni pclmulqdq vmx ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx512vbmi umip avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid fsrm flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==2.0.1
[pip3] torchvision==0.15.2
[conda] blas                      1.0                         mkl    anaconda
[conda] cpuonly                   2.0                           0    pytorch
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640    anaconda
[conda] mkl-service               2.4.0           py310h7f8727e_0    anaconda
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0    anaconda
[conda] mkl_random                1.2.2           py310h00e6091_0    anaconda
[conda] numpy                     1.23.5          py310hd5efca6_0    anaconda
[conda] numpy-base                1.23.5          py310h8e6c178_0    anaconda
[conda] pytorch                   2.0.1              py3.10_cpu_0    pytorch
[conda] pytorch-mutex             1.0                         cpu    pytorch
[conda] torchvision               0.15.2                py310_cpu    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
@@ -275,6 +275,8 @@ def _arguments_preproc(self, *args: DataType, data_keys: List[DataKey]) -> List[
                 inp.append(self._preproc_keypoints(arg, dcate))
             elif DataKey.get(dcate) in _BOXES_OPTIONS:
                 inp.append(self._preproc_boxes(arg, dcate))
+            elif DataKey.get(dcate) is DataKey.CLASS:
+                inp.append(arg)
             else:
                 raise NotImplementedError(f"input type of {dcate} is not implemented.")
         return inp
@@ -307,6 +309,9 @@ def _arguments_postproc(
                         _out_b = _out_b.type(in_arg.dtype)
                 out.append(_out_b)
 
+            elif DataKey.get(dcate) is DataKey.CLASS:
+                out.append(out_arg)
+
             else:
                 raise NotImplementedError(f"input type of {dcate} is not implemented.")
 
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
@@ -91,6 +91,8 @@ def _get_op(self, data_key: DataKey) -> Type[SequentialOpsInterface[Any]]:
             return BoxSequentialOps
         if data_key == DataKey.KEYPOINTS:
             return KeypointSequentialOps
+        if data_key == DataKey.CLASS:
+            return ClassSequentialOps
         raise RuntimeError(f"Operation for `{data_key.name}` is not found.")
 
     def transform(
@@ -195,6 +197,22 @@ def inverse(cls, input: Tensor, module: Module, param: ParamItem, extra_args: Di
         return input
 
 
+class ClassSequentialOps(SequentialOpsInterface[Tensor]):
+    """Apply and inverse transformations for class labels if needed."""
+
+    @classmethod
+    def transform(cls, input: Tensor, module: Module, param: ParamItem, extra_args: Dict[str, Any] = {}) -> Tensor:
+        if isinstance(module, K.MixAugmentationBaseV2):
+            raise NotImplementedError(
+                "The support for class labels for mix augmentations that change the class label is not yet supported."
+            )
+        return input
+
+    @classmethod
+    def inverse(cls, input: Tensor, module: Module, param: ParamItem, extra_args: Dict[str, Any] = {}) -> Tensor:
+        return input
+
+
 class MaskSequentialOps(SequentialOpsInterface[Tensor]):
     """Apply and inverse transformations for mask tensors."""
 


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
@@ -505,6 +505,14 @@ def test_bbox(self, bbox, augmentation, device, dtype):
         for i in range(len(bbox)):
             assert len(bboxes_transformed[i]) == len(bbox[i])
 
+    def test_class(self, device, dtype):
+        img = torch.zeros((5, 1, 5, 5))
+        labels = torch.randint(0, 10, size=(5, 1))
+        aug = K.AugmentationSequential(K.RandomCrop((3, 3), pad_if_needed=True), data_keys=['input', 'class'])
+
+        _, out_labels = aug(img, labels)
+        assert labels is out_labels
+
     @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)