import torch


def rgb2yuv(rgb, device):
    """
    Convert RGB image into YUV https://en.wikipedia.org/wiki/YUV
    """
    rgb = ((rgb + 1) / 2) * 255.0
    rgb_ = rgb.transpose(1, 3)  # input is 3*n*n   default
    A = torch.tensor([[0.299, -0.14714119, 0.61497538],
                      [0.587, -0.28886916, -0.51496512],
                      [0.114, 0.43601035, -0.10001026]]).to(device)  # from  Wikipedia
    yuv = torch.tensordot(rgb_, A, 1).transpose(1, 3)
    yuv = yuv / 255.0
    return yuv