File size: 4,279 Bytes
d670799 |
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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
# Copyright (c) OpenMMLab. All rights reserved.
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))))
|