Spaces:
Sleeping
Sleeping
File size: 4,127 Bytes
7e2a2a5 |
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 |
import sys
sys.path.append('/share/group_machongyang/limengtian/libycnn2/lib')
import argparse
import os
import cv2
import numpy as np
from multiprocessing import Process
from tqdm import tqdm
from py_ycnn import face_detection
detector = face_detection()
detector.init_face_detection('/share/group_machongyang/limengtian/libycnn2/model/')
def work(lst, output_folder, layout, wid):
need_A = 'A' in layout
need_B = 'B' in layout
if need_A:
os.mkdir(os.path.join(output_folder, f'lmk_A_{wid:02d}'))
if need_B:
os.mkdir(os.path.join(output_folder, f'lmk_B_{wid:02d}'))
res_list = []
for i, (A_path, B_path) in tqdm(enumerate(lst)):
if os.path.isfile(A_path) and os.path.isfile(B_path):
if need_A:
A = cv2.imread(A_path)
A = detector.detectArray(A)
if len(A) < 1:
continue
A = np.array(A[0], dtype=np.float32)
A = A[:202].reshape(101, 2)
np.save(os.path.join(output_folder, f'lmk_A_{wid:02d}', f'{i:06d}.npy'), A)
if need_B:
B = cv2.imread(B_path)
B = detector.detectArray(B)
if len(B) < 1:
continue
B = np.array(B[0], dtype=np.float32)
B = B[:202].reshape(101, 2)
np.save(os.path.join(output_folder, f'lmk_B_{wid:02d}', f'{i:06d}.npy'), B)
if layout[0] == 'A':
lmk_A_path = os.path.join(output_folder, f'lmk_A_{wid:02d}', f'{i:06d}.npy')
else:
lmk_A_path = os.path.join(output_folder, f'lmk_B_{wid:02d}', f'{i:06d}.npy')
if layout[1] == 'A':
lmk_B_path = os.path.join(output_folder, f'lmk_A_{wid:02d}', f'{i:06d}.npy')
else:
lmk_B_path = os.path.join(output_folder, f'lmk_B_{wid:02d}', f'{i:06d}.npy')
res_list.append((A_path, B_path, lmk_A_path, lmk_B_path))
with open(os.path.join(output_folder, f'lmk_{wid:02d}.txt'), 'wt') as f:
for A_path, B_path, lmk_A_path, lmk_B_path in res_list:
f.write(f'{A_path} {B_path} {lmk_A_path} {lmk_B_path}\n')
if __name__ == '__main__':
# parse arguments
parser = argparse.ArgumentParser(description='Landmark Generator')
parser.add_argument('--input_list', type=str, required=True)
parser.add_argument('--output_folder', type=str, required=True)
parser.add_argument('--layout', type=str, default='AB')
parser.add_argument('--worker_count', type=int, default=16)
args = parser.parse_args()
# check arguments
assert os.path.isfile(args.input_list), f'FileNotFound: {args.input_list}.'
assert not os.path.exists(args.output_folder), 'Output folder exists.'
assert args.layout in ['AA', 'AB', 'BA', 'BB'], 'Layout error.'
assert args.worker_count >= 1 and args.worker_count <= 64, 'Worker count error.'
os.makedirs(args.output_folder)
# read input list and split to each process
lsts = [ [] for _ in range(args.worker_count) ]
with open(args.input_list, 'rt') as f:
line = f.readline().strip()
total_cnt = 0
while line:
A_path, B_path = line.split(' ')
lsts[total_cnt % args.worker_count].append((A_path, B_path))
line = f.readline().strip()
total_cnt += 1
print('Total images:', total_cnt)
# generate landmark
pool = []
for i in range(args.worker_count):
proc = Process(target=work, args=(lsts[i], args.output_folder, args.layout, i))
proc.start()
pool.append(proc)
for proc in pool:
proc.join()
# combine lists
new_fn = os.path.splitext(os.path.split(args.input_list)[1])[0] + '_with_lmk.txt'
with open(os.path.join(args.output_folder, new_fn), 'wt') as fw:
for i in range(args.worker_count):
with open(os.path.join(args.output_folder, f'lmk_{i:02d}.txt'), 'rt') as fr:
line = fr.readline()
while line:
fw.write(line)
line = fr.readline()
|