Adding text labels to confusion matrix in Tensorflow for Tensorboard

Question

I am customizing the code from Tensorflow's example retrain.py, to train on my own images by adding additional dense layers, dropouts, momentum Gradient descent etc.

I wanted to add a confusion matrix to Tensorboard so I followed the first answer (Jerod's) from this post (I had also tried the second answer but facing some debugging issues) and added a few lines to add_evaluation_step function. So now it looks like :

def add_evaluation_step(result_tensor, ground_truth_tensor):

  with tf.name_scope('accuracy'):
    with tf.name_scope('correct_prediction'):
      prediction = tf.argmax(result_tensor, 1)
      correct_prediction = tf.equal(
          prediction, tf.argmax(ground_truth_tensor, 1))
    with tf.name_scope('accuracy'):
      evaluation_step = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
  tf.summary.scalar('accuracy', evaluation_step)
  print('prediction shape :: {}'.format(ground_truth_tensor))

  #Add confusion matrix
  batch_confusion = tf.confusion_matrix(tf.argmax(ground_truth_tensor, 1), prediction,
                                             num_classes=7,
                                             name='batch_confusion')
  # Create an accumulator variable to hold the counts
  confusion = tf.Variable( tf.zeros([7,7], 
                                          dtype=tf.int32 ),
                                 name='confusion' )
  # Create the update op for doing a "+=" accumulation on the batch
  confusion_update = confusion.assign( confusion + batch_confusion )
  # Cast counts to float so tf.summary.image renormalizes to [0,255]
  confusion_image = tf.reshape( tf.cast( confusion_update, tf.float32),
                                 [1, 7, 7, 1])

  tf.summary.image('confusion',confusion_image)

  return evaluation_step, prediction

This gives me :

My question is how can I add the labels to the rows(actual class) and columns(Predicted class). To get something like :

Answer 1

Jerod's answer contains almost everything you need, along for instance yauheni_selivonchyk's other answer on how to add custom images to Tensorboard.

It is then only a matter of putting everything together, i.e.:

1. Implementing methods to pass plotted images to summaries (as RGB arrays)
2. Implementing a method to convert the matrix data into a prettified confusion image
3. Defining your running evaluation operations to obtain the confusion matrix data (along other metrics) and preparing a placeholder and summary to receive the plotted image
4. Using everything together

---

1. Implementing methods to pass plotted images to summaries

import matplotlib
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import numpy as np
import tensorflow as tf

# Inspired by yauheni_selivonchyk on SO (https://stackoverflow.com/a/42815564/624547)

def get_figure(figsize=(10, 10), dpi=300):
    """
    Return a pyplot figure
    :param figsize:
    :param dpi:
    :return:
    """
    fig = plt.figure(num=0, figsize=figsize, dpi=dpi)
    fig.clf()
    return fig


def fig_to_rgb_array(fig, expand=True):
    """
    Convert figure into a RGB array
    :param fig:         PyPlot Figure
    :param expand:      Flag to expand
    :return:            RGB array
    """
    fig.canvas.draw()
    buf = fig.canvas.tostring_rgb()
    ncols, nrows = fig.canvas.get_width_height()
    shape = (nrows, ncols, 3) if not expand else (1, nrows, ncols, 3)
    return np.fromstring(buf, dtype=np.uint8).reshape(shape)


def figure_to_summary(fig, summary, place_holder):
    """
    Convert figure into TF summary
    :param fig:             Figure
    :param summary:         Summary to eval
    :param place_holder:    Summary image placeholder
    :return:                Summary
    """
    image = fig_to_rgb_array(fig)
    return summary.eval(feed_dict={place_holder: image})

2. Converting matrix data into a prettified confusion image

(here is an example, but it's up to what you want)

def confusion_matrix_to_image_summary(confusion_matrix, summary, place_holder, 
                                      list_classes, figsize=(9, 9)):
    """
    Plot confusion matrix and return as TF summary
    :param matrix:          Confusion matrix (N x N)
    :param filename:        Filename
    :param list_classes:    List of classes (N)
    :param figsize:         Pyplot figsize for the confusion image
    :return:                /
    """
    fig = get_figure(figsize=(9, 9))
    df = pd.DataFrame(confusion_matrix, index=list_classes, columns=list_classes)
    ax = sns.heatmap(df, annot=True, fmt='.0%')
    # Whatever embellishments you want:
    plt.title('Confusion matrix')
    plt.xticks(rotation=90)
    plt.yticks(rotation=0)
    image_sum = figure_to_summary(fig, summary, place_holder)
    return image_sum

3. Defining your evaluation operations & Preparing placeholder

# Inspired by Jerod's answer on SO (https://stackoverflow.com/a/42857070/624547)    
def add_evaluation_step(result_tensor, ground_truth_tensor, num_classes, confusion_matrix_figsize=(9, 9)):
    """
    Sets up the evaluation operations, computing the running accuracy and confusion image
    :param result_tensor:               Output tensor
    :param ground_truth_tensor:         Target class tensor
    :param num_classes:                 Number of classes
    :param confusion_matrix_figsize:    Pyplot figsize for the confusion image
    :return:                            TF operations, summaries and placeholders (see usage below)
    """
    scope = "evaluation"
    with tf.name_scope(scope):
        predictions = tf.argmax(result_tensor, 1, name="prediction")

        # Streaming accuracy (lookup and update tensors):
        accuracy, accuracy_update = tf.metrics.accuracy(ground_truth_tensor, predictions, name='accuracy')
        # Per-batch confusion matrix:
        batch_confusion = tf.confusion_matrix(ground_truth_tensor, predictions, num_classes=num_classes,
                                              name='batch_confusion')

        # Aggregated confusion matrix:
        confusion_matrix = tf.Variable(tf.zeros([num_classes, num_classes], dtype=tf.int32),
                                       name='confusion')
        confusion_update = confusion_matrix.assign(confusion_matrix + batch_confusion)

        # We suppose each batch contains a complete class, to directly normalize by its size:
        evaluate_streaming_metrics_op = tf.group(accuracy_update, confusion_update)

        # Confusion image from matrix (need to extend dims + cast to float so tf.summary.image renormalizes to [0,255]):
        confusion_image = tf.reshape(tf.cast(confusion_update, tf.float32), [1, num_classes, num_classes, 1])

        # Summaries:
        tf.summary.scalar('accuracy', accuracy, collections=[scope])
        summary_op = tf.summary.merge_all(scope)

        # Preparing placeholder for confusion image (so that we can pass the plotted image to it):
        #      (we basically pre-allocate a plot figure and pass its RGB array to a placeholder)
        confusion_image_placeholder = tf.placeholder(tf.uint8,
                                                     fig_to_rgb_array(get_figure(figsize=confusion_matrix_figsize)).shape)
        confusion_image_summary = tf.summary.image('confusion_image', confusion_image_placeholder)

    # Isolating all the variables stored by the metric operations:
    running_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=scope)
    running_vars += tf.get_collection(tf.GraphKeys.LOCAL_VARIABLES, scope=scope)

    # Initializer op to start/reset running variables
    reset_streaming_metrics_op = tf.variables_initializer(var_list=running_vars)

    return evaluate_streaming_metrics_op, reset_streaming_metrics_op, summary_op, confusion_image_summary, \
           confusion_image_placeholder, confusion_image

4. Putting everything together

A quick example how to use this, though it needs to be adapted to your training procedure, etc.

classes = ["obj1", "obj2", "obj3"]
num_classes = len(classes)
model = your_network(...)

evaluate_streaming_metrics_op, reset_streaming_metrics_op, summary_op,
confusion_image_summary,  confusion_image_placeholder, confusion_image = \
add_evaluation_step(model.output, model.target, num_classes)

def evaluate(session, model, eval_data_gen):
    """
    Evaluate the model
    :param session:         TF session
    :param eval_data_gen:   Data to evaluate on
    :return:                Evaluation summaries for Tensorboard
    """
    # Resetting streaming vars:
    session.run(reset_streaming_metrics_op)

    # Evaluating running ops over complete eval dataset, e.g.:
    for batch in eval_data_gen:
        feed_dict = {model.inputs: batch}
        session.run(evaluate_streaming_metrics_op, feed_dict=feed_dict)

    # Obtaining the final results:
    summary_str, confusion_results = session.run([summary_op, confusion_image])

    # Converting confusion data into plot into summary:
    confusion_img_str = confusion_matrix_to_image_summary(
        confusion_results[0,:,:,0], confusion_image_summary, confusion_image_placeholder, classes)
    summary_str += confusion_img_str

    return summary_str # to be given to a SummaryWriter