Source code for kornia.color.yuv


import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple


[docs]class RgbToYuv(nn.Module): r"""Convert image from RGB to YUV The image data is assumed to be in the range of (0, 1). args: image (torch.Tensor): RGB image to be converted to YUV. returns: torch.tensor: YUV version of the image. shape: - image: :math:`(*, 3, H, W)` - output: :math:`(*, 3, H, W)` Examples:: >>> input = torch.rand(2, 3, 4, 5) >>> yuv = kornia.color.RgbToYuv() >>> output = yuv(input) # 2x3x4x5 Reference:: [1] https://es.wikipedia.org/wiki/YUV#RGB_a_Y'UV """ def __init__(self) -> None: super(RgbToYuv, self).__init__() def forward( # type: ignore self, input: torch.Tensor) -> torch.Tensor: return rgb_to_yuv(input)
[docs]def rgb_to_yuv(input: torch.Tensor) -> torch.Tensor: r"""Convert an RGB image to YUV The image data is assumed to be in the range of (0, 1). Args: input (torch.Tensor): RGB Image to be converted to YUV. Returns: torch.Tensor: YUV version of the image. See :class:`~kornia.color.RgbToYuv` for details.""" if not torch.is_tensor(input): raise TypeError("Input type is not a torch.Tensor. Got {type(input)}") if not(len(input.shape) == 3 or len(input.shape) == 4): raise ValueError(f"Input size must have a shape of (*, 3, H, W) or (3, H, W). Got {input.shape}") if input.shape[-3] != 3: raise ValueError(f"Expected input to have 3 channels, got {input.shape[-3]}") r, g, b = torch.chunk(input, chunks=3, dim=-3) y: torch.Tensor = 0.299 * r + 0.587 * g + 0.114 * b u: torch.Tensor = -0.147 * r - 0.289 * g + 0.436 * b v: torch.Tensor = 0.615 * r - 0.515 * g - 0.100 * b yuv_img: torch.Tensor = torch.cat((y, u, v), -3) return yuv_img
[docs]class YuvToRgb(nn.Module): r"""Convert image from YUV to RGB The image data is assumed to be in the range of (0, 1). args: image (torch.Tensor): YUV image to be converted to RGB. returns: torch.tensor: RGB version of the image. shape: - image: :math:`(*, 3, H, W)` - output: :math:`(*, 3, H, W)` Examples:: >>> input = torch.rand(2, 3, 4, 5) >>> rgb = kornia.color.YuvToRgb() >>> output = rgb(input) # 2x3x4x5 """ def __init__(self) -> None: super(YuvToRgb, self).__init__() def forward( # type: ignore self, input: torch.Tensor) -> torch.Tensor: return yuv_to_rgb(input)
[docs]def yuv_to_rgb(input: torch.Tensor) -> torch.Tensor: r"""Convert an YUV image to RGB The image data is assumed to be in the range of (0, 1). Args: input (torch.Tensor): YUV Image to be converted to RGB. Returns: torch.Tensor: RGB version of the image. See :class:`~kornia.color.YuvToRgb` for details.""" if not torch.is_tensor(input): raise TypeError("Input type is not a torch.Tensor. Got {type(input)}") if not(len(input.shape) == 3 or len(input.shape) == 4): raise ValueError(f"Input size must have a shape of (*, 3, H, W) or (3, H, W). Got {input.shape}") if input.shape[-3] != 3: raise ValueError(f"Expected input to have 3 channels, got {input.shape[-3]}") y, u, v = torch.chunk(input, chunks=3, dim=-3) r: torch.Tensor = y + 1.14 * v # coefficient for g is 0 g: torch.Tensor = y + -0.396 * u - 0.581 * v b: torch.Tensor = y + 2.029 * u # coefficient for b is 0 rgb_img: torch.Tensor = torch.cat((r, g, b), -3) return rgb_img