Spaces:
Runtime error
Runtime error
File size: 3,178 Bytes
2fa4776 |
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 |
# 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
)
|