RabbitRUI/ruispace
0
1# coding: utf-82 3import os4import pickle5 6import matplotlib7import pandas as pd8 9matplotlib.use('Agg')10import matplotlib.pyplot as plt11import timeit12import sklearn13import argparse14import cv215import numpy as np16import torch17from skimage import transform as trans18from backbones import get_model19from sklearn.metrics import roc_curve, auc20 21from menpo.visualize.viewmatplotlib import sample_colours_from_colourmap22from prettytable import PrettyTable23from pathlib import Path24 25import sys26import warnings27 28sys.path.insert(0, "../")29warnings.filterwarnings("ignore")30 31parser = argparse.ArgumentParser(description='do ijb test')32# general33parser.add_argument('--model-prefix', default='', help='path to load model.')34parser.add_argument('--image-path', default='', type=str, help='')35parser.add_argument('--result-dir', default='.', type=str, help='')36parser.add_argument('--batch-size', default=128, type=int, help='')37parser.add_argument('--network', default='iresnet50', type=str, help='')38parser.add_argument('--job', default='insightface', type=str, help='job name')39parser.add_argument('--target', default='IJBC', type=str, help='target, set to IJBC or IJBB')40args = parser.parse_args()41 42target = args.target43model_path = args.model_prefix44image_path = args.image_path45result_dir = args.result_dir46gpu_id = None47use_norm_score = True # if Ture, TestMode(N1)48use_detector_score = True # if Ture, TestMode(D1)49use_flip_test = True # if Ture, TestMode(F1)50job = args.job51batch_size = args.batch_size52 53 54class Embedding(object):55 def __init__(self, prefix, data_shape, batch_size=1):56 image_size = (112, 112)57 self.image_size = image_size58 weight = torch.load(prefix)59 resnet = get_model(args.network, dropout=0, fp16=False).cuda()60 resnet.load_state_dict(weight)61 model = torch.nn.DataParallel(resnet)62 self.model = model63 self.model.eval()64 src = np.array([65 [30.2946, 51.6963],66 [65.5318, 51.5014],67 [48.0252, 71.7366],68 [33.5493, 92.3655],69 [62.7299, 92.2041]], dtype=np.float32)70 src[:, 0] += 8.071 self.src = src72 self.batch_size = batch_size73 self.data_shape = data_shape74 75 def get(self, rimg, landmark):76 77 assert landmark.shape[0] == 68 or landmark.shape[0] == 578 assert landmark.shape[1] == 279 if landmark.shape[0] == 68:80 landmark5 = np.zeros((5, 2), dtype=np.float32)81 landmark5[0] = (landmark[36] + landmark[39]) / 282 landmark5[1] = (landmark[42] + landmark[45]) / 283 landmark5[2] = landmark[30]84 landmark5[3] = landmark[48]85 landmark5[4] = landmark[54]86 else:87 landmark5 = landmark88 tform = trans.SimilarityTransform()89 tform.estimate(landmark5, self.src)90 M = tform.params[0:2, :]91 img = cv2.warpAffine(rimg,92 M, (self.image_size[1], self.image_size[0]),93 borderValue=0.0)94 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)95 img_flip = np.fliplr(img)96 img = np.transpose(img, (2, 0, 1)) # 3*112*112, RGB97 img_flip = np.transpose(img_flip, (2, 0, 1))98 input_blob = np.zeros((2, 3, self.image_size[1], self.image_size[0]), dtype=np.uint8)99 input_blob[0] = img100 input_blob[1] = img_flip101 return input_blob102 103 @torch.no_grad()104 def forward_db(self, batch_data):105 imgs = torch.Tensor(batch_data).cuda()106 imgs.div_(255).sub_(0.5).div_(0.5)107 feat = self.model(imgs)108 feat = feat.reshape([self.batch_size, 2 * feat.shape[1]])109 return feat.cpu().numpy()110 111 112# 将一个list尽量均分成n份，限制len(list)==n，份数大于原list内元素个数则分配空list[]113def divideIntoNstrand(listTemp, n):114 twoList = [[] for i in range(n)]115 for i, e in enumerate(listTemp):116 twoList[i % n].append(e)117 return twoList118 119 120def read_template_media_list(path):121 # ijb_meta = np.loadtxt(path, dtype=str)122 ijb_meta = pd.read_csv(path, sep=' ', header=None).values123 templates = ijb_meta[:, 1].astype(np.int)124 medias = ijb_meta[:, 2].astype(np.int)125 return templates, medias126 127 128# In[ ]:129 130 131def read_template_pair_list(path):132 # pairs = np.loadtxt(path, dtype=str)133 pairs = pd.read_csv(path, sep=' ', header=None).values134 # print(pairs.shape)135 # print(pairs[:, 0].astype(np.int))136 t1 = pairs[:, 0].astype(np.int)137 t2 = pairs[:, 1].astype(np.int)138 label = pairs[:, 2].astype(np.int)139 return t1, t2, label140 141 142# In[ ]:143 144 145def read_image_feature(path):146 with open(path, 'rb') as fid:147 img_feats = pickle.load(fid)148 return img_feats149 150 151# In[ ]:152 153 154def get_image_feature(img_path, files_list, model_path, epoch, gpu_id):155 batch_size = args.batch_size156 data_shape = (3, 112, 112)157 158 files = files_list159 print('files:', len(files))160 rare_size = len(files) % batch_size161 faceness_scores = []162 batch = 0163 img_feats = np.empty((len(files), 1024), dtype=np.float32)164 165 batch_data = np.empty((2 * batch_size, 3, 112, 112))166 embedding = Embedding(model_path, data_shape, batch_size)167 for img_index, each_line in enumerate(files[:len(files) - rare_size]):168 name_lmk_score = each_line.strip().split(' ')169 img_name = os.path.join(img_path, name_lmk_score[0])170 img = cv2.imread(img_name)171 lmk = np.array([float(x) for x in name_lmk_score[1:-1]],172 dtype=np.float32)173 lmk = lmk.reshape((5, 2))174 input_blob = embedding.get(img, lmk)175 176 batch_data[2 * (img_index - batch * batch_size)][:] = input_blob[0]177 batch_data[2 * (img_index - batch * batch_size) + 1][:] = input_blob[1]178 if (img_index + 1) % batch_size == 0:179 print('batch', batch)180 img_feats[batch * batch_size:batch * batch_size +181 batch_size][:] = embedding.forward_db(batch_data)182 batch += 1183 faceness_scores.append(name_lmk_score[-1])184 185 batch_data = np.empty((2 * rare_size, 3, 112, 112))186 embedding = Embedding(model_path, data_shape, rare_size)187 for img_index, each_line in enumerate(files[len(files) - rare_size:]):188 name_lmk_score = each_line.strip().split(' ')189 img_name = os.path.join(img_path, name_lmk_score[0])190 img = cv2.imread(img_name)191 lmk = np.array([float(x) for x in name_lmk_score[1:-1]],192 dtype=np.float32)193 lmk = lmk.reshape((5, 2))194 input_blob = embedding.get(img, lmk)195 batch_data[2 * img_index][:] = input_blob[0]196 batch_data[2 * img_index + 1][:] = input_blob[1]197 if (img_index + 1) % rare_size == 0:198 print('batch', batch)199 img_feats[len(files) -200 rare_size:][:] = embedding.forward_db(batch_data)201 batch += 1202 faceness_scores.append(name_lmk_score[-1])203 faceness_scores = np.array(faceness_scores).astype(np.float32)204 # img_feats = np.ones( (len(files), 1024), dtype=np.float32) * 0.01205 # faceness_scores = np.ones( (len(files), ), dtype=np.float32 )206 return img_feats, faceness_scores207 208 209# In[ ]:210 211 212def image2template_feature(img_feats=None, templates=None, medias=None):213 # ==========================================================214 # 1. face image feature l2 normalization. img_feats:[number_image x feats_dim]215 # 2. compute media feature.216 # 3. compute template feature.217 # ==========================================================218 unique_templates = np.unique(templates)219 template_feats = np.zeros((len(unique_templates), img_feats.shape[1]))220 221 for count_template, uqt in enumerate(unique_templates):222 223 (ind_t,) = np.where(templates == uqt)224 face_norm_feats = img_feats[ind_t]225 face_medias = medias[ind_t]226 unique_medias, unique_media_counts = np.unique(face_medias,227 return_counts=True)228 media_norm_feats = []229 for u, ct in zip(unique_medias, unique_media_counts):230 (ind_m,) = np.where(face_medias == u)231 if ct == 1:232 media_norm_feats += [face_norm_feats[ind_m]]233 else: # image features from the same video will be aggregated into one feature234 media_norm_feats += [235 np.mean(face_norm_feats[ind_m], axis=0, keepdims=True)236 ]237 media_norm_feats = np.array(media_norm_feats)238 # media_norm_feats = media_norm_feats / np.sqrt(np.sum(media_norm_feats ** 2, -1, keepdims=True))239 template_feats[count_template] = np.sum(media_norm_feats, axis=0)240 if count_template % 2000 == 0:241 print('Finish Calculating {} template features.'.format(242 count_template))243 # template_norm_feats = template_feats / np.sqrt(np.sum(template_feats ** 2, -1, keepdims=True))244 template_norm_feats = sklearn.preprocessing.normalize(template_feats)245 # print(template_norm_feats.shape)246 return template_norm_feats, unique_templates247 248 249# In[ ]:250 251 252def verification(template_norm_feats=None,253 unique_templates=None,254 p1=None,255 p2=None):256 # ==========================================================257 # Compute set-to-set Similarity Score.258 # ==========================================================259 template2id = np.zeros((max(unique_templates) + 1, 1), dtype=int)260 for count_template, uqt in enumerate(unique_templates):261 template2id[uqt] = count_template262 263 score = np.zeros((len(p1),)) # save cosine distance between pairs264 265 total_pairs = np.array(range(len(p1)))266 batchsize = 100000 # small batchsize instead of all pairs in one batch due to the memory limiation267 sublists = [268 total_pairs[i:i + batchsize] for i in range(0, len(p1), batchsize)269 ]270 total_sublists = len(sublists)271 for c, s in enumerate(sublists):272 feat1 = template_norm_feats[template2id[p1[s]]]273 feat2 = template_norm_feats[template2id[p2[s]]]274 similarity_score = np.sum(feat1 * feat2, -1)275 score[s] = similarity_score.flatten()276 if c % 10 == 0:277 print('Finish {}/{} pairs.'.format(c, total_sublists))278 return score279 280 281# In[ ]:282def verification2(template_norm_feats=None,283 unique_templates=None,284 p1=None,285 p2=None):286 template2id = np.zeros((max(unique_templates) + 1, 1), dtype=int)287 for count_template, uqt in enumerate(unique_templates):288 template2id[uqt] = count_template289 score = np.zeros((len(p1),)) # save cosine distance between pairs290 total_pairs = np.array(range(len(p1)))291 batchsize = 100000 # small batchsize instead of all pairs in one batch due to the memory limiation292 sublists = [293 total_pairs[i:i + batchsize] for i in range(0, len(p1), batchsize)294 ]295 total_sublists = len(sublists)296 for c, s in enumerate(sublists):297 feat1 = template_norm_feats[template2id[p1[s]]]298 feat2 = template_norm_feats[template2id[p2[s]]]299 similarity_score = np.sum(feat1 * feat2, -1)300 score[s] = similarity_score.flatten()301 if c % 10 == 0:302 print('Finish {}/{} pairs.'.format(c, total_sublists))303 return score304 305 306def read_score(path):307 with open(path, 'rb') as fid:308 img_feats = pickle.load(fid)309 return img_feats310 311 312# # Step1: Load Meta Data313 314# In[ ]:315 316assert target == 'IJBC' or target == 'IJBB'317 318# =============================================================319# load image and template relationships for template feature embedding320# tid --> template id, mid --> media id321# format:322# image_name tid mid323# =============================================================324start = timeit.default_timer()325templates, medias = read_template_media_list(326 os.path.join('%s/meta' % image_path,327 '%s_face_tid_mid.txt' % target.lower()))328stop = timeit.default_timer()329print('Time: %.2f s. ' % (stop - start))330 331# In[ ]:332 333# =============================================================334# load template pairs for template-to-template verification335# tid : template id, label : 1/0336# format:337# tid_1 tid_2 label338# =============================================================339start = timeit.default_timer()340p1, p2, label = read_template_pair_list(341 os.path.join('%s/meta' % image_path,342 '%s_template_pair_label.txt' % target.lower()))343stop = timeit.default_timer()344print('Time: %.2f s. ' % (stop - start))345 346# # Step 2: Get Image Features347 348# In[ ]:349 350# =============================================================351# load image features352# format:353# img_feats: [image_num x feats_dim] (227630, 512)354# =============================================================355start = timeit.default_timer()356img_path = '%s/loose_crop' % image_path357img_list_path = '%s/meta/%s_name_5pts_score.txt' % (image_path, target.lower())358img_list = open(img_list_path)359files = img_list.readlines()360# files_list = divideIntoNstrand(files, rank_size)361files_list = files362 363# img_feats364# for i in range(rank_size):365img_feats, faceness_scores = get_image_feature(img_path, files_list,366 model_path, 0, gpu_id)367stop = timeit.default_timer()368print('Time: %.2f s. ' % (stop - start))369print('Feature Shape: ({} , {}) .'.format(img_feats.shape[0],370 img_feats.shape[1]))371 372# # Step3: Get Template Features373 374# In[ ]:375 376# =============================================================377# compute template features from image features.378# =============================================================379start = timeit.default_timer()380# ==========================================================381# Norm feature before aggregation into template feature?382# Feature norm from embedding network and faceness score are able to decrease weights for noise samples (not face).383# ==========================================================384# 1. FaceScore （Feature Norm）385# 2. FaceScore （Detector）386 387if use_flip_test:388 # concat --- F1389 # img_input_feats = img_feats390 # add --- F2391 img_input_feats = img_feats[:, 0:img_feats.shape[1] //392 2] + img_feats[:, img_feats.shape[1] // 2:]393else:394 img_input_feats = img_feats[:, 0:img_feats.shape[1] // 2]395 396if use_norm_score:397 img_input_feats = img_input_feats398else:399 # normalise features to remove norm information400 img_input_feats = img_input_feats / np.sqrt(401 np.sum(img_input_feats ** 2, -1, keepdims=True))402 403if use_detector_score:404 print(img_input_feats.shape, faceness_scores.shape)405 img_input_feats = img_input_feats * faceness_scores[:, np.newaxis]406else:407 img_input_feats = img_input_feats408 409template_norm_feats, unique_templates = image2template_feature(410 img_input_feats, templates, medias)411stop = timeit.default_timer()412print('Time: %.2f s. ' % (stop - start))413 414# # Step 4: Get Template Similarity Scores415 416# In[ ]:417 418# =============================================================419# compute verification scores between template pairs.420# =============================================================421start = timeit.default_timer()422score = verification(template_norm_feats, unique_templates, p1, p2)423stop = timeit.default_timer()424print('Time: %.2f s. ' % (stop - start))425 426# In[ ]:427save_path = os.path.join(result_dir, args.job)428# save_path = result_dir + '/%s_result' % target429 430if not os.path.exists(save_path):431 os.makedirs(save_path)432 433score_save_file = os.path.join(save_path, "%s.npy" % target.lower())434np.save(score_save_file, score)435 436# # Step 5: Get ROC Curves and TPR@FPR Table437 438# In[ ]:439 440files = [score_save_file]441methods = []442scores = []443for file in files:444 methods.append(Path(file).stem)445 scores.append(np.load(file))446 447methods = np.array(methods)448scores = dict(zip(methods, scores))449colours = dict(450 zip(methods, sample_colours_from_colourmap(methods.shape[0], 'Set2')))451x_labels = [10 ** -6, 10 ** -5, 10 ** -4, 10 ** -3, 10 ** -2, 10 ** -1]452tpr_fpr_table = PrettyTable(['Methods'] + [str(x) for x in x_labels])453fig = plt.figure()454for method in methods:455 fpr, tpr, _ = roc_curve(label, scores[method])456 roc_auc = auc(fpr, tpr)457 fpr = np.flipud(fpr)458 tpr = np.flipud(tpr) # select largest tpr at same fpr459 plt.plot(fpr,460 tpr,461 color=colours[method],462 lw=1,463 label=('[%s (AUC = %0.4f %%)]' %464 (method.split('-')[-1], roc_auc * 100)))465 tpr_fpr_row = []466 tpr_fpr_row.append("%s-%s" % (method, target))467 for fpr_iter in np.arange(len(x_labels)):468 _, min_index = min(469 list(zip(abs(fpr - x_labels[fpr_iter]), range(len(fpr)))))470 tpr_fpr_row.append('%.2f' % (tpr[min_index] * 100))471 tpr_fpr_table.add_row(tpr_fpr_row)472plt.xlim([10 ** -6, 0.1])473plt.ylim([0.3, 1.0])474plt.grid(linestyle='--', linewidth=1)475plt.xticks(x_labels)476plt.yticks(np.linspace(0.3, 1.0, 8, endpoint=True))477plt.xscale('log')478plt.xlabel('False Positive Rate')479plt.ylabel('True Positive Rate')480plt.title('ROC on IJB')481plt.legend(loc="lower right")482fig.savefig(os.path.join(save_path, '%s.pdf' % target.lower()))483print(tpr_fpr_table)484