Spaces:
Sleeping
Sleeping
File size: 1,874 Bytes
f075308 |
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 |
"""
Takes a dataset directory and repeats the frames to include only a random frame from each video
This is needed to calculate same-frame FVD and DiFID
"""
import os
import random
import argparse
from typing import List
import shutil
from tqdm import tqdm
def construct_static_videos_dataset(videos_dir: os.PathLike, max_len: int=None, output_dir: os.PathLike=None, force_len: int=None):
output_dir = output_dir if not output_dir is None else f'{videos_dir}_freeze'
clips_paths = [os.path.join(videos_dir, d) for d in os.listdir(videos_dir)]
print(f'Saving into {output_dir}')
for video_idx, clip_path in enumerate(tqdm(clips_paths)):
frames_paths = os.listdir(clip_path)
frame_to_repeat = random.choice(frames_paths)
curr_output_dir = os.path.join(output_dir, f'{video_idx:05d}')
os.makedirs(curr_output_dir, exist_ok=True)
num_frames_to_create = force_len if not force_len is None else min(len(frames_paths), max_len)
for i in range(num_frames_to_create):
ext = os.path.splitext(frame_to_repeat)[1].lower()
target_file_path = os.path.join(curr_output_dir, f'{i:06d}{ext}')
shutil.copy(os.path.join(clip_path, frame_to_repeat), target_file_path)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('-d', '--directory', type=str, help='Directory with video frames')
parser.add_argument('-o', '--output_dir', type=None, help='Where to save the file?.')
parser.add_argument('-l', '--max_len', type=int, help='Max video length')
parser.add_argument('-fl', '--force_len', type=int, help='Force video length')
args = parser.parse_args()
construct_static_videos_dataset(
videos_dir=args.directory,
max_len=args.max_len,
output_dir=args.output_dir,
force_len=args.force_len,
)
|