Spaces:
Runtime error
Runtime error
| # make_training_vid("outputs/zero123/64_teddy_rgba.png@20230627-195615", frames_per_vid=30, fps=20, max_iters=200) | |
| import argparse | |
| import glob | |
| import os | |
| import imageio | |
| import numpy as np | |
| from PIL import Image, ImageDraw | |
| from tqdm import tqdm | |
| def draw_text_in_image(img, texts): | |
| img = Image.fromarray(img) | |
| draw = ImageDraw.Draw(img) | |
| black, white = (0, 0, 0), (255, 255, 255) | |
| for i, text in enumerate(texts): | |
| draw.text((2, (img.size[1] // len(texts)) * i + 1), f"{text}", white) | |
| draw.text((0, (img.size[1] // len(texts)) * i + 1), f"{text}", white) | |
| draw.text((2, (img.size[1] // len(texts)) * i - 1), f"{text}", white) | |
| draw.text((0, (img.size[1] // len(texts)) * i - 1), f"{text}", white) | |
| draw.text((1, (img.size[1] // len(texts)) * i), f"{text}", black) | |
| return np.asarray(img) | |
| def make_training_vid(exp, frames_per_vid=1, fps=3, max_iters=None, max_vids=None): | |
| # exp = "/admin/home-vikram/git/threestudio/outputs/zero123/64_teddy_rgba.png@20230627-195615" | |
| files = glob.glob(os.path.join(exp, "save", "*.mp4")) | |
| if os.path.join(exp, "save", "training_vid.mp4") in files: | |
| files.remove(os.path.join(exp, "save", "training_vid.mp4")) | |
| its = [int(os.path.basename(file).split("-")[0].split("it")[-1]) for file in files] | |
| it_sort = np.argsort(its) | |
| files = list(np.array(files)[it_sort]) | |
| its = list(np.array(its)[it_sort]) | |
| max_vids = max_iters // its[0] if max_iters is not None else max_vids | |
| files, its = files[:max_vids], its[:max_vids] | |
| frames, i = [], 0 | |
| for it, file in tqdm(zip(its, files), total=len(files)): | |
| vid = imageio.mimread(file) | |
| for _ in range(frames_per_vid): | |
| frame = vid[i % len(vid)] | |
| frame = draw_text_in_image(frame, [str(it)]) | |
| frames.append(frame) | |
| i += 1 | |
| # Save | |
| imageio.mimwrite(os.path.join(exp, "save", "training_vid.mp4"), frames, fps=fps) | |
| def join(file1, file2, name): | |
| # file1 = "/admin/home-vikram/git/threestudio/outputs/zero123/OLD_64_dragon2_rgba.png@20230629-023028/save/it200-val.mp4" | |
| # file2 = "/admin/home-vikram/git/threestudio/outputs/zero123/64_dragon2_rgba.png@20230628-152734/save/it200-val.mp4" | |
| vid1 = imageio.mimread(file1) | |
| vid2 = imageio.mimread(file2) | |
| frames = [] | |
| for f1, f2 in zip(vid1, vid2): | |
| frames.append( | |
| np.concatenate([f1[:, : f1.shape[0]], f2[:, : f2.shape[0]]], axis=1) | |
| ) | |
| imageio.mimwrite(name, frames) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--exp", help="directory of experiment") | |
| parser.add_argument( | |
| "--frames_per_vid", type=int, default=1, help="# of frames from each val vid" | |
| ) | |
| parser.add_argument("--fps", type=int, help="max # of iters to save") | |
| parser.add_argument("--max_iters", type=int, help="max # of iters to save") | |
| parser.add_argument( | |
| "--max_vids", | |
| type=int, | |
| help="max # of val videos to save. Will be overridden by max_iters", | |
| ) | |
| args = parser.parse_args() | |
| make_training_vid( | |
| args.exp, args.frames_per_vid, args.fps, args.max_iters, args.max_vids | |
| ) | |