logger.py

# Code referenced from https://gist.github.com/gyglim/1f8dfb1b5c82627ae3efcfbbadb9f514
# tensorboard --logdir='./logs' --port=6006
# tensorboard --logdir=old:out/2019-07-20___18-13-10/tensorboard,new:out/2019-07-21___14-45-41/tensorboard --port=6006

import os
import sys
import argparse

import tensorflow as tf
import numpy as np
import scipy.misc 
try:
    from StringIO import StringIO  # Python 2.7
except ImportError:
    from io import BytesIO         # Python 3.x


class Logger(object):
    
    def __init__(self, log_dir):
        """Create a summary writer logging to log_dir."""
        self.writer = tf.summary.FileWriter(log_dir)

    def scalar_summary(self, tag, value, step):
        """Log a scalar variable."""
        summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)])
        self.writer.add_summary(summary, step)

    def image_summary(self, tag, images, step):
        """Log a list of images."""

        img_summaries = []
        for i, img in enumerate(images):
            # Write the image to a string
            try:
                s = StringIO()
            except:
                s = BytesIO()
            scipy.misc.toimage(img).save(s, format="png")

            # Create an Image object
            img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(),
                                       height=img.shape[0],
                                       width=img.shape[1])
            # Create a Summary value
            img_summaries.append(tf.Summary.Value(tag='%s/%d' % (tag, i), image=img_sum))

        # Create and write Summary
        summary = tf.Summary(value=img_summaries)
        self.writer.add_summary(summary, step)
        
    def histo_summary(self, tag, values, step, bins=1000):
        """Log a histogram of the tensor of values."""

        # Create a histogram using numpy
        counts, bin_edges = np.histogram(values, bins=bins)

        # Fill the fields of the histogram proto
        hist = tf.HistogramProto()
        hist.min = float(np.min(values))
        hist.max = float(np.max(values))
        hist.num = int(np.prod(values.shape))
        hist.sum = float(np.sum(values))
        hist.sum_squares = float(np.sum(values**2))

        # Drop the start of the first bin
        bin_edges = bin_edges[1:]

        # Add bin edges and counts
        for edge in bin_edges:
            hist.bucket_limit.append(edge)
        for c in counts:
            hist.bucket.append(c)

        # Create and write Summary
        summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)])
        self.writer.add_summary(summary, step)
        self.writer.flush()



def main(ARGS):
    out_dir = os.path.expanduser(ARGS.out_dir)
    assert os.path.isdir(out_dir), "Path {} does not exist ".format(ARGS.out_dir)

    print("\nStarting Tensorboard")

    date_folders = []
    files = os.listdir(out_dir)
    if len(files):
        for f in files:
            date_folders.append(f)

    assert date_folders, "No tensorboard folders"

    command_string = "tensorboard --logdir="
    date_folders.sort()  
    for i,folder in enumerate(date_folders):
        tensorboard_path = os.path.join(os.path.join(out_dir, folder), 'tensorboard')
        # print("Full Tensorboard Path: {}".format(tensorboard_path))
        command_string += str(i+1) + ":" + tensorboard_path
        if i != len(date_folders)-1:
            command_string += ","

    # Port
    command_string += " --port={}".format(ARGS.tensorboard_port)

    print("FINAL Command String: {}".format(command_string))
    print("\n")
    os.system(command_string)

    
def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--out_dir', type=str,  help='Directory where models and event logs are stored.', default='./out')
    parser.add_argument('--tensorboard_port', type=int, help='Tensorboard port for command', default=6006)
    return parser.parse_args(argv)
  
if __name__ == '__main__':
    main(parse_arguments(sys.argv[1:]))