How to configure a tf.data.Dataset for variable size images?

Viewed 1091

I'm setting up a image data pipeline on Tensorflow 2.1. I'm using a dataset with RGB images of variable shapes (h, w, 3) and I can't find a way to make it work. I get the following error when I call tf.data.Dataset.batch() :

tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot batch tensors with different shapes in component 0. First element had shape [256,384,3] and element 3 had shape [160,240,3]

I found the padded_batch method but I don't want my images to be padded to the same shape.

EDIT:

I think that I found a little workaround to this by using the function tf.data.experimental.dense_to_ragged_batch (which convert the dense tensor representation to a ragged one).

Unlike tf.data.Dataset.batch, the input elements to be batched may have different shapes, and each batch will be encoded as a tf.RaggedTensor

But then I have another problem. My dataset contains images and their corresponding labels. When I use the function like this:

ds = ds.map(
      lambda x: tf.data.experimental.dense_to_ragged_batch(batch_size)
  ) 

I get the following error because it tries to map the function to the entire dataset (thus to images and labels), which is not possible because it can only be applied to a 1 single tensor (not 2).

TypeError: <lambda>() takes 1 positional argument but 2 were given

Is there a way to specify which element of the two I want the transformation to be applied to ?

2 Answers

I just hit the same problem. The solution turned out to be loading the data as 2 datasets and then using dataet.zip() to merge them.

images = dataset.map(parse_images, num_parallel_calls=tf.data.experimental.AUTOTUNE)
images = dataset_images.apply(
    tf.data.experimental.dense_to_ragged_batch(batch_size=batch_size, drop_remainder=True))

dataset_total_cost = dataset.map(get_total_cost)
dataset_total_cost = dataset_total_cost.batch(batch_size, drop_remainder=True)

dataset = dataset.zip((dataset_images, dataset_total_cost))

If you do not want to resize your images, you can only use a batch size of 1 and not bigger than that. Thus you can train your model one image at at time. The error you reported clearly says that you are using a batch size bigger than 1 and trying to put two images of different shape/size in a batch. You could either resize your images to a fixed shape (or pad your images), or use batch size of 1 as follows:

my_data = tf.data.Dataset(....) # with whatever arguments you use here
my_data = my_data.batch(1)
Related