File size: 1,638 Bytes
df07554 |
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 |
import sys
import dlib
import os
import cv2
import face_alignment
import time
from multiprocessing import Pool, Process, Queue
def run(gpu, files):
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False, device='cuda')
print('gpu={},n_files={}'.format(gpu, len(files)))
tic = time.time()
count = 0
for (img_name, savename) in files:
I = cv2.imread(img_name)
points_list = fa.get_landmarks(I)
with open(savename, 'w') as f:
if(points_list is not None):
for points in points_list:
for (x, y) in points:
f.write('({}, {})\t'.format(x, y))
f.write('\n')
count += 1
if(count % 1000 == 0):
print('dst={},eta={}'.format(savename, (time.time()-tic)/(count) * (len(files) - count) / 3600.0))
if(__name__ == '__main__'):
with open('imgs.txt', 'r') as f:
data = [line.strip() for line in f.readlines()]
data = [(name, name.replace('.jpg', '.txt')) for name in data]
for (_, dst) in data:
dir, _ = os.path.split(dst)
if(not os.path.exists(dir)):
os.makedirs(dir)
processes = []
n_p = 3
gpus = ['1', '2', '3']
bs = len(data) // n_p
for i in range(n_p):
if(i == n_p - 1):
bs = len(data)
p = Process(target=run, args=(gpus[i],data[:bs],))
data = data[bs:]
p.start()
processes.append(p)
assert(len(data) == 0)
for p in processes:
p.join()
|