CM_func.py

  # Define the confusion matrix logging function
  def cm_logger(epoch, logs):

    # Create a name for this image
    i_name='Confusion Matrix: {}; {}'.format(model_name, ds.dataset_name)

    # Create the per-epoch callback for the confusion matrix
    cm_image=image_funcs.plot_to_image(metrics_funcs.generate_heatmap(
                                       dataset=ds,
                                       model=model).figure)
    # Write it to the logs
    with cm_file_writer.as_default():
      tf.summary.image(name=i_name, data=cm_image, step=epoch)