How to add an image to summary during evaluation when using Estimator?

I run an evaluation at the end of each epoch and need to show an image calculated from the features and labels arguments of the model function model_fn. Including a tf.summary.image(name, image) in evaluation part of the model function does not help and it looks to me that the only way to do so is to pass the correct eval_metric_ops to construct the EstimatorSpec for mode EVAL. So I first sub-class Estimator so that it considers images. The following code is mostly from estimator.py; the only change is the few lines marked by "my change" inside _write_dict_to_summary:

import logging
import io
import numpy as np
import matplotlib.pyplot as plt
import six
from google.protobuf import message
import tensorflow as tf
from tensorflow.python.training import evaluation
from tensorflow.python import ops
from tensorflow.python.estimator.estimator import _dict_to_str, _write_checkpoint_path_to_summary
from tensorflow.core.framework import summary_pb2
from tensorflow.python.framework import tensor_util
from tensorflow.python.summary.writer import writer_cache


def dump_as_image(a):
    vmin = np.min(a)
    vmax = np.max(a)
    img = np.squeeze((img - vmin) / (vmax - vmin) * 255).astype(np.uint8)
    s = io.BytesIO()
    plt.imsave(s, img, format='png', vmin=0, vmax=255, cmap='gray')
    return s.getvalue()


# see https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/estimator/estimator.py

def _write_dict_to_summary(output_dir, dictionary, current_global_step):
    logging.info('Saving dict for global step %d: %s', current_global_step, _dict_to_str(dictionary))
    summary_writer = writer_cache.FileWriterCache.get(output_dir)
    summary_proto = summary_pb2.Summary()
    for key in dictionary:
        if dictionary[key] is None:
            continue
        if key == 'global_step':
            continue
        if (isinstance(dictionary[key], np.float32) or
            isinstance(dictionary[key], float)):
            summary_proto.value.add(tag=key, simple_value=float(dictionary[key]))
        elif (isinstance(dictionary[key], np.int64) or
              isinstance(dictionary[key], np.int32) or
              isinstance(dictionary[key], int)):
            summary_proto.value.add(tag=key, simple_value=int(dictionary[key]))
        elif isinstance(dictionary[key], six.binary_type):
            try:
                summ = summary_pb2.Summary.FromString(dictionary[key])
                for i, img_bytes in enumerate(summ.value):
                    summ.value[i].tag = '%s/%d' % (key, i)
                summary_proto.value.extend(summ.value)
            except message.DecodeError:
                logging.warn('Skipping summary for %s, cannot parse string to Summary.', key)
                continue
        elif isinstance(dictionary[key], np.ndarray):
            value = summary_proto.value.add()
            value.tag = key
            value.node_name = key
            array = dictionary[key]

            # my change begins
            if array.ndim == 2:
                buffer = dump_as_image(array)
                value.image.encoded_image_string = buffer
            # my change ends

            else:
                tensor_proto = tensor_util.make_tensor_proto(array)
                value.tensor.CopyFrom(tensor_proto)

                logging.info(
                    'Summary for np.ndarray is not visible in Tensorboard by default. '
                    'Consider using a Tensorboard plugin for visualization (see '
                    'https://github.com/tensorflow/tensorboard-plugin-example/blob/master/README.md'
                    ' for more information).')
        else:
            logging.warn(
                'Skipping summary for %s, must be a float, np.float32, np.int64, '
                'np.int32 or int or np.ndarray or a serialized string of Summary.',
                key)
    summary_writer.add_summary(summary_proto, current_global_step)
    summary_writer.flush()


class ImageMonitoringEstimator(tf.estimator.Estimator):

    def __init__(self, *args, **kwargs):
        tf.estimator.Estimator._assert_members_are_not_overridden = lambda self: None
        super(ImageMonitoringEstimator, self).__init__(*args, **kwargs)

    def _evaluate_run(self, checkpoint_path, scaffold, update_op, eval_dict, all_hooks, output_dir):

        eval_results = evaluation._evaluate_once(
            checkpoint_path=checkpoint_path,
            master=self._config.evaluation_master,
            scaffold=scaffold,
            eval_ops=update_op,
            final_ops=eval_dict,
            hooks=all_hooks,
            config=self._session_config)

        current_global_step = eval_results[ops.GraphKeys.GLOBAL_STEP]

        _write_dict_to_summary(
            output_dir=output_dir,
            dictionary=eval_results,
            current_global_step=current_global_step)

        if checkpoint_path:
            _write_checkpoint_path_to_summary(
                output_dir=output_dir,
                checkpoint_path=checkpoint_path,
                current_global_step=current_global_step)

        return eval_results

the model function is like --

def model_func(features, labels, mode):
    # calculate network_output
    if mode == tf.estimator.ModeKeys.TRAIN:
        # training
    elif mode == tf.estimator.ModeKeys.EVAL:

        # make_image consists of slicing and concatenations
        images = tf.map_fn(make_image, (features, network_output, labels), dtype=features.dtype)

        eval_metric_ops = images, tf.no_op() # not working
        return tf.estimator.EstimatorSpec(mode, loss=loss)
                                          eval_metric_ops={'images': eval_metric_ops})
    else:
        # prediction

And the main part --

# mon_features and mon_labels are np.ndarray
estimator = ImageMonitoringEstimator(model_fn=model_func,...)
mon_input_func = tf.estimator.inputs.numpy_input_fn(mon_features,
                                                    mon_labels, 
                                                    shuffle=False,
                                                    num_epochs=num_epochs,
                                                    batch_size=len(mon_features))

for _ in range(num_epochs):
    estimator.train(...)
    estimator.evaluate(input_fn=mon_input_func)

The code above will give a warning (later an error):

WARNING:tensorflow:An OutOfRangeError or StopIteration exception is raised by the code in FinalOpsHook. This typically means the Ops running by the FinalOpsHook have a dependency back to some input source, which should not happen. For example, for metrics in tf.estimator.Estimator, all metrics functions return two Ops: value_op and update_op. Estimator.evaluate calls the update_op for each batch of the data in input source and, once it is exhausted, it call the value_op to get the metric values. The value_op here should have dependency back to variables reading only, rather than reading another batch from input. Otherwise, the value_op, executed by FinalOpsHook, triggers another data reading, which ends OutOfRangeError/StopIteration. Please fix that.

Looks like I didn't set the eval_metric_ops correctly. I guess tf.map_fn touches another batch as the warning message hints; maybe I need some stacking operation as the update_op to build the images used for monitoring incrementally? But I am not sure how to do that. So how to add an image to summary during evaluation when using Estimator?