|
|
|
|
|
import argparse
|
|
|
import glob
|
|
|
import os
|
|
|
import os.path as osp
|
|
|
import sys
|
|
|
from multiprocessing import Pool
|
|
|
|
|
|
|
|
|
def encode_video(frame_dir_item):
|
|
|
"""Encode frames to video using ffmpeg.
|
|
|
|
|
|
Args:
|
|
|
frame_dir_item (list): Rawframe item containing raw frame directory
|
|
|
full path, rawframe directory (short) path, rawframe directory id.
|
|
|
|
|
|
Returns:
|
|
|
bool: Whether synthesize video successfully.
|
|
|
"""
|
|
|
full_path, frame_dir_path, frame_dir_id = frame_dir_item
|
|
|
out_full_path = args.out_dir
|
|
|
|
|
|
img_name_tmpl = args.filename_tmpl + '.' + args.in_format
|
|
|
img_path = osp.join(full_path, img_name_tmpl)
|
|
|
|
|
|
out_vid_name = frame_dir_path + '.' + args.ext
|
|
|
out_vid_path = osp.join(out_full_path, out_vid_name)
|
|
|
|
|
|
cmd = osp.join(
|
|
|
f"ffmpeg -start_number {args.start_idx} -r {args.fps} -i '{img_path}' "
|
|
|
f"-vcodec {args.vcodec} '{out_vid_path}'")
|
|
|
os.system(cmd)
|
|
|
|
|
|
print(f'{frame_dir_id} {frame_dir_path} done')
|
|
|
sys.stdout.flush()
|
|
|
return True
|
|
|
|
|
|
|
|
|
def parse_args():
|
|
|
parser = argparse.ArgumentParser(description='synthesize videos')
|
|
|
parser.add_argument('src_dir', type=str, help='source rawframe directory')
|
|
|
parser.add_argument('out_dir', type=str, help='output video directory')
|
|
|
parser.add_argument(
|
|
|
'--fps', type=int, default=30, help='fps of videos to be synthesized')
|
|
|
parser.add_argument(
|
|
|
'--level',
|
|
|
type=int,
|
|
|
choices=[1, 2],
|
|
|
default=2,
|
|
|
help='directory level of data')
|
|
|
parser.add_argument(
|
|
|
'--num-worker',
|
|
|
type=int,
|
|
|
default=8,
|
|
|
help='number of workers to build videos')
|
|
|
parser.add_argument(
|
|
|
'--in-format',
|
|
|
type=str,
|
|
|
default='jpg',
|
|
|
choices=['jpg', 'png'],
|
|
|
help='input format')
|
|
|
parser.add_argument(
|
|
|
'--start-idx', type=int, default=0, help='starting index of rawframes')
|
|
|
parser.add_argument(
|
|
|
'--filename-tmpl',
|
|
|
type=str,
|
|
|
default='img_%05d',
|
|
|
help='filename template of rawframes')
|
|
|
parser.add_argument(
|
|
|
'--vcodec', type=str, default='mpeg4', help='coding method of videos')
|
|
|
parser.add_argument(
|
|
|
'--ext',
|
|
|
type=str,
|
|
|
default='mp4',
|
|
|
choices=['mp4', 'avi'],
|
|
|
help='video file extensions')
|
|
|
parser.add_argument('--num-gpu', type=int, default=8, help='number of GPU')
|
|
|
parser.add_argument(
|
|
|
'--resume',
|
|
|
action='store_true',
|
|
|
default=False,
|
|
|
help='resume optical flow extraction instead of overwriting')
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
return args
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
args = parse_args()
|
|
|
|
|
|
if not osp.isdir(args.out_dir):
|
|
|
print(f'Creating folder: {args.out_dir}')
|
|
|
os.makedirs(args.out_dir)
|
|
|
|
|
|
if args.level == 2:
|
|
|
classes = os.listdir(args.src_dir)
|
|
|
for classname in classes:
|
|
|
new_dir = osp.join(args.out_dir, classname)
|
|
|
if not osp.isdir(new_dir):
|
|
|
print(f'Creating folder: {new_dir}')
|
|
|
os.makedirs(new_dir)
|
|
|
|
|
|
print('Reading rgb frames from folder: ', args.src_dir)
|
|
|
print('Input format of rgb frames: ', args.in_format)
|
|
|
fullpath_list = glob.glob(args.src_dir + '/*' * args.level)
|
|
|
done_fullpath_list = glob.glob(args.src_dir + '/*' * args.level + '.' +
|
|
|
args.ext)
|
|
|
print('Total number of rgb frame folders found: ', len(fullpath_list))
|
|
|
|
|
|
if args.resume:
|
|
|
fullpath_list = set(fullpath_list).difference(set(done_fullpath_list))
|
|
|
fullpath_list = list(fullpath_list)
|
|
|
print('Resuming. number of videos to be synthesized: ',
|
|
|
len(fullpath_list))
|
|
|
|
|
|
if args.level == 2:
|
|
|
frame_dir_list = list(
|
|
|
map(
|
|
|
lambda p: osp.join(
|
|
|
osp.basename(osp.dirname(p)), osp.basename(p)),
|
|
|
fullpath_list))
|
|
|
elif args.level == 1:
|
|
|
frame_dir_list = list(map(osp.basename, fullpath_list))
|
|
|
|
|
|
pool = Pool(args.num_worker)
|
|
|
pool.map(encode_video,
|
|
|
zip(fullpath_list, frame_dir_list, range(len(frame_dir_list))))
|
|
|
|
