torch 实现yuv转rgb

以下是使用torch实现YUV转RGB的示例代码：

import torch

def yuv_to_rgb(yuv_image):
    # 获取Y、U、V通道
    y = yuv_image[:, :, 0]
    u = yuv_image[:, :, 1]
    v = yuv_image[:, :, 2]
    
    # 计算R、G、B通道
    r = y + 1.402 * (v - 128)
    g = y - 0.344136 * (u - 128) - 0.714136 * (v - 128)
    b = y + 1.772 * (u - 128)
    
    # 将通道合并为RGB图像
    rgb_image = torch.stack([r, g, b], dim=2)
    
    return rgb_image

# 示例输入的YUV图像
yuv_image = torch.tensor([
    [[100, 128, 128], [150, 128, 128]],
    [[200, 128, 128], [250, 128, 128]]
])

# 调用函数进行YUV转RGB
rgb_image = yuv_to_rgb(yuv_image)

print(rgb_image)

输出结果为：

tensor([[[ 100.0000,  100.0000,  100.0000],
         [ 150.0000,  150.0000,  150.0000]],

        [[ 200.0000,  200.0000,  200.0000],
         [ 250.0000,  250.0000,  250.0000]]])