Tensorflow custom preprocessing with tf.py_function losing shape

user_9209390 0
Viewed 639

Im writing a model and doing the preprocessing part: I have a method which preprocesses my tensorflow dataset by calling:

ds = ds.map(process_path, num_parallel_calls=AUTOTUNE)

I followed the tensorflow documentation and got this code for process_path:

def process_path(filename):
  label = get_label(filename)

  image = tf.io.read_file(filename)
  image = tf.image.decode_jpeg(image, channels=3)
  image = tf.image.rgb_to_grayscale(image)
  image = tf.image.convert_image_dtype(image, tf.float32)
  image = tf.image.resize(image, [224, 224])
  
  return image, label

Then I want to add my own preprocessing, such as rotating the image so I created a rotate method wrapped with py_function as the documentation suggests:

def rotate_image(image):
  return tfa.image.rotate(image, random.randrange(-5, 5)/1.0)

def tf_rotate_image(image, label):
  [image,] = tf.py_function(rotate_image, [image], [tf.float32])
  return image, label

However when I add this to my process_path the model seems to break and freezes... I added print statements with image.shape after each adjustment and it shows that after the rotate method the image shape becomes <unknown> so I believe this to be the error:

def process_path(filename):
  label = get_label(filename)

  image = tf.io.read_file(filename)
  print(image.shape)
  image = tf.image.decode_jpeg(image, channels=3)
  print(image.shape)
  image = tf.image.rgb_to_grayscale(image)
  print(image.shape)
  image = tf.image.convert_image_dtype(image, tf.float32)
  print(image.shape)
  image = tf.image.resize(image, [224, 224])
  print(image.shape)

  image, label = tf_rotate_image(image, label)
  print(image.shape)
  
  return image, label

Output:

()
(None, None, 3)
(None, None, 1)
(None, None, 1)
(224, 224, 1)
<unknown>

Any help is greatly appreciated.

tensorflow2.0 tensorflow machine-learning python
1 Answers

You are losing your shape because of the call to tf.py_function.

Using a function based on tensorflow ops

In that case, because you are using tfa, you can simply call the function without wrapping it in a tf.py_function as you are using tf ops :

def process_path(filename):
  label = 1

  image = tf.io.read_file(filename)
  print(image.shape)
  image = tf.image.decode_jpeg(image, channels=3)
  print(image.shape)
  image = tf.image.rgb_to_grayscale(image)
  print(image.shape)
  image = tf.image.convert_image_dtype(image, tf.float32)
  print(image.shape)
  image = tf.image.resize(image, [224, 224])
  print(image.shape)

  image, label = (lambda x,y : (rotate_image(x), y))(image, label)
  print(image.shape)

will result in :

>>> ds = ds.map(process_path)
()
(None, None, 3)
(None, None, 1)
(None, None, 1)
(224, 224, 1)
(224, 224, 1)

Using a function that does not consist of tensorflow ops

If you want to use a function that does not consist on tensorflow ops, then you can use tf.py_function, and set the shape explicitly. This what is done in the guide tf.data: Build TensorFlow input pipelines. Taking the example from that guide :

def tf_random_rotate_image(image, label):
  im_shape = image.shape
  [image,] = tf.py_function(random_rotate_image, [image], [tf.float32])
  # the shape is set explicitly because tensorflow can not ensure
  # that the shape is not modified during the execution of the function
  image.set_shape(im_shape)
  return image, label

However, when doing that, one assumption that is made by tensorflow is that the shape you are setting is actually correct! The following example will crash, because the function lambda x:1 does not preserve the shape of the input.

def not_shape_preserving(image, label):
    im_shape = image.shape
    # this function does not preserve the shape
    [image,] = tf.py_function(lambda x: 1., [image], [tf.float32])
    image.set_shape(im_shape)
    return image, label

Creating the dataset will work, because tensorflow trusts you. However, when trying to use it, you will be met with an error akin to this :

Incompatible shapes at component 0: expected [224,224,1] but got [].
user_73701530
Related