here  is my custom callback to show some images

class DisplaySamples(tf.keras.callbacks.Callback):
    def __init__(self, log_dir='./logs/tmp/', get_samples=None):
        super(DisplaySamples, self).__init__()
        self.step = 0
        self.samples = get_samples
        self.writer = tf.summary.FileWriter(log_dir)

    def on_batch_end(self, batch, logs=None):
        logs = logs or {}
        self.step += 1
        if self.step % 2 == 0:
            # images using TensorBorad;
            summary_str = []
            for i in range(len(self.samples)):
                image = self.samples[i]
                keypoints = np.squeeze(self.model.predict(image)[0], axis=0)

                summary_str.append(tf.Summary.Value(tag='plot/image/{}'.format(i),
                                                    image=tf.summary.image("image", image)))
                summary_str.append(tf.Summary.Value(tag='plot/keypts/{}'.format(i),
                                                    image=tf.summary.image("keyps", keypoints)))

            self.writer.add_summary(tf.Summary(value=summary_str), global_step=self.step)


def read_images():
    """

    Returns
    -------

    """

    image_list = []
    for filename in glob.glob('../samples/*.jpg'):
        im = cv2.imread(filename)
        im = cv2.resize(im, (512, 512, 3))
        im = np.reshape(im, (1, 512, 512, 3))
        image_list.append(im)
    return image_list