Repository: kornia/kornia

Issue:
rgb_to_yuv420 does not support (*, 3, H, W) format
### Describe the bug

The documentation of the `rgb_to_yuv420` function states that the `(*, 3, H, W)` format is supported. However, the average pooling function (torch.nn.functional.avg_pool2d) used after converting to YUV does not support an arbitrary amount of dimensions in the front.

### Reproduction steps

```bash
import torch
import kornia
from torch import Tensor

images: Tensor = torch.rand(2, 2, 3, 16, 16)
images_y, images_uv = kornia.color.rgb_to_yuv420(images)  # type: Tensor, Tensor
```


### Expected behavior

RuntimeError: Expected 3D or 4D (batch mode) tensor with optional 0 dim batch size for input, but got:[2, 2, 2, 16, 16]

### Environment

```shell
- PyTorch Version: `1.12.0`
- How you installed PyTorch: `conda`
- Python version: `3.9.12`
- CUDA/cuDNN version: CPU utilized
- Any other relevant information:
```


### Additional context

_No response_


Discussion:
@ChristophReich1996 thanks for pointing this out, feel free to open a PR with the fix.

Resolving patch:
diff --git a/kornia/color/yuv.py b/kornia/color/yuv.py
--- a/kornia/color/yuv.py
+++ b/kornia/color/yuv.py
@@ -68,7 +68,7 @@ def rgb_to_yuv420(image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
 
     yuvimage = rgb_to_yuv(image)
 
-    return (yuvimage[..., :1, :, :], torch.nn.functional.avg_pool2d(yuvimage[..., 1:3, :, :], (2, 2)))
+    return (yuvimage[..., :1, :, :], yuvimage[..., 1:3, :, :].unfold(-2, 2, 2).unfold(-2, 2, 2).mean((-1, -2)))
 
 
 def rgb_to_yuv422(image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -99,7 +99,7 @@ def rgb_to_yuv422(image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
 
     yuvimage = rgb_to_yuv(image)
 
-    return (yuvimage[..., :1, :, :], torch.nn.functional.avg_pool2d(yuvimage[..., 1:3, :, :], (1, 2)))
+    return (yuvimage[..., :1, :, :], yuvimage[..., 1:3, :, :].unfold(-1, 2, 2).mean(-1))
 
 
 def yuv_to_rgb(image: torch.Tensor) -> torch.Tensor:


Test patch:
diff --git a/test/color/test_yuv.py b/test/color/test_yuv.py
--- a/test/color/test_yuv.py
+++ b/test/color/test_yuv.py
@@ -72,7 +72,7 @@ def test_smoke(self, device, dtype):
         assert isinstance(kornia.color.rgb_to_yuv420(img)[0], torch.Tensor)
         assert isinstance(kornia.color.rgb_to_yuv420(img)[1], torch.Tensor)
 
-    @pytest.mark.parametrize("shape", [(1, 3, 4, 4), (2, 3, 2, 4), (3, 3, 4, 2), (3, 2, 2)])
+    @pytest.mark.parametrize("shape", [(1, 3, 4, 4), (2, 3, 2, 4), (3, 3, 4, 2), (3, 2, 2), (3, 3, 3, 4, 4)])
     def test_cardinality(self, device, dtype, shape):
         img = torch.ones(shape, device=device, dtype=dtype)
         shapey = list(shape)
@@ -229,7 +229,7 @@ def test_smoke(self, device, dtype):
         assert isinstance(kornia.color.rgb_to_yuv422(img)[0], torch.Tensor)
         assert isinstance(kornia.color.rgb_to_yuv422(img)[1], torch.Tensor)
 
-    @pytest.mark.parametrize("shape", [(1, 3, 4, 4), (2, 3, 2, 4), (3, 3, 4, 2), (3, 2, 2)])
+    @pytest.mark.parametrize("shape", [(1, 3, 4, 4), (2, 3, 2, 4), (3, 3, 4, 2), (3, 2, 2), (3, 3, 3, 4, 4)])
     def test_cardinality(self, device, dtype, shape):
         img = torch.ones(shape, device=device, dtype=dtype)
         shapey = list(shape)