tensorboard(X) のログから gif 画像を作るスクリプト

record_reader.py

'''
references:
    https://github.com/lanpa/tensorboardX/blob/master/tensorboardX/event_file_writer.py
'''
import struct

from tensorboardX.crc32c import crc32c
from tensorboardX.record_writer import masked_crc32c
from tensorboardX.proto.event_pb2 import Event


class RecordReader():
    def __init__(self, path):
        self.path = path
        self._reader = open(path, 'rb')

    def read(self):
        length_b = self._reader.read(8)
        if len(length_b) == 0:
            raise EOFError()
        length = struct.unpack('Q', length_b)[0]
        length_masked_crc32c_b = self._reader.read(4)
        length_masked_crc32c = struct.unpack('I', length_masked_crc32c_b)[0]
        assert length_masked_crc32c == masked_crc32c(length_b), 'length crc32c'

        event_str = self._reader.read(length)
        event_masked_crc32c_b = self._reader.read(4)
        event_masked_crc32c = struct.unpack('I', event_masked_crc32c_b)[0]
        assert event_masked_crc32c == masked_crc32c(event_str), 'event crc32c'

        res = Event()
        res.ParseFromString(event_str)
        return res

    def close(self):
        self._reader.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def __iter__(self):
        return self

    def __next__(self):
        try:
            return self.read()
        except EOFError:
            raise StopIteration()

tf_evet_log_gif.py

import os

from argparse import ArgumentParser
from record_reader import RecordReader
from tensorboardX import FileWriter
from PIL import Image, ImageDraw
import io

parser = ArgumentParser()
parser.add_argument('src', type=str, help='path to event file')
parser.add_argument('tag', type=str, help='tag name of image')
parser.add_argument('--output', type=str, default='out.gif')
parser.add_argument('--keep_freq', type=str, default='5', help='ex) 5,8 then epoch 5n and 8n will be kept')
parser.add_argument('--keep_epoch', type=str, default=None, help='ex) 5,8 then epoch 5 and 8 will be kept')
parser.add_argument('--duration', type=int, default=500, help='ms')
args = parser.parse_args()

args.keep_freq = list(map(int, args.keep_freq.split(',')))
args.keep_epoch = list(map(int, args.keep_epoch.split(','))) if args.keep_epoch else []

images = []
reader = RecordReader(args.src)
for event in reader:
    save = False
    save = save or any([event.step % freq == 0 for freq in args.keep_freq])
    save = save or event.step in args.keep_epoch
    save = save and (len(event.summary.value) > 0 and event.summary.value[0].WhichOneof('value') == 'image')
    save = save and event.summary.value[0].tag == args.tag
    if not save:
        continue
    image = event.summary.value[0].image
    image = Image.open(io.BytesIO(image.encoded_image_string))
    # 画像に epoch を表示する
    draw = ImageDraw.Draw(image)
    draw.text((5, 5), 'step {}'.format(event.step), fill=128)
    images.append(image)
    #
    print('{:5}'.format(len(images)), end='\r')

reader.close()

# gif 生成
images[0].save(args.output, save_all=True, append_images=images[1:], duration=args.duration, loop=0)
print('generated gif from {} images.'.format(len(images)))