Repository: kornia/kornia

Issue:
[Feat] Add support to other color representations YUV
## 🚀 Feature
 Add support to other color representations YUV 420, 444


Discussion:
YUV 420/422 is typically separated in 2 or 3 planes, do you have thoughts on how best to represent this as torch tensors. E.g. Options seems to be
1. [B,1,H*1,5,W] typical one buffer format that mimics file layout. Quite ugly
2. [B,3,H,W] represent as 444 but just do the subsample operations so that e.g. every 2x2 pixels have the same value. Probably not what you are after
3. [B,1,H,W] (Y) + [B,2,H/2,W/2] for UV (or UV split as well into 3 different tensors). Somewhat sensible but having multiple tensors breaks with other formats.
Do you have any reference material for what you describe ? This was an external feature request, after quick search I found the following snippet.
https://gist.github.com/Quasimondo/c3590226c924a06b276d606f4f189639
> 3. [B,1,H,W] (Y) + [B,2,H/2,W/2] for UV (or UV split as well into 3 different tensors). Somewhat sensible but having multiple tensors breaks with other formats.

This option looks most reasonable.  Looks like opencv also supports some other YUV formats and has a function to;
"Converts an image from one color space to another **where the source image is stored in two planes**."
https://docs.opencv.org/4.5.3/d8/d01/group__imgproc__color__conversions.html#ga8e873314e72a1a6c0252375538fbf753

Would something like that work?


@edgarriba: Right, that link describes the color conversion itself which you already support (at least in a limited fashion, there are multiple different ways of doing this, i.e. rec 709, 2020 as well as 601). Problems comes form doing subsampling as in the case of YUV 420 where we end up with chroma channels that are lower resolution. 420 is 2x reduction both horizontal vertical while 422 split only horizontal. https://en.wikipedia.org/wiki/Chroma_subsampling is a good start.
@cceyda: Agreed that probably does make most  sense. I guess we can limit to just 2 planes as in that example and perhaps do 420 and 422 with just one chrome offset mode (i.e. defining were the chroma samples reside in relation to the luma samples)
if understand correctly we are targeting the following:
```python
def rgb_to_yub(img: Tensor) -> Tuple[Tensor, Tensor]:
    """Converts a tensor image from RGB to YUV color space.

    Args:
       img: the input RGB image with shape (*,3,H,W).

    Returns:
       A tuple containing the converted YUV image.
       - Tensor containing the Y plane with shape (*,1,H,W)
       - Tensor containing the UV planes with shape (*,2,H,W)
    """
    ....
```
The only concern is that later to compose functions will make usability a bit complicated
```python
pipe = nn.Sequential(
    RgbToYuv(),
    RandomRotate(...),  # this expects consistent tensor image
)
out = pipe(torch.rand(1,3,224,224)
```
I suppose the only thing that is compatible with e.g. the augmentations ops would be option 2 (keep same same tensor format just pretend we do subsampling). I suppose that is somewhat useful as an augmentation but I guess the more common use cases (without knowing the kornia use cases deeply):
1. YUV -> RGB transform when YUV data is subsampled, which can happen from compressed formats or, I guess most commonly, if having lower level access to a camera that happens to do YUV 420/422 output. I assume reading from compresed format usually happens through pytorch or opencv which won't serve you subsampled yuv data.
2. When we need to compress the results by e.g. passing data to a device that prefers subsampled YUV (hardware image/video encoders are probably the most common?) in this case this would typically be done as the last step of the pipeline. E.g. read data, do some computer based image manipulaiton, save the data.
3. RGB -> YUV when you algorithm calls for doing something in YUV colorspace instead, typically you would not use subsampled data but if you used it you would typically need som heavy tailoring of your algorithm anyway
Are there other uses? Otherwise the function you suggest (perhaps named rgb_to_yuv420 and yuv420_to_rgb?)
@oskarflordal the general use case for kornia is to have anything useful to be used during and after the training of neural networks with geometric vision focus. However, it's interesting for us to understand other use cases and how we can make kornia more transversal to different vision applications.

I understand the use cases you mention but overall for our purposes would make this operator not really compatible with the rest of kornia components. In general how we assume things is through TITO (tensor-in-tensor-out) and returning tuples produced conflicts in the past within our API. /cc @ducha-aiki @shijianjian 

