|
|
import os, glob |
|
|
from utils.commons.multiprocess_utils import multiprocess_run_tqdm |
|
|
|
|
|
from data_gen.utils.path_converter import PathConverter, pc |
|
|
|
|
|
|
|
|
|
|
|
def extract_img_job(video_name, raw_img_dir=None): |
|
|
if raw_img_dir is not None: |
|
|
out_path = raw_img_dir |
|
|
else: |
|
|
out_path = pc.to(video_name.replace(".mp4", ""), "vid", "gt") |
|
|
os.makedirs(out_path, exist_ok=True) |
|
|
ffmpeg_path = "/usr/bin/ffmpeg" |
|
|
cmd = f'{ffmpeg_path} -i {video_name} -vf fps={25},scale=w=512:h=512 -qmin 1 -q:v 1 -start_number 0 -v quiet {os.path.join(out_path, "%8d.jpg")}' |
|
|
os.system(cmd) |
|
|
|
|
|
if __name__ == '__main__': |
|
|
import argparse, glob, tqdm, random |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument("--vid_dir", default='/home/tiger/datasets/raw/CelebV-HQ/video') |
|
|
parser.add_argument("--ds_name", default='CelebV-HQ') |
|
|
parser.add_argument("--num_workers", default=64, type=int) |
|
|
parser.add_argument("--process_id", default=0, type=int) |
|
|
parser.add_argument("--total_process", default=1, type=int) |
|
|
args = parser.parse_args() |
|
|
vid_dir = args.vid_dir |
|
|
ds_name = args.ds_name |
|
|
if ds_name in ['lrs3_trainval']: |
|
|
mp4_name_pattern = os.path.join(vid_dir, "*/*.mp4") |
|
|
elif ds_name in ['TH1KH_512', 'CelebV-HQ']: |
|
|
vid_names = glob.glob(os.path.join(vid_dir, "*.mp4")) |
|
|
elif ds_name in ['lrs2', 'lrs3', 'voxceleb2']: |
|
|
vid_name_pattern = os.path.join(vid_dir, "*/*/*.mp4") |
|
|
vid_names = glob.glob(vid_name_pattern) |
|
|
elif ds_name in ["RAVDESS", 'VFHQ']: |
|
|
vid_name_pattern = os.path.join(vid_dir, "*/*/*/*.mp4") |
|
|
vid_names = glob.glob(vid_name_pattern) |
|
|
vid_names = sorted(vid_names) |
|
|
|
|
|
process_id = args.process_id |
|
|
total_process = args.total_process |
|
|
if total_process > 1: |
|
|
assert process_id <= total_process -1 |
|
|
num_samples_per_process = len(vid_names) // total_process |
|
|
if process_id == total_process: |
|
|
vid_names = vid_names[process_id * num_samples_per_process : ] |
|
|
else: |
|
|
vid_names = vid_names[process_id * num_samples_per_process : (process_id+1) * num_samples_per_process] |
|
|
|
|
|
for i, res in multiprocess_run_tqdm(extract_img_job, vid_names, num_workers=args.num_workers, desc="extracting images"): |
|
|
pass |
|
|
|
|
|
|
