如何在Tensorflow中将图像分割为RGB通道

Question

我想将RGB图像分成3个通道，对它们执行一些操作（转换为YUV，更准确），然后再次连接它们。这是我到目前为止所做的（这段代码不起作用，因为我得到尺寸错误）：

    def rgb_to_yuv(rgb):
        with tf.name_scope("rgb_to_yuv"):
            rgb = check_image(rgb) #check the number of channels
            rgb_pixels = tf.reshape(rgb, [-1, 3]) #check if the image is srgb
            r, g, b= tf.split(3, 3, rgb_pixels)
            Y = 0.299 * r + 0.587 * g + 0.114 * b
            U = 0.492 * (b - Y)
            V = 0.877 * (r - Y)
            yuv_pixels = tf.concat([Y, U, V], 1)
            return tf.reshape(yuv_pixels, tf.shape(rgb)) 

这是check_image（来自pix2pix implementation）：

def check_image(image):
    assertion = tf.assert_equal(tf.shape(image)[-1], 3, message="image must have 3 color channels")
    with tf.control_dependencies([assertion]):
        image = tf.identity(image)

    if image.get_shape().ndims not in (3, 4):
        raise ValueError("image must be either 3 or 4 dimensions")

    # make the last dimension 3 so that you can unstack the colors
    shape = list(image.get_shape())
    shape[-1] = 3
    image.set_shape(shape)
    return image