Repository: kornia/kornia

Issue:
ColorJiggle only works on 3-channel images, leaving no high-level abstraction for brightness/contrast for non-3ch images
### Describe the bug

Kornia's `ColorJiggle` throws an error when used on non-3-channel images, even when hue is not changed. As a consequence, Kornia doesn't seem to have any higher-level abstraction for creating random brightness or contrast augmentations unless an image has 3 channels.

### Reproduction steps

1. Demo code to reproduce:

```py
demo_aug = K.AugmentationSequential(
        K.ColorJiggle(brightness=0.2 / 255., p=0.9),
        data_keys=["input"],
        same_on_batch=False,
    )
print(demo_aug)

demo_aug(torch.rand((4, 2, 192, 192)))
```

The output:

```py
AugmentationSequential(
  (ColorJiggle_0): ColorJiggle(ColorJiggle(brightness=0.0007843137254901962, contrast=0.0, saturation=0.0, hue=0.0, p=0.9, p_batch=1.0, same_on_batch=False))
)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_121730/3049129453.py in <module>
      6 print(demo_aug)
      7 
----> 8 demo_aug(torch.rand((4, 2, 192, 192)))

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1108         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1110             return forward_call(*input, **kwargs)
   1111         # Do not call functions when jit is used
   1112         full_backward_hooks, non_full_backward_hooks = [], []

/opt/conda/lib/python3.7/site-packages/kornia/augmentation/container/augment.py in forward(self, label, params, data_keys, *args)
    377                 _inp = args[idx]
    378 
--> 379                 _out = super().forward(_inp, label, params=params, extra_args=extra_args)  # type: ignore
    380                 self._transform_matrix = self.get_transformation_matrix(_inp, params=params)
    381                 if self.return_label:

/opt/conda/lib/python3.7/site-packages/kornia/augmentation/container/image.py in forward(self, input, label, params, extra_args)
    379         for param in params:
    380             module = self.get_submodule(param.name)
--> 381             input, label = self.apply_to_input(input, label, module, param=param, extra_args=extra_args)  # type: ignore
    382             if isinstance(module, (_AugmentationBase, MixAugmentationBase, SequentialBase)):
    383                 param = ParamItem(param.name, module._params)

/opt/conda/lib/python3.7/site-packages/kornia/augmentation/container/image.py in apply_to_input(self, input, label, module, param, extra_args)
    213         if module is None:
    214             module = self.get_submodule(param.name)
--> 215         return self.apply_inverse_func.apply_trans(input, label, module, param, extra_args)  # type: ignore
    216 
    217     def forward_parameters(self, batch_shape: torch.Size) -> List[ParamItem]:

/opt/conda/lib/python3.7/site-packages/kornia/augmentation/container/utils.py in apply_trans(cls, input, label, module, param, extra_args)
    181             input, label = module(input, label=label, params=param.data)
    182         elif isinstance(module, (_AugmentationBase,)):
--> 183             input = module(input, params=param.data, **extra_args)
    184         elif isinstance(module, kornia.augmentation.ImageSequential):
    185             temp = module.apply_inverse_func

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1108         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1110             return forward_call(*input, **kwargs)
   1111         # Do not call functions when jit is used
   1112         full_backward_hooks, non_full_backward_hooks = [], []

/opt/conda/lib/python3.7/site-packages/kornia/augmentation/base.py in forward(self, input, params, **kwargs)
    181         params, flags = self._process_kwargs_to_params_and_flags(params, self.flags, **kwargs)
    182 
--> 183         output = self.apply_func(in_tensor, params, flags)
    184         return self.transform_output_tensor(output, input_shape) if self.keepdim else output
    185 

/opt/conda/lib/python3.7/site-packages/kornia/augmentation/base.py in apply_func(self, in_tensor, params, flags)
    254         elif to_apply.all():
    255             trans_matrix = self.compute_transformation(in_tensor, params=params, flags=flags)
--> 256             output = self.apply_transform(in_tensor, params=params, flags=flags, transform=trans_matrix)
    257         else:
    258             output = in_tensor.clone()

/opt/conda/lib/python3.7/site-packages/kornia/augmentation/_2d/intensity/color_jiggle.py in apply_transform(self, input, params, flags, transform)
     91         for idx in params["order"].tolist():
     92             t = transforms[idx]
---> 93             jittered = t(jittered)
     94 
     95         return jittered

/opt/conda/lib/python3.7/site-packages/kornia/augmentation/_2d/intensity/color_jiggle.py in <lambda>(img)
     84             lambda img: adjust_brightness(img, params["brightness_factor"] - 1),
     85             lambda img: adjust_contrast(img, params["contrast_factor"]),
---> 86             lambda img: adjust_saturation(img, params["saturation_factor"]),
     87             lambda img: adjust_hue(img, params["hue_factor"] * 2 * pi),
     88         ]

/opt/conda/lib/python3.7/site-packages/kornia/enhance/adjust.py in adjust_saturation(image, factor)
    129 
    130     # convert the rgb image to hsv
--> 131     x_hsv: Tensor = rgb_to_hsv(image)
    132 
    133     # perform the conversion

/opt/conda/lib/python3.7/site-packages/kornia/color/hsv.py in rgb_to_hsv(image, eps)
     32 
     33     if len(image.shape) < 3 or image.shape[-3] != 3:
---> 34         raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
     35 
     36     max_rgb, argmax_rgb = image.max(-3)

ValueError: Input size must have a shape of (*, 3, H, W). Got torch.Size([4, 2, 192, 192])
```