Besides, I like the your naming proposal. Would you be open to send a PR ? we can assist you in the process.
@edgarriba Yes, I will prepare a PR (code is minimal but plenty of comments/tests to fill in...). I notice the documentation/implementation is inconsistent for the yuv to rgb functions themselves, I will fix that as well (input is supposed to be [-0.5,0.5] for uv instead of [0,1] for the transformation to work properly (I suggest to change the documentation since it is mostly the expected behaviour).
@oskarflordal that's awesome - definitely we can iterate over that.

Resolving patch:
diff --git a/kornia/color/yuv.py b/kornia/color/yuv.py
--- a/kornia/color/yuv.py
+++ b/kornia/color/yuv.py
@@ -50,8 +50,8 @@ def rgb_to_yuv420(image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         image: RGB Image to be converted to YUV with shape :math:`(*, 3, H, W)`.
 
     Returns:
-       Tensor containing the Y plane with shape :math:`(*, 1, H, W)
-       Tensor containing the UV planes with shape :math:`(*, 2, H/2, W/2)
+        A Tensor containing the Y plane with shape :math:`(*, 1, H, W)`
+        A Tensor containing the UV planes with shape :math:`(*, 2, H/2, W/2)`
 
     Example:
         >>> input = torch.rand(2, 3, 4, 6)
@@ -67,9 +67,8 @@ def rgb_to_yuv420(image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         raise ValueError(f"Input H&W must be evenly disible by 2. Got {image.shape}")
 
     yuvimage = rgb_to_yuv(image)
-    yuvchunks = torch.chunk(yuvimage, 3, dim=-3)
 
-    return (yuvchunks[0], torch.nn.functional.avg_pool2d(torch.cat(yuvchunks[1:3], dim=-3), (2, 2)))
+    return (yuvimage[..., :1, :, :], torch.nn.functional.avg_pool2d(yuvimage[..., 1:3, :, :], (2, 2)))
 
 
 def rgb_to_yuv422(image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -82,8 +81,8 @@ def rgb_to_yuv422(image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         image: RGB Image to be converted to YUV with shape :math:`(*, 3, H, W)`.
 
     Returns:
-       Tensor containing the Y plane with shape :math:`(*, 1, H, W)
-       Tensor containing the UV planes with shape :math:`(*, 2, H, W/2)
+       A Tensor containing the Y plane with shape :math:`(*, 1, H, W)`
+       A Tensor containing the UV planes with shape :math:`(*, 2, H, W/2)`
 
     Example:
         >>> input = torch.rand(2, 3, 4, 6)
@@ -100,10 +99,7 @@ def rgb_to_yuv422(image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
 
     yuvimage = rgb_to_yuv(image)
 
-    # use chunk to handle 3 and 4 dimensions
-    yuvchunks = torch.chunk(yuvimage, 3, dim=-3)
-
-    return (yuvchunks[0], torch.nn.functional.avg_pool2d(torch.cat(yuvchunks[1:3], dim=-3), (1, 2)))
+    return (yuvimage[..., :1, :, :], torch.nn.functional.avg_pool2d(yuvimage[..., 1:3, :, :], (1, 2)))
 
 
 def yuv_to_rgb(image: torch.Tensor) -> torch.Tensor:
@@ -149,7 +145,7 @@ def yuv420_to_rgb(imagey: torch.Tensor, imageuv: torch.Tensor) -> torch.Tensor:
 
     Args:
         imagey: Y (luma) Image plane to be converted to RGB with shape :math:`(*, 1, H, W)`.
-        imageuv: UV (luma) Image planes to be converted to RGB with shape :math:`(*, 2, H/2, W/2)`.
+        imageuv: UV (chroma) Image planes to be converted to RGB with shape :math:`(*, 2, H/2, W/2)`.
 
     Returns:
         RGB version of the image with shape :math:`(*, 3, H, W)`.


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
@@ -107,9 +107,62 @@ def test_exception(self, device, dtype):
             img = torch.ones(3, 1, 2, device=device, dtype=dtype)
             assert kornia.color.rgb_to_yuv420(img)
 
-    # TODO: investigate and implement me
-    # def test_unit(self, device, dtype):
-    #    pass
+    # Test max/min values. This is essentially testing the transform rather than the subsampling
+    # ref values manually checked vs rec 601
+    def test_unit_white(self, device, dtype):  # skipcq: PYL-R0201
+        rgb = torch.tensor([[[255, 255], [255, 255]], [[255, 255], [255, 255]], [[255, 255], [255, 255]]],
+                           device=device, dtype=torch.uint8).type(dtype) / 255.0
+        refy = torch.tensor([[[255, 255], [255, 255]]], device=device, dtype=torch.uint8)
+        refuv = torch.tensor([[[0]], [[0]]], device=device, dtype=torch.int8)
+
+        resy = (kornia.color.rgb_to_yuv420(rgb)[0] * 255.0).round().type(torch.uint8)
+        resuv = (kornia.color.rgb_to_yuv420(rgb)[1] * 255.0).round().clamp(-128, 127).type(torch.int8)
+        assert_close(refy, resy)
+        assert_close(refuv, resuv)
+
+    def test_unit_black(self, device, dtype):  # skipcq: PYL-R0201
+        rgb = torch.tensor([[[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]],
+                           device=device, dtype=torch.uint8).type(dtype) / 255.0
+        refy = torch.tensor([[[0, 0], [0, 0]]], device=device, dtype=torch.uint8)
+        refuv = torch.tensor([[[0]], [[0]]], device=device, dtype=torch.int8)
+
+        resy = (kornia.color.rgb_to_yuv420(rgb)[0] * 255.0).round().type(torch.uint8)
+        resuv = (kornia.color.rgb_to_yuv420(rgb)[1] * 255.0).round().clamp(-128, 127).type(torch.int8)
+        assert_close(refy, resy)
+        assert_close(refuv, resuv)
+
+    def test_unit_gray(self, device, dtype):  # skipcq: PYL-R0201
+        rgb = torch.tensor([[[127, 127], [127, 127]], [[127, 127], [127, 127]], [[127, 127], [127, 127]]],
+                           device=device, dtype=torch.uint8).type(dtype) / 255.0
+        refy = torch.tensor([[[127, 127], [127, 127]]], device=device, dtype=torch.uint8)
+        refuv = torch.tensor([[[0]], [[0]]], device=device, dtype=torch.int8)
+
+        resy = (kornia.color.rgb_to_yuv420(rgb)[0] * 255.0).round().type(torch.uint8)
+        resuv = (kornia.color.rgb_to_yuv420(rgb)[1] * 255.0).round().clamp(-128, 127).type(torch.int8)
+        assert_close(refy, resy)
+        assert_close(refuv, resuv)
+
+    def test_unit_red(self, device, dtype):  # skipcq: PYL-R0201
+        rgb = torch.tensor([[[255, 255], [255, 255]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]],
+                           device=device, dtype=torch.uint8).type(dtype) / 255.0
+        refy = torch.tensor([[[76, 76], [76, 76]]], device=device, dtype=torch.uint8)
+        refuv = torch.tensor([[[-37]], [[127]]], device=device, dtype=torch.int8)
+
+        resy = (kornia.color.rgb_to_yuv420(rgb)[0] * 255.0).round().type(torch.uint8)
+        resuv = (kornia.color.rgb_to_yuv420(rgb)[1] * 255.0).round().clamp(-128, 127).type(torch.int8)
+        assert_close(refy, resy)
+        assert_close(refuv, resuv)
+
+    def test_unit_blue(self, device, dtype):  # skipcq: PYL-R0201
+        rgb = torch.tensor([[[0, 0], [0, 0]], [[0, 0], [0, 0]], [[255, 255], [255, 255]]],
+                           device=device, dtype=torch.uint8).type(dtype) / 255.0
+        refy = torch.tensor([[[29, 29], [29, 29]]], device=device, dtype=torch.uint8)
+        refuv = torch.tensor([[[111]], [[-25]]], device=device, dtype=torch.int8)
+
+        resy = (kornia.color.rgb_to_yuv420(rgb)[0] * 255.0).type(torch.uint8)
+        resuv = (kornia.color.rgb_to_yuv420(rgb)[1] * 255.0).clamp(-128, 127).type(torch.int8)
+        assert_close(refy, resy)
+        assert_close(refuv, resuv)
 
     # This measures accuracy, given the impact of the subsampling we will avoid the issue by
     # repeating a 2x2 pattern for which mean will match what we get from upscaling again
@@ -182,10 +235,6 @@ def test_exception(self, device, dtype):
             img = torch.ones(3, 2, 1, device=device, dtype=dtype)
             assert kornia.color.rgb_to_yuv422(img)
 
-    # TODO: investigate and implement me
-    # def test_unit(self, device, dtype):
-    #    pass
-
     # This measures accuracy, given the impact of the subsampling we will avoid the issue by
     # repeating a 2x2 pattern for which mena will match what we get from upscaling again
     def test_forth_and_back(self, device, dtype):  # skipcq: PYL-R0201
@@ -327,9 +376,25 @@ def test_exception(self, device, dtype):
             imguv = torch.ones(3, 0, 1, device=device, dtype=dtype)
             assert kornia.color.yuv420_to_rgb(imgy, imguv)
 
-    # TODO: investigate and implement me
-    # def test_unit(self, device, dtype):
-    #    pass
+    # Test max/min values. This is essentially testing the transform rather than the subsampling
+    # ref values manually checked vs rec 601
+    def test_unit_white(self, device, dtype):  # skipcq: PYL-R0201
+        refrgb = torch.tensor([[[255, 255], [255, 255]], [[255, 255], [255, 255]], [[255, 255], [255, 255]]],
+                              device=device, dtype=torch.uint8)
+        y = torch.tensor([[[255, 255], [255, 255]]], device=device, dtype=torch.uint8).type(dtype) / 255.0
+        uv = torch.tensor([[[0]], [[0]]], device=device, dtype=torch.int8).type(torch.float) / 255.0
+
+        resrgb = (kornia.color.yuv420_to_rgb(y, uv) * 255.0).round().type(torch.uint8)
+        assert_close(refrgb, resrgb)
+
+    def test_unit_red(self, device, dtype):  # skipcq: PYL-R0201
+        refrgb = torch.tensor([[[221, 221], [221, 221]], [[17, 17], [17, 17]], [[1, 1], [1, 1]]],
+                              device=device, dtype=torch.uint8)
+        y = torch.tensor([[[76, 76], [76, 76]]], device=device, dtype=torch.uint8).type(dtype) / 255.0
+        uv = torch.tensor([[[-37]], [[127]]], device=device, dtype=torch.int8).type(torch.float) / 255.0
+
+        resrgb = (kornia.color.yuv420_to_rgb(y, uv) * 255.0).round().type(torch.uint8)
+        assert_close(refrgb, resrgb)
 
     # TODO: improve accuracy
     def test_forth_and_back(self, device, dtype):  # skipcq: PYL-R0201