make_gif_from_tb_images.py

#!/usr/bin/python3
import argparse
import os
import io
import numpy as np
from string import ascii_lowercase
from random import sample

from PIL import Image
import moviepy.editor as mpy
from tensorboard.backend.event_processing import event_accumulator


parser = argparse.ArgumentParser(description='Convert Images from TensorBoard to gif')
parser.add_argument('--logdir', help='Path to TensorBoard logging folder')
parser.add_argument('--tag', type=str, default='Batch_of_Reconstructions', help='Tag of images')
parser.add_argument('--fps', type=int, default=5, help='frames per second')
parser.add_argument('--savedir', type=str, default='.', help='Where to save')
parser.add_argument('--savepath', help='Where to save')
parser.add_argument('--max_images', type=int, default=0, help='max number of images to extract')


def generate_random_name():
    letters = sample(ascii_lowercase, 10)
    return ''.join(letters) + '.gif'


class frames_generator:
    def __init__(self, images, with_steps=False):
        self.images = images
        self.num_images = len(images)
        self.c = 0

    def get_frame(self, t):
        if self.c >= self.num_images:
            img = self.images[-1]
        else:
            img = self.images[self.c]
        a = img.encoded_image_string
        image = Image.open(io.BytesIO(a))
        self.c += 1
        return np.array(image)


def main():
    args = parser.parse_args()
    accumulator = event_accumulator.EventAccumulator(
        args.logdir, size_guidance={event_accumulator.IMAGES: args.max_images})
    accumulator.Reload()
    images = accumulator.Images(args.tag)
    clip = mpy.VideoClip(frames_generator(images).get_frame,
                         duration=len(images) // args.fps)
    if args.savepath is None:
        save_path = os.path.join(args.savedir, generate_random_name())
    else:
        save_path = args.savepath
    clip.write_gif(save_path, fps=args.fps)


if __name__ == '__main__':
    main()