Tensorflow_io: ValueError: Cannot infer argument `num` from shape (None, None, None)

I am trying to read and decode tiff images in tensorflow. I am using tensrflow_io package as follows, I am getting this error that I cant figure out.

import tensorflow as tf
import tensorflow_io as tfio
import os

def process_image(image):

  image = tf.io.read_file(image)
  image = tfio.experimental.image.decode_tiff(image)
  image = tfio.experimental.color.rgba_to_rgb(image)
  return image

path = os.path.join(os.curdir, '*.TIF')
files = tf.data.Dataset.list_files(path)

Output:

for file in files.take(5):
  print(file)

tf.Tensor(b'./SIMCEPImages_A01_C1_F1_s10_w1.TIF', shape=(), dtype=string)
tf.Tensor(b'./SIMCEPImages_A01_C1_F1_s04_w1.TIF', shape=(), dtype=string)
tf.Tensor(b'./SIMCEPImages_A01_C1_F1_s12_w1.TIF', shape=(), dtype=string)
tf.Tensor(b'./SIMCEPImages_A01_C1_F1_s04_w2.TIF', shape=(), dtype=string)
tf.Tensor(b'./SIMCEPImages_A01_C1_F1_s11_w1.TIF', shape=(), dtype=string)

Now if I call:

dataset = files.map(process_image, num_parallel_calls=tf.data.experimental.AUTOTUNE)

for img in dataset.take(5):
  print(img.shape)

ValueError: in user code:

    File "<ipython-input-4-1d2deab36c6d>", line 5, in process_image  *
        image = tfio.experimental.color.rgba_to_rgb(image)
    File "/usr/local/lib/python3.7/dist-packages/tensorflow_io/python/experimental/color_ops.py", line 80, in rgba_to_rgb  *
        rgba = tf.unstack(input, axis=-1)

    ValueError: Cannot infer argument `num` from shape (None, None, None)

>Solution :

The problem is that tfio.experimental.color.rgba_to_rgb uses unstack under the hood, which cannot work in graph mode. One solution would be to manually index the channels you want according to the source code for rgba_to_rgb. Here is a working example:

import numpy as np
from PIL import Image
import tensorflow as tf
import tensorflow_io as tfio
import os

# Create dummy data
data = np.random.randint(0, 255, (10,10)).astype(np.uint8)
im = Image.fromarray(data)
im.save('image1.tif')
im.save('image2.tif')

def process_image(image):

  image = tf.io.read_file(image)
  image = tfio.experimental.image.decode_tiff(image)
  r, g, b = image[:, :, 0], image[:, :, 1], image[:, :, 2]
  return tf.stack([r, g, b], axis=-1)

path = os.path.join(os.curdir, '*.tif')
files = tf.data.Dataset.list_files(path)

for file in files.take(5):
  print(file)

dataset = files.map(process_image, num_parallel_calls=tf.data.experimental.AUTOTUNE)
for img in dataset.take(5):
  print(img.shape)

tf.Tensor(b'./image2.tif', shape=(), dtype=string)
tf.Tensor(b'./image1.tif', shape=(), dtype=string)
(10, 10, 3)
(10, 10, 3)

If you really want to use tfio.experimental.color.rgba_to_rgb, it will have be out of graph mode, using for example tf.py_function.