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
