使用 webp 和 jpeg 图像具有不同数量的通道

Question

我正在处理计算机视觉项目，我的图像是 webp 和 jpeg 的组合。我正在使用 tensorflow '2.3.2'
你可以这样想我的目录：

IMAGES
 |-img1.jpeg
 |-img2.webp

对于阅读 webp，我使用 tfio.image.decode_webp，在阅读 jpeg 时，我使用 tf.image.decode_jpeg(img, channels=3)。这是代码：

def load(file_path):
    img = tf.io.read_file(file_path)
    extension = tf.strings.split(file_path,sep=".")    
    if extension[-1] == "webp" :
        img = tfio.image.decode_webp(img)
    else :
        img = tf.image.decode_jpeg(img, channels=3)
    #img preprocess here
    return img

def create_dataset(df,batch_size):
    image = df["image_path"]
    # I'm working on MultiTaskLearning so I have multiple targets
    target1 = df["target1"].to_numpy()
    target2 = df["target2"].to_numpy()
    
    ds = tf.data.Dataset.from_tensor_slices((image,target1,target2))
    ds = ds.map(lambda image, target1,target2: (load(image),  {"target1":target1, "target2":target2}), num_parallel_calls=tf.data.experimental.AUTOTUNE)
    
    ds = ds.batch(batch_size)
    ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
    return ds

dataset = create_dataset(df,100)

问题是，webp 被转换为 4 通道（RGBA）张量，其中解码 jpeg 位于 3 通道（RGB）中。这会在我的数据集中造成不一致，因为该模型仅包含 3 通道图像。

我能想到的一个解决方案是通过this 将我所有的 webp 转换为 jpeg。但是有没有更好的解决方案呢？比如在 TensorFlow 中将 4 通道转换为 3 通道，或者在 TensorFlow 中将 webp 读取为 3 通道，或者我可以将解决方案放入我的 python 脚本中的其他任何东西？

Answer 1

如果您想用 *.jpeg 和 *.webp 图像训练单个模型，那么您应该为两者创建相同的输入层布局。

为此，您基本上需要将 RGB 转换为 RGBA 或（我会做的）RGBA 到 RGB。如果您想简单地删除 alpha 通道，您可以使用 tensorflow 的 rgba_to_rgb（如 @Lescurel pointed out in the comments）。

但是alpha compositing 的 RBGA 到 RGB 转换不是很复杂，你可以直接在调用load 得到的张量上进行操作。

这是针对this SO 问题提出的numpy implementation 的张量流改编：

def rgba2rgb(rgba, background=(255,255,255)):
    row, col, ch = tf.shape(rgba)
    if ch == 3:
        return rgba
    assert ch == 4, 'RGBA image has 4 channels.'
    r, g, b, a = tf.unstack(tf.cast(rgba, tf.float32), axis=-1)
    a =  tf.cast(a, tf.float32) / 255.0
    R, G, B = background
    r = r * a + (1.0 - a) * R
    g = g * a + (1.0 - a) * G
    b = b * a + (1.0 - a) * B
    rgb = tf.stack([r,g,b], axis=-1)
    return tf.cast(rgb, tf.uint8)

使用它可以避免在张量上调用 .numpy() 和 potential issues related to that。

所以基本上rgba2rgb(load(image)) 应该可以解决问题。