def infer_numpy(self, faces, target_embs): ''' faces : list of PIL Image target_embs : [n, 512] computed embeddings of faces in facebank names : recorded names of faces in facebank tta : test time augmentation (hfilp, that's all) ''' embs = [] for img in faces: if self.tta: with torch.no_grad(): mirror = trans.functional.hflip(img) emb = self.model(self.test_transform(img).to(self.conf.device).unsqueeze(0)) emb_mirror = self.model(self.test_transform(mirror).to(self.conf.device).unsqueeze(0)) embs.append(l2_norm(emb + emb_mirror).data.cpu().numpy()) else: with torch.no_grad(): embs.append(self.model(self.test_transform(img).to(self.conf.device).unsqueeze(0)).data.cpu().numpy()) source_embs = np.array(embs) diff = source_embs - np.expand_dims(target_embs, 0) dist = np.sum(np.power(diff, 2), axis=2) minimum = np.amin(dist, axis=1) min_idx = np.argmin(dist, axis=1) # minimum, min_idx = torch.min(dist, dim=1) min_idx[minimum > self.threshold] = -1 # if no match, set idx to -1 return min_idx, minimum, source_embs
def infer_tensor(self, faces, target_embs): ''' faces : list of PIL Image target_embs : [n, 512] computed embeddings of faces in facebank names : recorded names of faces in facebank tta : test time augmentation (hfilp, that's all) ''' embs = [] for img in faces: if self.tta: with torch.no_grad(): mirror = trans.functional.hflip(img) emb = self.model( self.test_transform(img).to( self.conf.device).unsqueeze(0)) emb_mirror = self.model( self.test_transform(mirror).to( self.conf.device).unsqueeze(0)) embs.append(l2_norm(emb + emb_mirror)) else: with torch.no_grad(): embs.append( self.model( self.test_transform(img).to( self.conf.device).unsqueeze(0))) source_embs = torch.cat(embs) diff = source_embs.unsqueeze(-1) - target_embs.transpose( 1, 0).unsqueeze(0).to(self.conf.device) dist = torch.sum(torch.pow(diff, 2), dim=1) minimum, min_idx = torch.min(dist, dim=1) min_idx[minimum > self.threshold] = -1 # if no match, set idx to -1 return min_idx, minimum, source_embs
def load_single_face(self, image, name='Unknow'): embeddings = [] names = ['Unknown'] embs = [] assert not image is None, 'None is not image, please enter image path!' try: if isinstance(image, np.ndarray): img = Image.fromarray(image) elif isinstance(image, str): assert os.path.isfile(image), 'No such image name: %s' % image img = Image.open(image) else: img = image except: pass if img.size != (112, 112): img = self.mtcnn.align(img) with torch.no_grad(): if self.tta: mirror = trans.functional.hflip(img) emb = self.model( self.test_transform(img).to(self.conf.device).unsqueeze(0)) emb_mirror = self.model( self.test_transform(mirror).to( self.conf.device).unsqueeze(0)) if self.use_tensor: embs.append(l2_norm(emb + emb_mirror)) else: embs.append(l2_norm(emb + emb_mirror).data.cpu().numpy()) else: embs.append( self.model( self.test_transform(img).to( self.conf.device).unsqueeze(0))) if not len(embs) == 0: names.append(name) names = np.array(names) if self.use_tensor: embedding = torch.cat(embs).mean(0, keepdim=True) embeddings.append(embedding) embeddings = torch.cat(embeddings) else: embedding = np.mean(embs, axis=0) embeddings.append(embedding[0]) return embeddings, names
def prepare_facebank_np(conf, model, mtcnn, tta=True): model.eval() embeddings = [] names = ['Unknown'] for path in Path(conf.facebank_path).iterdir(): if path.is_file(): continue else: embs = [] for file in path.iterdir(): if not file.is_file(): continue else: try: img = Image.open(file) image = np.array(img) if image.shape[2] > 3: img = Image.fromarray(image[..., :3]) except: continue if img.size != (112, 112): img = mtcnn.align(img) if img is None: continue with torch.no_grad(): if tta: mirror = trans.functional.hflip(img) emb = model( conf.test_transform(img).to( conf.device).unsqueeze(0)) emb_mirror = model( conf.test_transform(mirror).to( conf.device).unsqueeze(0)) embs.append( l2_norm(emb + emb_mirror).data.cpu().numpy()) else: embs.append( model( conf.test_transform(img).to( conf.device).unsqueeze(0))) if len(embs) == 0: continue embedding = np.mean(embs, axis=0) embeddings.append(embedding[0]) names.append(path.name) embeddings = np.array(embeddings) names = np.array(names) torch.save(embeddings, '%s/facebank.pth' % conf.facebank_path) np.save('%s/names' % conf.facebank_path, names) return embeddings, names