Spaces:
Sleeping
Sleeping
| """ | |
| Converts a directory of video frames into an mp4-grid | |
| """ | |
| import sys; sys.path.extend(['.']) | |
| import os | |
| import argparse | |
| import random | |
| import numpy as np | |
| import torch | |
| from torch import Tensor | |
| import torchvision.transforms.functional as TVF | |
| from torchvision import utils | |
| from PIL import Image | |
| from tqdm import tqdm | |
| import torchvision | |
| def frames_to_video_grid(videos_dir: os.PathLike, num_videos: int, length: int, fps: int, output_path: os.PathLike, select_random: bool=False, random_seed: int=None): | |
| clips_paths = [os.path.join(videos_dir, d) for d in os.listdir(videos_dir)] | |
| # bad_idx = [0, 9, 11, 16] | |
| # clips_paths = [c for i, c in enumerate(clips_paths) if not i in bad_idx] | |
| if select_random: | |
| random.seed(random_seed) | |
| clips_paths = random.sample(clips_paths, k=num_videos) | |
| else: | |
| clips_paths = clips_paths[:num_videos] | |
| videos = [read_first_n_frames(d, length) for d in tqdm(clips_paths, desc='Reading data...')] # [num_videos, length, c, h, w] | |
| videos = [fill_with_black_squares(v, length) for v in tqdm(videos, desc='Adding empty frames')] # [num_videos, length, c, h, w] | |
| frame_grids = torch.stack(videos).permute(1, 0, 2, 3, 4) # [video_len, num_videos, c, h, w] | |
| frame_grids = [utils.make_grid(fs, nrow=int(np.ceil(np.sqrt(num_videos)))) for fs in tqdm(frame_grids, desc='Making grids')] | |
| if os.path.dirname(output_path) != "": | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| frame_grids = (torch.stack(frame_grids) * 255).to(torch.uint8).permute(0, 2, 3, 1) # [T, H, W, C] | |
| torchvision.io.write_video(output_path, frame_grids, fps=fps, video_codec='h264', options={'crf': '10'}) | |
| def read_first_n_frames(d: os.PathLike, num_frames: int) -> Tensor: | |
| images = [Image.open(os.path.join(d, f)) for f in sorted(os.listdir(d))[:num_frames]] | |
| images = [TVF.to_tensor(x) for x in images] | |
| return torch.stack(images) | |
| def fill_with_black_squares(video, desired_len: int) -> Tensor: | |
| if len(video) >= desired_len: | |
| return video | |
| return torch.cat([ | |
| video, | |
| torch.zeros_like(video[0]).unsqueeze(0).repeat(desired_len - len(video), 1, 1, 1), | |
| ], dim=0) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('-d', '--directory', type=str, help='Directory with video frames') | |
| parser.add_argument('-n', '--num_videos', type=int, help='Number of videos to consider') | |
| parser.add_argument('-l', '--length', type=int, help='Video length (in frames)') | |
| parser.add_argument('--fps', type=int, default=25, help='FPS to save with.') | |
| parser.add_argument('-o', '--output_path', type=str, help='Where to save the file?.') | |
| parser.add_argument('--select_random', action='store_true', help='Select videos at random?') | |
| parser.add_argument('--random_seed', type=int, default=None, help='Random seed when selecting videos at random') | |
| args = parser.parse_args() | |
| frames_to_video_grid( | |
| videos_dir=args.directory, | |
| num_videos=args.num_videos, | |
| length=args.length, | |
| fps=args.fps, | |
| output_path=args.output_path, | |
| select_random=args.select_random, | |
| random_seed=args.random_seed, | |
| ) | |