[Bug] RandomJPEG does fails if sides are not divisible by 16

[ternaus]

Describe the bug

RandomJPEG does fails if sides are not divisible by 16

Reproduction steps

In [1]: from  kornia.augmentation import RandomJPEG

In [2]: import torch
   ...: rng = torch.manual_seed(0)
   ...: images = 0.1904 * torch.ones(2, 3, 33, 37)
   ...: aug = RandomJPEG(jpeg_quality=(1.0, 50.0), p=1.)
   ...: images_jpeg = aug(images)
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
Cell In[2], line 5
      3 images = 0.1904 * torch.ones(2, 3, 33, 37)
      4 aug = RandomJPEG(jpeg_quality=(1.0, 50.0), p=1.)
----> 5 images_jpeg = aug(images)

File ~/anaconda3/envs/albumentations_benchmark/lib/python3.10/site-packages/torch/nn/modules/module.py:1511, in Module._wrapped_call_impl(self, *args, **kwargs)
   1509     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1510 else:
-> 1511     return self._call_impl(*args, **kwargs)

File ~/anaconda3/envs/albumentations_benchmark/lib/python3.10/site-packages/torch/nn/modules/module.py:1520, in Module._call_impl(self, *args, **kwargs)
   1515 # If we don't have any hooks, we want to skip the rest of the logic in
   1516 # this function, and just call forward.
   1517 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1518         or _global_backward_pre_hooks or _global_backward_hooks
   1519         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1520     return forward_call(*args, **kwargs)
   1522 try:
   1523     result = None

File ~/anaconda3/envs/albumentations_benchmark/lib/python3.10/site-packages/kornia/augmentation/base.py:210, in _BasicAugmentationBase.forward(self, input, params, **kwargs)
    206     params["batch_prob"] = tensor([True] * batch_shape[0])
    208 params, flags = self._process_kwargs_to_params_and_flags(params, self.flags, **kwargs)
--> 210 output = self.apply_func(in_tensor, params, flags)
    211 return self.transform_output_tensor(output, input_shape) if self.keepdim else output

File ~/anaconda3/envs/albumentations_benchmark/lib/python3.10/site-packages/kornia/augmentation/_2d/base.py:129, in RigidAffineAugmentationBase2D.apply_func(self, in_tensor, params, flags)
    126     flags = self.flags
    128 trans_matrix = self.generate_transformation_matrix(in_tensor, params, flags)
--> 129 output = self.transform_inputs(in_tensor, params, flags, trans_matrix)
    130 self._transform_matrix = trans_matrix
    132 return output

File ~/anaconda3/envs/albumentations_benchmark/lib/python3.10/site-packages/kornia/augmentation/base.py:261, in _AugmentationBase.transform_inputs(self, input, params, flags, transform, **kwargs)
    259 self.validate_tensor(in_tensor)
    260 if to_apply.all():
--> 261     output = self.apply_transform(in_tensor, params, flags, transform=transform)
    262 elif not to_apply.any():
    263     output = self.apply_non_transform(in_tensor, params, flags, transform=transform)

File ~/anaconda3/envs/albumentations_benchmark/lib/python3.10/site-packages/kornia/augmentation/_2d/intensity/jpeg.py:56, in RandomJPEG.apply_transform(self, input, params, flags, transform)
     53 def apply_transform(
     54     self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
     55 ) -> Tensor:
---> 56     jpeg_output: Tensor = jpeg_codec_differentiable(input, params["jpeg_quality"])
     57     return jpeg_output

File ~/anaconda3/envs/albumentations_benchmark/lib/python3.10/site-packages/kornia/utils/image.py:231, in perform_keep_shape_image.<locals>._wrapper(input, *args, **kwargs)
    229 input_shape = input.shape
    230 input = _to_bchw(input)  # view input as (B, C, H, W)
--> 231 output = f(input, *args, **kwargs)
    232 if len(input_shape) == 3:
    233     output = output[0]

File ~/anaconda3/envs/albumentations_benchmark/lib/python3.10/site-packages/kornia/enhance/jpeg.py:440, in jpeg_codec_differentiable(image_rgb, jpeg_quality, quantization_table_y, quantization_table_c)
    438 # Check shape of inputs
    439 KORNIA_CHECK_SHAPE(image_rgb, ["*", "3", "H", "W"])
--> 440 KORNIA_CHECK(
    441     (image_rgb.shape[-1] % 16 == 0) and (image_rgb.shape[-2] % 16 == 0),
    442     f"image dimension must be divisible by 16. Got the shape {image_rgb.shape}.",
    443 )
    444 KORNIA_CHECK_SHAPE(jpeg_quality, ["B"])
    445 # Add batch dimension to quantization tables if needed

File ~/anaconda3/envs/albumentations_benchmark/lib/python3.10/site-packages/kornia/core/check.py:103, in KORNIA_CHECK(condition, msg, raises)
    101 if not condition:
    102     if raises:
--> 103         raise Exception(f"{condition} not true.\n{msg}")
    104     return False
    105 return True

Exception: False not true.
image dimension must be divisible by 16. Got the shape torch.Size([2, 3, 33, 37]).

### Expected behavior

-

### Environment

```shell
-

Additional context

No response

[edgarriba]

@ternaus thanks to rise this issue. @ChristophReich1996 is this a bug or a limitation of the algorithm ?

[ChristophReich1996]

Hi @ternaus & @edgarriba JPEG performs coding based on 8 x 8 patches. As we downsample the chroma by a factor of two this corresponds to a patch size of 16 x 16 w.r.t. the original image. So the input image dimensions must be devisable by 16. I would consider padding the image, performing JPEG coding, and clipping the image again. @edgarriba we could also add this to the differentiable JPEG implementation if you like.

[edgarriba]

@ChristophReich1996 sure, feel free to add the patch