Spaces:
Sleeping
Sleeping
| from argparse import ArgumentParser | |
| import time | |
| import numpy as np | |
| import os | |
| import json | |
| import sys | |
| from PIL import Image | |
| import multiprocessing as mp | |
| import math | |
| import torch | |
| import torchvision.transforms as trans | |
| sys.path.append(".") | |
| sys.path.append("..") | |
| from models.mtcnn.mtcnn import MTCNN | |
| from models.encoders.model_irse import IR_101 | |
| from configs.paths_config import model_paths | |
| CIRCULAR_FACE_PATH = model_paths['circular_face'] | |
| def chunks(lst, n): | |
| """Yield successive n-sized chunks from lst.""" | |
| for i in range(0, len(lst), n): | |
| yield lst[i:i + n] | |
| def extract_on_paths(file_paths): | |
| facenet = IR_101(input_size=112) | |
| facenet.load_state_dict(torch.load(CIRCULAR_FACE_PATH)) | |
| facenet.cuda() | |
| facenet.eval() | |
| mtcnn = MTCNN() | |
| id_transform = trans.Compose([ | |
| trans.ToTensor(), | |
| trans.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) | |
| ]) | |
| pid = mp.current_process().name | |
| print('\t{} is starting to extract on {} images'.format(pid, len(file_paths))) | |
| tot_count = len(file_paths) | |
| count = 0 | |
| scores_dict = {} | |
| for res_path, gt_path in file_paths: | |
| count += 1 | |
| if count % 100 == 0: | |
| print('{} done with {}/{}'.format(pid, count, tot_count)) | |
| if True: | |
| input_im = Image.open(res_path) | |
| input_im, _ = mtcnn.align(input_im) | |
| if input_im is None: | |
| print('{} skipping {}'.format(pid, res_path)) | |
| continue | |
| input_id = facenet(id_transform(input_im).unsqueeze(0).cuda())[0] | |
| result_im = Image.open(gt_path) | |
| result_im, _ = mtcnn.align(result_im) | |
| if result_im is None: | |
| print('{} skipping {}'.format(pid, gt_path)) | |
| continue | |
| result_id = facenet(id_transform(result_im).unsqueeze(0).cuda())[0] | |
| score = float(input_id.dot(result_id)) | |
| scores_dict[os.path.basename(gt_path)] = score | |
| return scores_dict | |
| def parse_args(): | |
| parser = ArgumentParser(add_help=False) | |
| parser.add_argument('--num_threads', type=int, default=4) | |
| parser.add_argument('--data_path', type=str, default='results') | |
| parser.add_argument('--gt_path', type=str, default='gt_images') | |
| args = parser.parse_args() | |
| return args | |
| def run(args): | |
| file_paths = [] | |
| for f in os.listdir(args.data_path): | |
| image_path = os.path.join(args.data_path, f) | |
| gt_path = os.path.join(args.gt_path, f) | |
| if f.endswith(".jpg") or f.endswith('.png'): | |
| file_paths.append([image_path, gt_path.replace('.png','.jpg')]) | |
| file_chunks = list(chunks(file_paths, int(math.ceil(len(file_paths) / args.num_threads)))) | |
| pool = mp.Pool(args.num_threads) | |
| print('Running on {} paths\nHere we goooo'.format(len(file_paths))) | |
| tic = time.time() | |
| results = pool.map(extract_on_paths, file_chunks) | |
| scores_dict = {} | |
| for d in results: | |
| scores_dict.update(d) | |
| all_scores = list(scores_dict.values()) | |
| mean = np.mean(all_scores) | |
| std = np.std(all_scores) | |
| result_str = 'New Average score is {:.2f}+-{:.2f}'.format(mean, std) | |
| print(result_str) | |
| out_path = os.path.join(os.path.dirname(args.data_path), 'inference_metrics') | |
| if not os.path.exists(out_path): | |
| os.makedirs(out_path) | |
| with open(os.path.join(out_path, 'stat_id.txt'), 'w') as f: | |
| f.write(result_str) | |
| with open(os.path.join(out_path, 'scores_id.json'), 'w') as f: | |
| json.dump(scores_dict, f) | |
| toc = time.time() | |
| print('Mischief managed in {}s'.format(toc - tic)) | |
| if __name__ == '__main__': | |
| args = parse_args() | |
| run(args) | |