The issue seems to be a consequence of the call to `adjust_hue` even when no hue is set:
```py
lambda img: adjust_hue(img, params["hue_factor"] * 2 * pi),
```

### Expected behavior

I would expect either that `ColorJiggle` works with any number of channels, or that the documentation would explain that it only works with 3 channel images. Because the issue seems to be that `adjust_hue` is called even when no hue changes are requested, it seems possible to fix this to permit changing brightness/contrast for non-3ch images.

### Environment

```shell
Kornia version 0.6.5.


Collecting environment information...
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
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.16.3
Libc version: glibc-2.10

Python version: 3.7.12 | packaged by conda-forge | (default, Oct 26 2021, 06:08:53)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.13.0-1033-gcp-x86_64-with-debian-bullseye-sid
Is CUDA available: True
CUDA runtime version: 11.3.109
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB
Nvidia driver version: 470.129.06
cuDNN version: Probably one of the following:
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_train.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_train.so.8
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.19.5
[pip3] torch==1.11.0
[pip3] torch-lr-finder==0.2.1
[pip3] torch-xla==1.11
[pip3] torchmetrics==0.9.1
[pip3] torchvision==0.12.0+cu113
[conda] blas                      2.114                       mkl    conda-forge
[conda] blas-devel                3.9.0            14_linux64_mkl    conda-forge
[conda] cudatoolkit               11.3.1               ha36c431_9    nvidia
[conda] dlenv-pytorch-1-11-gpu    1.0.20220525     py37hc1c1d6d_0    file:///tmp/conda-pkgs
[conda] libblas                   3.9.0            14_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            14_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            14_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            14_linux64_mkl    conda-forge
[conda] mkl                       2022.0.1           h06a4308_117  
[conda] mkl-devel                 2022.0.1           h66538d2_117  
[conda] mkl-include               2022.0.1           h06a4308_117  
[conda] numpy                     1.19.5           py37h3e96413_3    conda-forge
[conda] pytorch                   1.11.0          py3.7_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch-lr-finder           0.2.1                    pypi_0    pypi
[conda] torchmetrics              0.9.1                    pypi_0    pypi
[conda] torchvision               0.12.0+cu113             pypi_0    pypi
```
```


### Additional context

_No response_


Discussion:
A demonstration that removes the hue and saturation is here. Note that this is not really production-level (the comment is not updated and the index traversal is just replaced with [0, 1]).

```py

from typing import Any, Dict, List, Optional, Tuple, Union, cast

from torch import Tensor

from kornia.augmentation import random_generator as rg
from kornia.augmentation._2d.intensity.base import IntensityAugmentationBase2D
from kornia.constants import pi
from kornia.enhance import adjust_brightness, adjust_contrast, adjust_hue, adjust_saturation


