|
import importlib |
|
import os |
|
import os.path as osp |
|
import shutil |
|
import sys |
|
from pathlib import Path |
|
|
|
import av |
|
import numpy as np |
|
import torch |
|
import torchvision |
|
from einops import rearrange |
|
from PIL import Image |
|
|
|
|
|
def seed_everything(seed): |
|
import random |
|
|
|
import numpy as np |
|
|
|
torch.manual_seed(seed) |
|
torch.cuda.manual_seed_all(seed) |
|
np.random.seed(seed % (2**32)) |
|
random.seed(seed) |
|
|
|
|
|
def import_filename(filename): |
|
spec = importlib.util.spec_from_file_location("mymodule", filename) |
|
module = importlib.util.module_from_spec(spec) |
|
sys.modules[spec.name] = module |
|
spec.loader.exec_module(module) |
|
return module |
|
|
|
|
|
def delete_additional_ckpt(base_path, num_keep): |
|
dirs = [] |
|
for d in os.listdir(base_path): |
|
if d.startswith("checkpoint-"): |
|
dirs.append(d) |
|
num_tot = len(dirs) |
|
if num_tot <= num_keep: |
|
return |
|
|
|
del_dirs = sorted(dirs, key=lambda x: int(x.split("-")[-1]))[: num_tot - num_keep] |
|
for d in del_dirs: |
|
path_to_dir = osp.join(base_path, d) |
|
if osp.exists(path_to_dir): |
|
shutil.rmtree(path_to_dir) |
|
|
|
|
|
def save_videos_from_pil(pil_images, path, fps=8): |
|
import av |
|
|
|
save_fmt = Path(path).suffix |
|
os.makedirs(os.path.dirname(path), exist_ok=True) |
|
width, height = pil_images[0].size |
|
|
|
if save_fmt == ".mp4": |
|
codec = "libx264" |
|
container = av.open(path, "w") |
|
stream = container.add_stream(codec, rate=fps) |
|
|
|
stream.width = width |
|
stream.height = height |
|
stream.pix_fmt = 'yuv420p' |
|
stream.bit_rate = 10000000 |
|
stream.options["crf"] = "18" |
|
|
|
|
|
|
|
for pil_image in pil_images: |
|
|
|
av_frame = av.VideoFrame.from_image(pil_image) |
|
container.mux(stream.encode(av_frame)) |
|
container.mux(stream.encode()) |
|
container.close() |
|
|
|
elif save_fmt == ".gif": |
|
pil_images[0].save( |
|
fp=path, |
|
format="GIF", |
|
append_images=pil_images[1:], |
|
save_all=True, |
|
duration=(1 / fps * 1000), |
|
loop=0, |
|
) |
|
else: |
|
raise ValueError("Unsupported file type. Use .mp4 or .gif.") |
|
|
|
|
|
def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8): |
|
videos = rearrange(videos, "b c t h w -> t b c h w") |
|
height, width = videos.shape[-2:] |
|
outputs = [] |
|
|
|
for x in videos: |
|
x = torchvision.utils.make_grid(x, nrow=n_rows) |
|
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) |
|
if rescale: |
|
x = (x + 1.0) / 2.0 |
|
x = (x * 255).numpy().astype(np.uint8) |
|
x = Image.fromarray(x) |
|
|
|
outputs.append(x) |
|
|
|
os.makedirs(os.path.dirname(path), exist_ok=True) |
|
|
|
save_videos_from_pil(outputs, path, fps) |
|
|
|
|
|
def read_frames(video_path): |
|
container = av.open(video_path) |
|
|
|
video_stream = next(s for s in container.streams if s.type == "video") |
|
frames = [] |
|
for packet in container.demux(video_stream): |
|
for frame in packet.decode(): |
|
image = Image.frombytes( |
|
"RGB", |
|
(frame.width, frame.height), |
|
frame.to_rgb().to_ndarray(), |
|
) |
|
frames.append(image) |
|
|
|
return frames |
|
|
|
|
|
def get_fps(video_path): |
|
container = av.open(video_path) |
|
video_stream = next(s for s in container.streams if s.type == "video") |
|
fps = video_stream.average_rate |
|
container.close() |
|
return fps |
|
|
