Source code for kornia.color.ycbcr

import torch
import torch.nn as nn


[docs]class YcbcrToRgb(nn.Module): r"""Convert image from YCbCr to Rgb The image data is assumed to be in the range of (0, 1). args: image (torch.Tensor): YCbCr 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:: >>> import torch >>> import kornia >>> input = torch.rand(2, 3, 4, 5) >>> rgb = kornia.color.YcbcrToRgb() >>> output = rgb(input) # 2x3x4x5 """ def __init__(self) -> None: super(YcbcrToRgb, self).__init__() def forward(self, image: torch.Tensor) -> torch.Tensor: # type: ignore return ycbcr_to_rgb(image)
[docs]def ycbcr_to_rgb(image: torch.Tensor) -> torch.Tensor: r"""Convert an YCbCr image to RGB The image data is assumed to be in the range of (0, 1). Args: image (torch.Tensor): YCbCr Image to be converted to RGB. Returns: torch.Tensor: RGB version of the image. """ if not torch.is_tensor(image): raise TypeError("Input type is not a torch.Tensor. Got {}".format( type(image))) if len(image.shape) < 3 or image.shape[-3] != 3: raise ValueError("Input size must have a shape of (*, 3, H, W). Got {}" .format(image.shape)) y: torch.Tensor = image[..., 0, :, :] cb: torch.Tensor = image[..., 1, :, :] cr: torch.Tensor = image[..., 2, :, :] delta: float = .5 cb_shifted: torch.Tensor = cb - delta cr_shifted: torch.Tensor = cr - delta r: torch.Tensor = y + 1.403 * cr_shifted g: torch.Tensor = y - .714 * cr_shifted - .344 * cb_shifted b: torch.Tensor = y + 1.773 * cb_shifted return torch.stack((r, g, b), -3)
[docs]class RgbToYcbcr(nn.Module): r"""Convert image from RGB to YCbCr The image data is assumed to be in the range of (0, 1). args: image (torch.Tensor): RGB image to be converted to YCbCr. returns: torch.tensor: YCbCr version of the image. shape: - image: :math:`(*, 3, H, W)` - output: :math:`(*, 3, H, W)` Examples:: >>> import torch >>> import kornia >>> input = torch.rand(2, 3, 4, 5) >>> ycbcr = kornia.color.RgbToYcbcr() >>> output = ycbcr(input) # 2x3x4x5 """ def __init__(self) -> None: super(RgbToYcbcr, self).__init__() def forward(self, image: torch.Tensor) -> torch.Tensor: # type: ignore return rgb_to_ycbcr(image)
[docs]def rgb_to_ycbcr(image: torch.Tensor) -> torch.Tensor: r"""Convert an RGB image to YCbCr. Args: image (torch.Tensor): RGB Image to be converted to YCbCr. Returns: torch.Tensor: YCbCr version of the image. """ if not torch.is_tensor(image): raise TypeError("Input type is not a torch.Tensor. Got {}".format( type(image))) if len(image.shape) < 3 or image.shape[-3] != 3: raise ValueError("Input size must have a shape of (*, 3, H, W). Got {}" .format(image.shape)) r: torch.Tensor = image[..., 0, :, :] g: torch.Tensor = image[..., 1, :, :] b: torch.Tensor = image[..., 2, :, :] delta = .5 y: torch.Tensor = .299 * r + .587 * g + .114 * b cb: torch.Tensor = (b - y) * .564 + delta cr: torch.Tensor = (r - y) * .713 + delta return torch.stack((y, cb, cr), -3)