| import sys |
| import dlib |
| import os |
| import cv2 |
| import face_alignment |
| from multiprocessing import Pool, Process, Queue |
| import time |
|
|
|
|
| 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() |
|
|
