import argparse import os import pickle import timeit import cv2 import mxnet as mx import numpy as np import pandas as pd import prettytable import skimage.transform from sklearn.metrics import roc_curve from sklearn.preprocessing import normalize from onnx_helper import ArcFaceORT SRC = np.array( [ [30.2946, 51.6963], [65.5318, 51.5014], [48.0252, 71.7366], [33.5493, 92.3655], [62.7299, 92.2041]] , dtype=np.float32) SRC[:, 0] += 8.0 class AlignedDataSet(mx.gluon.data.Dataset): def __init__(self, root, lines, align=True): self.lines = lines self.root = root self.align = align def __len__(self): return len(self.lines) def __getitem__(self, idx): each_line = self.lines[idx] name_lmk_score = each_line.strip().split(' ') name = os.path.join(self.root, name_lmk_score[0]) img = cv2.cvtColor(cv2.imread(name), cv2.COLOR_BGR2RGB) landmark5 = np.array([float(x) for x in name_lmk_score[1:-1]], dtype=np.float32).reshape((5, 2)) st = skimage.transform.SimilarityTransform() st.estimate(landmark5, SRC) img = cv2.warpAffine(img, st.params[0:2, :], (112, 112), borderValue=0.0) img_1 = np.expand_dims(img, 0) img_2 = np.expand_dims(np.fliplr(img), 0) output = np.concatenate((img_1, img_2), axis=0).astype(np.float32) output = np.transpose(output, (0, 3, 1, 2)) output = mx.nd.array(output) return output def extract(model_root, dataset): model = ArcFaceORT(model_path=model_root) model.check() feat_mat = np.zeros(shape=(len(dataset), 2 * model.feat_dim)) def batchify_fn(data): return mx.nd.concat(*data, dim=0) data_loader = mx.gluon.data.DataLoader( dataset, 128, last_batch='keep', num_workers=4, thread_pool=True, prefetch=16, batchify_fn=batchify_fn) num_iter = 0 for batch in data_loader: batch = batch.asnumpy() batch = (batch - model.input_mean) / model.input_std feat = model.session.run(model.output_names, {model.input_name: batch})[0] feat = np.reshape(feat, (-1, model.feat_dim * 2)) feat_mat[128 * num_iter: 128 * num_iter + feat.shape[0], :] = feat num_iter += 1 if num_iter % 50 == 0: print(num_iter) return feat_mat def read_template_media_list(path): ijb_meta = pd.read_csv(path, sep=' ', header=None).values templates = ijb_meta[:, 1].astype(np.int) medias = ijb_meta[:, 2].astype(np.int) return templates, medias def read_template_pair_list(path): pairs = pd.read_csv(path, sep=' ', header=None).values t1 = pairs[:, 0].astype(np.int) t2 = pairs[:, 1].astype(np.int) label = pairs[:, 2].astype(np.int) return t1, t2, label def read_image_feature(path): with open(path, 'rb') as fid: img_feats = pickle.load(fid) return img_feats def image2template_feature(img_feats=None, templates=None, medias=None): unique_templates = np.unique(templates) template_feats = np.zeros((len(unique_templates), img_feats.shape[1])) for count_template, uqt in enumerate(unique_templates): (ind_t,) = np.where(templates == uqt) face_norm_feats = img_feats[ind_t] face_medias = medias[ind_t] unique_medias, unique_media_counts = np.unique(face_medias, return_counts=True) media_norm_feats = [] for u, ct in zip(unique_medias, unique_media_counts): (ind_m,) = np.where(face_medias == u) if ct == 1: media_norm_feats += [face_norm_feats[ind_m]] else: # image features from the same video will be aggregated into one feature media_norm_feats += [np.mean(face_norm_feats[ind_m], axis=0, keepdims=True), ] media_norm_feats = np.array(media_norm_feats) template_feats[count_template] = np.sum(media_norm_feats, axis=0) if count_template % 2000 == 0: print('Finish Calculating {} template features.'.format( count_template)) template_norm_feats = normalize(template_feats) return template_norm_feats, unique_templates def verification(template_norm_feats=None, unique_templates=None, p1=None, p2=None): template2id = np.zeros((max(unique_templates) + 1, 1), dtype=int) for count_template, uqt in enumerate(unique_templates): template2id[uqt] = count_template score = np.zeros((len(p1),)) total_pairs = np.array(range(len(p1))) batchsize = 100000 sublists = [total_pairs[i: i + batchsize] for i in range(0, len(p1), batchsize)] total_sublists = len(sublists) for c, s in enumerate(sublists): feat1 = template_norm_feats[template2id[p1[s]]] feat2 = template_norm_feats[template2id[p2[s]]] similarity_score = np.sum(feat1 * feat2, -1) score[s] = similarity_score.flatten() if c % 10 == 0: print('Finish {}/{} pairs.'.format(c, total_sublists)) return score def verification2(template_norm_feats=None, unique_templates=None, p1=None, p2=None): template2id = np.zeros((max(unique_templates) + 1, 1), dtype=int) for count_template, uqt in enumerate(unique_templates): template2id[uqt] = count_template score = np.zeros((len(p1),)) # save cosine distance between pairs total_pairs = np.array(range(len(p1))) batchsize = 100000 # small batchsize instead of all pairs in one batch due to the memory limiation sublists = [total_pairs[i:i + batchsize] for i in range(0, len(p1), batchsize)] total_sublists = len(sublists) for c, s in enumerate(sublists): feat1 = template_norm_feats[template2id[p1[s]]] feat2 = template_norm_feats[template2id[p2[s]]] similarity_score = np.sum(feat1 * feat2, -1) score[s] = similarity_score.flatten() if c % 10 == 0: print('Finish {}/{} pairs.'.format(c, total_sublists)) return score def main(args): use_norm_score = True # if Ture, TestMode(N1) use_detector_score = True # if Ture, TestMode(D1) use_flip_test = True # if Ture, TestMode(F1) assert args.target == 'IJBC' or args.target == 'IJBB' start = timeit.default_timer() templates, medias = read_template_media_list( os.path.join('%s/meta' % args.image_path, '%s_face_tid_mid.txt' % args.target.lower())) stop = timeit.default_timer() print('Time: %.2f s. ' % (stop - start)) start = timeit.default_timer() p1, p2, label = read_template_pair_list( os.path.join('%s/meta' % args.image_path, '%s_template_pair_label.txt' % args.target.lower())) stop = timeit.default_timer() print('Time: %.2f s. ' % (stop - start)) start = timeit.default_timer() img_path = '%s/loose_crop' % args.image_path img_list_path = '%s/meta/%s_name_5pts_score.txt' % (args.image_path, args.target.lower()) img_list = open(img_list_path) files = img_list.readlines() dataset = AlignedDataSet(root=img_path, lines=files, align=True) img_feats = extract(args.model_root, dataset) faceness_scores = [] for each_line in files: name_lmk_score = each_line.split() faceness_scores.append(name_lmk_score[-1]) faceness_scores = np.array(faceness_scores).astype(np.float32) stop = timeit.default_timer() print('Time: %.2f s. ' % (stop - start)) print('Feature Shape: ({} , {}) .'.format(img_feats.shape[0], img_feats.shape[1])) start = timeit.default_timer() if use_flip_test: img_input_feats = img_feats[:, 0:img_feats.shape[1] // 2] + img_feats[:, img_feats.shape[1] // 2:] else: img_input_feats = img_feats[:, 0:img_feats.shape[1] // 2] if use_norm_score: img_input_feats = img_input_feats else: img_input_feats = img_input_feats / np.sqrt(np.sum(img_input_feats ** 2, -1, keepdims=True)) if use_detector_score: print(img_input_feats.shape, faceness_scores.shape) img_input_feats = img_input_feats * faceness_scores[:, np.newaxis] else: img_input_feats = img_input_feats template_norm_feats, unique_templates = image2template_feature( img_input_feats, templates, medias) stop = timeit.default_timer() print('Time: %.2f s. ' % (stop - start)) start = timeit.default_timer() score = verification(template_norm_feats, unique_templates, p1, p2) stop = timeit.default_timer() print('Time: %.2f s. ' % (stop - start)) save_path = os.path.join(args.result_dir, "{}_result".format(args.target)) if not os.path.exists(save_path): os.makedirs(save_path) score_save_file = os.path.join(save_path, "{}.npy".format(args.model_root)) np.save(score_save_file, score) files = [score_save_file] methods = [] scores = [] for file in files: methods.append(os.path.basename(file)) scores.append(np.load(file)) methods = np.array(methods) scores = dict(zip(methods, scores)) x_labels = [10 ** -6, 10 ** -5, 10 ** -4, 10 ** -3, 10 ** -2, 10 ** -1] tpr_fpr_table = prettytable.PrettyTable(['Methods'] + [str(x) for x in x_labels]) for method in methods: fpr, tpr, _ = roc_curve(label, scores[method]) fpr = np.flipud(fpr) tpr = np.flipud(tpr) tpr_fpr_row = [] tpr_fpr_row.append("%s-%s" % (method, args.target)) for fpr_iter in np.arange(len(x_labels)): _, min_index = min( list(zip(abs(fpr - x_labels[fpr_iter]), range(len(fpr))))) tpr_fpr_row.append('%.2f' % (tpr[min_index] * 100)) tpr_fpr_table.add_row(tpr_fpr_row) print(tpr_fpr_table) if __name__ == '__main__': parser = argparse.ArgumentParser(description='do ijb test') # general parser.add_argument('--model-root', default='', help='path to load model.') parser.add_argument('--image-path', default='', type=str, help='') parser.add_argument('--result-dir', default='.', type=str, help='') parser.add_argument('--target', default='IJBC', type=str, help='target, set to IJBC or IJBB') main(parser.parse_args())