Image summaries with TensorFlow's Dataset API

Viewed 716

I am using TensorFlow's Dataset API to load and preprocess images. I would like to add summaries of my preprocessed images to Tensorboard.

What is the recommended way to do this?

So far, my code looks like this:

def get_data():
  dataset = FixedLengthRecordDataset(...)
  dataset = dataset.map(dataset_parser, ...)
  if is_training:
    dataset = dataset.map(preprocess_for_train, ...)
  # Do shuffling, batching...
  return dataset

def preprocess_for_train(image, label):
  # Do preprocessing...
  image = tf.image.random_flip_left_right(image)
  # Add summary
  tf.summary.image('preprocessed_image', tf.expand_dims(image, 0))
  return image, label

Within preprocess_for_train my image is listed in the SUMMARIES collection, but it is no longer part of the graph when returning to the outer function. I think this happens because map uses a different thread and therefore refers to a different instance of tf.Graph.

Since this is not working, what other options do I have to display my images in Tensorboard?

1 Answers

I found a hack to add the preprocessed images to Tensorboard, using the output of my iterator:

train_dataset = get_data()
iterator = Iterator.from_structure(train_dataset.output_types, train_dataset.output_shapes)
# Batch consists of [image, label]
next_batch = iterator.get_next()
tf.summary.image('preprocessed_image', next_batch[0])

However, this will call next_batch a second time when running the summary_op. It helps for just debugging the image preprocessing, but it is not a solution for real training. Furthermore, only the fully preprocessed images can be observed, not intermediate preprocessing stages.

Related