Spaces:
Sleeping
Sleeping
File size: 4,258 Bytes
7e2a2a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import argparse
import os, sys
import cv2
import numpy as np
from natsort import natsorted
import subprocess
from tensorboard.backend.event_processing import event_accumulator
def decode_from_buffer(encoded_image_string):
s = np.frombuffer(encoded_image_string, dtype=np.uint8)
image = cv2.imdecode(s, cv2.IMREAD_COLOR)
return image
def main(args):
if args.log_dir == '':
print("Did not specify the log directory to compile video from. Please check again.")
sys.exit()
if not os.path.exists(args.output_dir):
os.mkdir(args.output_dir)
if not os.path.exists(os.path.join(args.output_dir, 'frames')):
os.mkdir(os.path.join(args.output_dir, 'frames'))
for filename in os.listdir(args.log_dir):
if 'events.out.tfevents' not in filename: # There should just be one file in the entire folder anyway, but just in case.
continue
event_file = os.path.join(args.log_dir, filename)
ea = event_accumulator.EventAccumulator(event_file)
ea.Reload()
keys = ea.images._buckets.keys()
entry_to_image_dict = {} # matches each entry of the form "epoch X iteration Y training_video img_category_name" to its associated images
epoch_dict = {} # each key contains all entries for that epoch. Keys are integers.
for entry in keys:
if 'epoch' not in entry or 'iteration' not in entry or 'training_video' not in entry:
continue
entry_to_image_dict[entry] = ea.images._buckets[entry].items[0].encoded_image_string
epoch = int(entry.split(" ")[1])
if epoch not in epoch_dict:
epoch_dict[epoch] = []
epoch_dict[epoch].append(entry)
epoch_list = epoch_dict.keys()
epoch_list_sorted = natsorted(epoch_list) # e.g. [1,2,3]
for epoch in epoch_list_sorted:
entries_in_this_epoch = epoch_dict[epoch]
entries_in_this_epoch = natsorted(entries_in_this_epoch)
# key are all the image category names. e.g. real_A, fake_B.
# values are the the entries with those names, in this epoch.
img_category_names = {}
for entry in entries_in_this_epoch:
img_category_name = entry.split(' ')[5]
if img_category_name not in img_category_names:
img_category_names[img_category_name] = []
img_category_names[img_category_name].append(entry)
epoch_img = [] # the final output image of this epoch
for img_category_name, entries_having_that_name in img_category_names.items():
imgs_for_one_category = []
for entry in entries_having_that_name:
img = decode_from_buffer(entry_to_image_dict[entry])
imgs_for_one_category.append(img)
imgs_for_one_category = np.concatenate(imgs_for_one_category, axis=1) # concatenate images along width
imgs_for_one_category = cv2.putText(imgs_for_one_category, text='epoch ' + str(epoch) + ' ' + img_category_name,
org=(0, imgs_for_one_category.shape[0] - 20), color=(0, 255, 0), thickness=2,
fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=2)
epoch_img.append(imgs_for_one_category)
epoch_img = np.concatenate(epoch_img, axis=0) # concatenate images along height
outpath = os.path.join(os.path.join(args.output_dir, 'frames'), '{:06}.jpg'.format(epoch))
cv2.imwrite(outpath, epoch_img)
command = 'ffmpeg -y -framerate ' + str(args.fps) + ' -i ' + \
os.path.join(os.path.join(args.output_dir, 'frames'), '%06d.jpg') + ' ' + os.path.join(args.output_dir, 'training_timeline.mp4')
sp = subprocess.Popen(command, shell=True)
while sp.poll() is None:
continue
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Style Master')
parser.add_argument('--log_dir', type=str, default='')
parser.add_argument('--output_dir', type=str, default='')
parser.add_argument('--fps', type=int, default=2)
args = parser.parse_args()
main(args)
|