torchnet / extract_lip_v2.py
milselarch's picture
push to main
df07554
raw
history blame contribute delete
No virus
1.45 kB
import os.path
import sys
from scripts import extract_lip
from multiprocessing import Pool, Process, Queue
from Loader import GridLoader
from dataset import GridDataset
from models.LipNet import LipNet
PARALLEL = False
loader = GridLoader()
video_paths = loader.load_video_paths(fetch_all_paths=False)
image_dirs, anno_dirs, target_dirs = [], [], []
for k in range(1, 35):
speaker_name = f's{k}'
image_dirpath = f'lip/GRID_imgs/{speaker_name}'
annos_dirpath = f'lip/GRID_aligns/{speaker_name}'
target_dirpath = f'lip/GRID_lip_imgs/{speaker_name}'
if not os.path.exists(image_dirpath):
continue
sentences = os.listdir(image_dirpath)
for sentence in sentences:
sentence_dir = f'{image_dirpath}/{sentence}'
print('SS', sentence_dir)
image_dirs.append(sentence_dir)
anno_dirs.append(annos_dirpath)
target_dirs.append(target_dirpath)
print(video_paths[:10])
for dst in target_dirs:
if not os.path.exists(dst):
os.makedirs(dst)
data = list(zip(image_dirs, anno_dirs, target_dirs))
processes = []
n_p = 8
bs = len(data) // n_p
if PARALLEL:
for i in range(n_p):
if i == n_p - 1:
bs = len(data)
p = Process(target=extract_lip.run, args=(data[:bs],))
data = data[bs:]
p.start()
processes.append(p)
assert (len(data) == 0)
for p in processes:
p.join()
else:
extract_lip.run(data)