How to display custom images in TensorBoard using Keras?

Question

I\'m working on a segmentation problem in Keras and I want to display segmentation results at the end of every training epoch.

I want something similar to Tensorflow

Answer 1

I believe I found a better way to log such custom images to tensorboard making use of the tf-matplotlib. Here is how...

class TensorBoardDTW(tf.keras.callbacks.TensorBoard):
    def __init__(self, **kwargs):
        super(TensorBoardDTW, self).__init__(**kwargs)
        self.dtw_image_summary = None

    def _make_histogram_ops(self, model):
        super(TensorBoardDTW, self)._make_histogram_ops(model)
        tf.summary.image('dtw-cost', create_dtw_image(model.output))

One just need to overwrite the _make_histogram_ops method from the TensorBoard callback class to add the custom summary. In my case, the create_dtw_image is a function that creates an image using the tf-matplotlib.

Regards,.

Answer 2

class customModelCheckpoint(Callback):
def __init__(self, log_dir='../logs/', feed_inputs_display=None):
      super(customModelCheckpoint, self).__init__()
      self.seen = 0
      self.feed_inputs_display = feed_inputs_display
      self.writer = tf.summary.FileWriter(log_dir)


def custom_set_feed_input_to_display(self, feed_inputs_display):
      self.feed_inputs_display = feed_inputs_display


# A callback has access to its associated model through the class property self.model.
def on_batch_end(self, batch, logs = None):
      logs = logs or {}
      self.seen += 1
      if self.seen % 8 == 0: # every 200 iterations or batches, plot the costumed images using TensorBorad;
          summary_str = []
          feature = self.feed_inputs_display[0][0]
          disp_gt = self.feed_inputs_display[0][1]
          disp_pred = self.model.predict_on_batch(feature)

          summary_str.append(tf.summary.image('disp_input/{}'.format(self.seen), feature, max_outputs=4))
          summary_str.append(tf.summary.image('disp_gt/{}'.format(self.seen), disp_gt, max_outputs=4))
          summary_str.append(tf.summary.image('disp_pred/{}'.format(self.seen), disp_pred, max_outputs=4))

          summary_st = tf.summary.merge(summary_str)
          summary_s = K.get_session().run(summary_st)
          self.writer.add_summary(summary_s, global_step=self.seen)
          self.writer.flush()

Then you can call your custom callback and write the image in tensorboard

callback_mc = customModelCheckpoint(log_dir='../logs/',  feed_inputs_display=[(a, b)])
callback_tb = TensorBoard(log_dir='../logs/', histogram_freq=0, write_graph=True, write_images=True)
callback = []
def data_gen(fr1, fr2):
while True:
    hdr_arr = []
    ldr_arr = []
    for i in range(args['batch_size']):
        try:
            ldr = pickle.load(fr2)           
            hdr = pickle.load(fr1)               
        except EOFError:
            fr1 = open(args['data_h_hdr'], 'rb')
            fr2 = open(args['data_h_ldr'], 'rb')
        hdr_arr.append(hdr)
        ldr_arr.append(ldr)
    hdr_h = np.array(hdr_arr)
    ldr_h = np.array(ldr_arr)
    gen = aug.flow(hdr_h, ldr_h, batch_size=args['batch_size'])
    out = gen.next()
    a = out[0]
    b = out[1]
    callback_mc.custom_set_feed_input_to_display(feed_inputs_display=[(a, b)])
    yield [a, b]

callback.append(callback_tb)
callback.append(callback_mc)
H = model.fit_generator(data_gen(fr1, fr2), steps_per_epoch=100,   epochs=args['epoch'], callbacks=callback)

picture