|
|
|
|
|
import argparse
|
|
|
import glob
|
|
|
import os
|
|
|
import os.path as osp
|
|
|
import sys
|
|
|
from multiprocessing import Pool
|
|
|
|
|
|
|
|
|
def resize_videos(vid_item):
|
|
|
"""Generate resized video cache.
|
|
|
|
|
|
Args:
|
|
|
vid_item (list): Video item containing video full path,
|
|
|
video relative path.
|
|
|
|
|
|
Returns:
|
|
|
bool: Whether generate video cache successfully.
|
|
|
"""
|
|
|
full_path, vid_path = vid_item
|
|
|
|
|
|
if args.to_mp4:
|
|
|
vid_path = vid_path.split('.')
|
|
|
assert len(vid_path) == 2, \
|
|
|
f"Video path '{vid_path}' contain more than one dot"
|
|
|
vid_path = vid_path[0] + '.mp4'
|
|
|
out_full_path = osp.join(args.out_dir, vid_path)
|
|
|
dir_name = osp.dirname(vid_path)
|
|
|
out_dir = osp.join(args.out_dir, dir_name)
|
|
|
if not osp.exists(out_dir):
|
|
|
os.makedirs(out_dir)
|
|
|
result = os.popen(
|
|
|
f'ffprobe -hide_banner -loglevel error -select_streams v:0 -show_entries stream=width,height -of csv=p=0 {full_path}'
|
|
|
)
|
|
|
w, h = [int(d) for d in result.readline().rstrip().split(',')]
|
|
|
if w > h:
|
|
|
cmd = (f'ffmpeg -hide_banner -loglevel error -i {full_path} '
|
|
|
f'-vf {"mpdecimate," if args.remove_dup else ""}'
|
|
|
f'scale=-2:{args.scale} '
|
|
|
f'{"-vsync vfr" if args.remove_dup else ""} '
|
|
|
f'-c:v libx264 {"-g 16" if args.dense else ""} '
|
|
|
f'-an {out_full_path} -y')
|
|
|
else:
|
|
|
cmd = (f'ffmpeg -hide_banner -loglevel error -i {full_path} '
|
|
|
f'-vf {"mpdecimate," if args.remove_dup else ""}'
|
|
|
f'scale={args.scale}:-2 '
|
|
|
f'{"-vsync vfr" if args.remove_dup else ""} '
|
|
|
f'-c:v libx264 {"-g 16" if args.dense else ""} '
|
|
|
f'-an {out_full_path} -y')
|
|
|
os.popen(cmd)
|
|
|
print(f'{vid_path} done')
|
|
|
sys.stdout.flush()
|
|
|
return True
|
|
|
|
|
|
|
|
|
def parse_args():
|
|
|
parser = argparse.ArgumentParser(
|
|
|
description='Generate the resized cache of original videos')
|
|
|
parser.add_argument('src_dir', type=str, help='source video directory')
|
|
|
parser.add_argument('out_dir', type=str, help='output video directory')
|
|
|
parser.add_argument(
|
|
|
'--dense',
|
|
|
action='store_true',
|
|
|
help='whether to generate a faster cache')
|
|
|
parser.add_argument(
|
|
|
'--level',
|
|
|
type=int,
|
|
|
choices=[1, 2],
|
|
|
default=2,
|
|
|
help='directory level of data')
|
|
|
parser.add_argument(
|
|
|
'--remove-dup',
|
|
|
action='store_true',
|
|
|
help='whether to remove duplicated frames')
|
|
|
parser.add_argument(
|
|
|
'--ext',
|
|
|
type=str,
|
|
|
default='mp4',
|
|
|
choices=['avi', 'mp4', 'webm', 'mkv'],
|
|
|
help='video file extensions')
|
|
|
parser.add_argument(
|
|
|
'--to-mp4',
|
|
|
action='store_true',
|
|
|
help='whether to output videos in mp4 format')
|
|
|
parser.add_argument(
|
|
|
'--scale',
|
|
|
type=int,
|
|
|
default=256,
|
|
|
help='resize image short side length keeping ratio')
|
|
|
parser.add_argument(
|
|
|
'--num-worker', type=int, default=8, help='number of workers')
|
|
|
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)
|
|
|
|
|
|
print('Reading videos from folder: ', args.src_dir)
|
|
|
print('Extension of videos: ', args.ext)
|
|
|
fullpath_list = glob.glob(args.src_dir + '/*' * args.level + '.' +
|
|
|
args.ext)
|
|
|
done_fullpath_list = glob.glob(args.out_dir + '/*' * args.level + args.ext)
|
|
|
print('Total number of videos found: ', len(fullpath_list))
|
|
|
print('Total number of videos transfer finished: ',
|
|
|
len(done_fullpath_list))
|
|
|
if args.level == 2:
|
|
|
vid_list = list(
|
|
|
map(
|
|
|
lambda p: osp.join(
|
|
|
osp.basename(osp.dirname(p)), osp.basename(p)),
|
|
|
fullpath_list))
|
|
|
elif args.level == 1:
|
|
|
vid_list = list(map(osp.basename, fullpath_list))
|
|
|
pool = Pool(args.num_worker)
|
|
|
pool.map(resize_videos, zip(fullpath_list, vid_list))
|
|
|
|
