Working with webp and jpeg images have different number of channels

Question

I'm working with computer vision project, where my images are combination of webp and jpeg. I'm using tensorflow '2.3.2'
You can think my directories like this :

IMAGES
 |-img1.jpeg
 |-img2.webp

For reading webp, I use tfio.image.decode_webp and when reading jpeg, I use tf.image.decode_jpeg(img, channels=3). Here's the code :

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)

The problem is, webp are converted to 4 channel(RGBA) tensor where decode jpeg is in 3 channel(RGB). This creates inconsistencies within my dataset since the model only except 3 channel images.

One solution I can think of is converting all my webp to jpeg through this. But is there any better solution for this? like converting the 4 channel into 3 channel in TensorFlow or reading webp as 3 channel in TensorFlow or anything else where I can just put the solution inside my python script?

Answer 1

If you want to train a single model with both *.jpeg and *.webp images then you should create the same input layer layout for both.

To do so, you basically need to convert either RGB to RGBA or (what I would do) RGBA to RGB. If you want to simply drop the alpha-channel you can use tensorflow's rgba_to_rgb (as @Lescurel pointed out in the comments).

But the RBGA to RGB conversion with alpha compositing is not very complex and you can do the operation directly on the tensor your get from calling load.

Here is a tensorflow adaptation of a numpy implementation proposed for this SO question:

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)

Using this avoids having to call .numpy() on the tensor and the potential issues related to that.

So basically rgba2rgb(load(image)) should then do the trick.