Spaces:
Sleeping
Sleeping
File size: 1,574 Bytes
d16b52d |
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 |
import os
import os.path as osp
import imageio
import numpy as np
import torch
import torchvision
from einops import rearrange
from PIL import Image
def read_video_frames(folder: str, height=None, width=None):
"""
Read video frames from the given folder.
Output:
frames, in [0, 255], uint8, THWC
"""
_SUPPORTED_EXTENSIONS = [".png", ".jpg", ".jpeg"]
frames = [f for f in os.listdir(folder) if osp.splitext(f)[1] in _SUPPORTED_EXTENSIONS]
# sort frames
sorted_frames = sorted(frames, key=lambda x: int(osp.splitext(x)[0]))
sorted_frames = [osp.join(folder, f) for f in sorted_frames]
if height is not None and width is not None:
sorted_frames = [np.array(Image.open(f).resize((width, height))) for f in sorted_frames]
else:
sorted_frames = [np.array(Image.open(f)) for f in sorted_frames]
sorted_frames = torch.stack([torch.from_numpy(f) for f in sorted_frames], dim=0)
return sorted_frames
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")
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 # -1,1 -> 0,1
x = (x * 255).numpy().astype(np.uint8)
outputs.append(x)
parent_dir = os.path.dirname(path)
if parent_dir != "":
os.makedirs(parent_dir, exist_ok=True)
imageio.mimsave(path, outputs, fps=fps, loop=0)
|