class ColorJiggleNoHueSat(IntensityAugmentationBase2D):
    r"""Apply a random transformation to the brightness, contrast, saturation and hue of a tensor image.
    .. image:: _static/img/ColorJiggle.png
    Args:
        p: probability of applying the transformation.
        brightness: The brightness factor to apply.
        contrast: The contrast factor to apply.
        same_on_batch: apply the same transformation across the batch.
        keepdim: whether to keep the output shape the same as input (True) or broadcast it
                 to the batch form (False).
    Shape:
        - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
        - Output: :math:`(B, C, H, W)`
    .. note::
        This function internally uses :func:`kornia.enhance.adjust_brightness`,
        :func:`kornia.enhance.adjust_contrast`. :func:`kornia.enhance.adjust_saturation`,
        :func:`kornia.enhance.adjust_hue`.
    Examples:
        >>> rng = torch.manual_seed(0)
        >>> inputs = torch.ones(1, 3, 3, 3)
        >>> aug = ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.)
        >>> aug(inputs)
        tensor([[[[0.9993, 0.9993, 0.9993],
                  [0.9993, 0.9993, 0.9993],
                  [0.9993, 0.9993, 0.9993]],
        <BLANKLINE>
                 [[0.9993, 0.9993, 0.9993],
                  [0.9993, 0.9993, 0.9993],
                  [0.9993, 0.9993, 0.9993]],
        <BLANKLINE>
                 [[0.9993, 0.9993, 0.9993],
                  [0.9993, 0.9993, 0.9993],
                  [0.9993, 0.9993, 0.9993]]]])
    To apply the exact augmenation again, you may take the advantage of the previous parameter state:
        >>> input = torch.randn(1, 3, 32, 32)
        >>> aug = ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.)
        >>> (aug(input) == aug(input, params=aug._params)).all()
        tensor(True)
    """

    def __init__(
        self,
        brightness: Union[Tensor, float, Tuple[float, float], List[float]] = 0.0,
        contrast: Union[Tensor, float, Tuple[float, float], List[float]] = 0.0,
        same_on_batch: bool = False,
        p: float = 1.0,
        keepdim: bool = False,
        return_transform: Optional[bool] = None
    ) -> None:
        super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
        self.brightness = brightness
        self.contrast = contrast
        self._param_generator = cast(
            rg.ColorJiggleGenerator,
            rg.ColorJiggleGenerator(brightness, contrast)
        )

    def apply_transform(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:

        transforms = [
            lambda img: adjust_brightness(img, params["brightness_factor"] - 1),
            lambda img: adjust_contrast(img, params["contrast_factor"])
        ]

        jittered = input
        for idx in [0, 1]:
            t = transforms[idx]
            jittered = t(jittered)

        return jittered
```

```py
import kornia_colorjiggle

demo_aug = K.AugmentationSequential(
        kornia_colorjiggle.ColorJiggleNoHueSat(brightness=0.2 / 255., p=0.9),
        data_keys=["input"],
        same_on_batch=False,
    )
print(demo_aug)

demo_aug(torch.rand((4, 2, 192, 192)))
```

Output:
```py

AugmentationSequential(
  (ColorJiggleNoHueSat_0): ColorJiggleNoHueSat(ColorJiggleNoHueSat(brightness=0.0007843137254901962, contrast=0.0, saturation=0.0, hue=0.0, p=0.9, p_batch=1.0, same_on_batch=False))
)
tensor([[[[0.9948, 0.0908, 0.9733,  ..., 0.7482, 0.2822, 0.0888],
          [0.2008, 0.8121, 0.6717,  ..., 0.9255, 0.0358, 0.4732],
          [0.3072, 0.4883, 0.7782,  ..., 0.0435, 0.2232, 0.9481],
          ...,
          [0.9865, 0.4192, 0.4377,  ..., 0.6264, 0.1472, 0.5743],
          [0.0999, 0.4580, 0.5749,  ..., 0.0701, 0.0949, 0.4180],
          [0.2859, 0.7441, 0.5989,  ..., 0.4774, 0.5582, 0.2087]],

         [[0.6786, 0.3145, 0.7621,  ..., 0.4631, 0.7427, 0.9016],
          [0.5416, 0.1805, 0.4715,  ..., 0.5679, 0.4411, 0.5031],
          [0.1880, 0.8481, 0.0808,  ..., 0.0818, 0.4355, 0.6351],
          ...,
          [0.7578, 0.8784, 0.9843,  ..., 0.8396, 0.9102, 0.5230],
          [0.8288, 0.9679, 0.7950,  ..., 0.5180, 0.9704, 0.2313],
          [0.2153, 0.3708, 0.9490,  ..., 0.9106, 0.1457, 0.1454]]],


        [[[0.0557, 0.3585, 0.7216,  ..., 0.9909, 0.5756, 0.2530],
          [0.9736, 0.8083, 0.0387,  ..., 0.3766, 0.0597, 0.6945],
          [0.0827, 0.4826, 0.8262,  ..., 0.6595, 0.2257, 0.3791],
          ...,
          [0.1235, 0.9182, 0.9092,  ..., 0.1502, 0.9667, 0.3505],
          [0.1002, 0.4241, 0.4998,  ..., 0.2244, 0.8463, 0.2512],
          [0.6523, 0.2789, 0.0168,  ..., 0.7983, 0.9491, 0.3584]],

         [[0.9077, 0.7763, 0.4239,  ..., 0.1885, 0.3604, 0.4574],
          [0.4077, 0.4302, 0.6814,  ..., 0.6126, 0.4072, 0.7745],
          [0.6935, 0.2169, 0.8654,  ..., 0.9455, 0.3389, 0.3330],
          ...,
          [0.8218, 0.8024, 0.4391,  ..., 0.0080, 0.0995, 0.9209],
          [0.3962, 0.5513, 0.3949,  ..., 0.3945, 0.2764, 0.9512],
          [0.5543, 0.8952, 0.3558,  ..., 0.4007, 0.7832, 0.2977]]],


        [[[0.6009, 0.4036, 0.3634,  ..., 0.0591, 0.4571, 0.1786],
          [0.4956, 0.4754, 0.1618,  ..., 0.3527, 0.7206, 0.4789],
          [0.7287, 0.1992, 0.0713,  ..., 0.2533, 0.1311, 0.3443],
          ...,
          [0.9445, 0.9497, 0.2767,  ..., 0.5779, 0.8997, 0.3855],
          [0.9510, 0.5950, 0.5739,  ..., 0.9350, 0.7366, 0.6642],
          [0.1304, 0.2102, 0.7245,  ..., 0.2969, 0.0855, 0.7167]],

         [[0.4867, 0.0896, 0.0609,  ..., 0.9234, 0.0282, 0.5666],
          [0.8105, 0.2168, 0.5883,  ..., 0.7007, 0.5261, 0.3199],
          [0.3213, 0.8741, 0.5830,  ..., 0.0487, 0.9679, 0.9515],
          ...,
          [0.3388, 0.3430, 0.7313,  ..., 0.1223, 0.7133, 0.0338],
          [0.5868, 0.8998, 0.7896,  ..., 0.9280, 0.5704, 0.0096],
          [0.7013, 0.5951, 0.2819,  ..., 0.2213, 0.7685, 0.7428]]],


        [[[0.0712, 0.2783, 0.5838,  ..., 0.9053, 0.8947, 0.2358],
          [0.9648, 0.6963, 0.1458,  ..., 0.2773, 0.1931, 0.6698],
          [0.3212, 0.7523, 0.5937,  ..., 0.6726, 0.9006, 0.9310],
          ...,
          [0.3301, 0.1614, 0.9285,  ..., 0.9260, 0.1509, 0.9834],
          [0.9186, 0.9510, 0.8868,  ..., 0.6636, 0.9761, 0.5583],
          [0.8006, 0.0120, 0.1133,  ..., 0.5790, 0.4920, 0.5489]],

         [[0.3572, 0.5710, 0.7197,  ..., 0.3005, 0.1512, 0.7109],
          [0.8920, 0.2346, 0.1962,  ..., 0.5867, 0.1685, 0.8015],
          [0.2155, 0.0182, 0.7572,  ..., 0.4229, 0.2679, 0.5710],
          ...,
          [0.6283, 0.8613, 0.3597,  ..., 0.2404, 0.7876, 0.8425],
          [0.5726, 0.5801, 0.6839,  ..., 0.0859, 0.3145, 0.0959],
          [0.6508, 0.2111, 0.8503,  ..., 0.7243, 0.0166, 0.2031]]]])

```
@carbocation i've never been a fan of  the color jitter operator exactly for this low flexibility, the design pattern was inherited from torchvision mostly to be cross compatible with that library.

I'd advocate as you did to create you own operator with the core basic augmentation operators according to your needs.

One hack to do here would be to define the factors as `Optional` defaulting to `None` and skip/non-skip based on that. However, we will face in any case the corner case of `hue!=None` and `img.chanels!=3`.

/cc @shijianjian any thoughts here ?
Hello, has this issue been addressed? I am running into the same problems when using only 1 input channel.

ok, i'm moving this to backlog and fix. If anyone here wants to take a lead on this issue let us know, otherwise will go in the priority queue of current tasks.
I will check it tonight 

Resolving patch:
diff --git a/kornia/augmentation/_2d/intensity/color_jiggle.py b/kornia/augmentation/_2d/intensity/color_jiggle.py
--- a/kornia/augmentation/_2d/intensity/color_jiggle.py
+++ b/kornia/augmentation/_2d/intensity/color_jiggle.py
@@ -75,10 +75,16 @@ def apply_transform(
         self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
     ) -> Tensor:
         transforms = [
-            lambda img: adjust_brightness(img, params["brightness_factor"] - 1),
-            lambda img: adjust_contrast(img, params["contrast_factor"]),
-            lambda img: adjust_saturation(img, params["saturation_factor"]),
-            lambda img: adjust_hue(img, params["hue_factor"] * 2 * pi),
+            lambda img: adjust_brightness(img, params["brightness_factor"] - 1)
+            if (params["brightness_factor"] - 1 != 0).any()
+            else img,
+            lambda img: adjust_contrast(img, params["contrast_factor"])
+            if (params["contrast_factor"] != 1).any()
+            else img,
+            lambda img: adjust_saturation(img, params["saturation_factor"])
+            if (params["saturation_factor"] != 1).any()
+            else img,
+            lambda img: adjust_hue(img, params["hue_factor"] * 2 * pi) if (params["hue_factor"] != 0).any() else img,
         ]
 
         jittered = input
diff --git a/kornia/augmentation/_2d/intensity/color_jitter.py b/kornia/augmentation/_2d/intensity/color_jitter.py
--- a/kornia/augmentation/_2d/intensity/color_jitter.py
+++ b/kornia/augmentation/_2d/intensity/color_jitter.py
@@ -88,10 +88,16 @@ def apply_transform(
         self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
     ) -> Tensor:
         transforms = [
-            lambda img: adjust_brightness_accumulative(img, params["brightness_factor"]),
-            lambda img: adjust_contrast_with_mean_subtraction(img, params["contrast_factor"]),
-            lambda img: adjust_saturation_with_gray_subtraction(img, params["saturation_factor"]),
-            lambda img: adjust_hue(img, params["hue_factor"] * 2 * pi),
+            lambda img: adjust_brightness_accumulative(img, params["brightness_factor"])
+            if (params["brightness_factor"] != 0).any()
+            else img,
+            lambda img: adjust_contrast_with_mean_subtraction(img, params["contrast_factor"])
+            if (params["contrast_factor"] != 1).any()
+            else img,
+            lambda img: adjust_saturation_with_gray_subtraction(img, params["saturation_factor"])
+            if (params["saturation_factor"] != 1).any()
+            else img,
+            lambda img: adjust_hue(img, params["hue_factor"] * 2 * pi) if (params["hue_factor"] != 0).any() else img,
         ]
 
         jittered = input


Test patch:
diff --git a/tests/augmentation/test_augmentation.py b/tests/augmentation/test_augmentation.py
--- a/tests/augmentation/test_augmentation.py
+++ b/tests/augmentation/test_augmentation.py
@@ -978,10 +978,11 @@ def test_smoke(self):
         )
         assert str(f) == repr
 
-    def test_color_jiggle(self, device, dtype):
+    @pytest.mark.parametrize("C", [1, 3])
+    def test_color_jiggle(self, device, dtype, C):
         f = ColorJiggle()
 
-        input = torch.rand(3, 5, 5, device=device, dtype=dtype).unsqueeze(0)  # 3 x 5 x 5
+        input = torch.rand(C, 5, 5, device=device, dtype=dtype).unsqueeze(0)  # 3 x 5 x 5
         expected = input
 
         expected_transform = torch.eye(3, device=device, dtype=dtype).unsqueeze(0)  # 3 x 3