|
import sys |
|
import time |
|
|
|
import face_alignment |
|
|
|
sys.path.append('..') |
|
|
|
import multiprocessing |
|
from multiprocessing import Pool |
|
|
|
from tqdm.auto import tqdm |
|
from Extractor import export_frames |
|
from helpers import * |
|
|
|
RUN_PARALLEL = True |
|
PARALLEL_USE_GPU = True |
|
NUM_WORKERS = 12 |
|
|
|
LRS2_dir = '/home/milselarch/projects/SUTD/50-035/LRS2/lrs2_v1' |
|
videos_dir = os.path.join(LRS2_dir, 'mvlrs_v1/main') |
|
images_dir = os.path.join(LRS2_dir, 'mvlrs_v1/main_images') |
|
|
|
usable_sets_paths = [ |
|
'../data/LRS2_CTC2_train.txt', |
|
'../data/LRS2_CTC2_val.txt', |
|
'../data/LRS2_CTC2_test.txt' |
|
] |
|
|
|
USABLE_VIDEO_GROUP_NAMES = [] |
|
for usable_sets_path in usable_sets_paths: |
|
USABLE_VIDEO_GROUP_NAMES.extend([ |
|
x.strip() for x in open(usable_sets_path).readlines() |
|
]) |
|
|
|
USABLE_VIDEO_GROUP_NAMES = set(USABLE_VIDEO_GROUP_NAMES) |
|
lock = multiprocessing.Lock() |
|
group_ids = os.listdir(videos_dir) |
|
video_paths = [] |
|
|
|
for group_id in group_ids: |
|
group_dirpath = os.path.join(videos_dir, group_id) |
|
filenames = os.listdir(group_dirpath) |
|
|
|
for filename in filenames: |
|
if not filename.endswith('.mp4'): |
|
continue |
|
|
|
basename, _ = os.path.splitext(filename) |
|
video_group_name = f'{group_id}/{basename}' |
|
if video_group_name not in USABLE_VIDEO_GROUP_NAMES: |
|
continue |
|
|
|
video_paths.append((group_id, filename)) |
|
|
|
|
|
def process_video( |
|
video_path, video_images_dir, use_gpu |
|
): |
|
export_frames( |
|
video_path, video_images_dir, |
|
recycle_landmarks=True, use_gpu=use_gpu, |
|
) |
|
|
|
|
|
|
|
return video_path, video_images_dir |
|
|
|
|
|
def callback(result): |
|
pbar.desc = str(result) |
|
pbar.update(1) |
|
|
|
|
|
|
|
|
|
if RUN_PARALLEL: |
|
pbar = tqdm(video_paths) |
|
pool = Pool(processes=NUM_WORKERS) |
|
jobs = [] |
|
|
|
for data_pair in video_paths: |
|
group_id, video_filename = data_pair |
|
pbar.desc = str(data_pair) |
|
|
|
basename, _ = os.path.splitext(video_filename) |
|
temp_id = f'{group_id}-{basename}' |
|
|
|
|
|
video_path = os.path.join(videos_dir, group_id, video_filename) |
|
group_images_dir = os.path.join(images_dir, group_id) |
|
video_images_dir = os.path.join(group_images_dir, basename) |
|
|
|
if not os.path.exists(group_images_dir): |
|
os.mkdir(group_images_dir) |
|
if not os.path.exists(video_images_dir): |
|
os.mkdir(video_images_dir) |
|
|
|
job_kwargs = kwargify( |
|
video_path=video_path, |
|
video_images_dir=video_images_dir, |
|
use_gpu=PARALLEL_USE_GPU |
|
) |
|
job = pool.apply_async( |
|
process_video, kwds=job_kwargs, |
|
callback=callback |
|
) |
|
|
|
jobs.append(job) |
|
|
|
print('ALL JOBS ADDED', len(jobs)) |
|
|
|
for job in jobs: |
|
job.wait() |
|
|
|
pool.close() |
|
pool.join() |
|
else: |
|
pbar = tqdm(video_paths) |
|
|
|
for data_pair in pbar: |
|
group_id, video_filename = data_pair |
|
|
|
basename, _ = os.path.splitext(video_filename) |
|
temp_id = f'{group_id}-{basename}' |
|
|
|
|
|
video_path = os.path.join(videos_dir, group_id, video_filename) |
|
group_images_dir = os.path.join(images_dir, group_id) |
|
video_images_dir = os.path.join(group_images_dir, basename) |
|
|
|
if not os.path.exists(group_images_dir): |
|
os.mkdir(group_images_dir) |
|
if not os.path.exists(video_images_dir): |
|
os.mkdir(video_images_dir) |
|
|
|
process_video( |
|
video_path=video_path, |
|
video_images_dir=video_images_dir, |
|
use_gpu=True |
|
) |
|
|
|
print('TOTAL VIDEOS', len(video_paths)) |
|
|
|
print('>>>') |