mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-18 22:57:49 +00:00
204 lines
7.7 KiB
Python
204 lines
7.7 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
from scipy import misc
|
|
import sys
|
|
import os
|
|
import argparse
|
|
import tensorflow as tf
|
|
import numpy as np
|
|
import mxnet as mx
|
|
import random
|
|
import sklearn
|
|
from sklearn.decomposition import PCA
|
|
from time import sleep
|
|
from easydict import EasyDict as edict
|
|
#import facenet
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'align'))
|
|
import detect_face
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'common'))
|
|
import face_image
|
|
import face_preprocess
|
|
|
|
def ch_dev(arg_params, aux_params, ctx):
|
|
new_args = dict()
|
|
new_auxs = dict()
|
|
for k, v in arg_params.items():
|
|
new_args[k] = v.as_in_context(ctx)
|
|
for k, v in aux_params.items():
|
|
new_auxs[k] = v.as_in_context(ctx)
|
|
return new_args, new_auxs
|
|
|
|
def do_flip(data):
|
|
for idx in xrange(data.shape[0]):
|
|
data[idx,:,:] = np.fliplr(data[idx,:,:])
|
|
|
|
class FaceModel:
|
|
def __init__(self, args):
|
|
model = edict()
|
|
with tf.Graph().as_default():
|
|
config = tf.ConfigProto()
|
|
config.gpu_options.per_process_gpu_memory_fraction = 0.2
|
|
sess = tf.Session(config=config)
|
|
#sess = tf.Session()
|
|
with sess.as_default():
|
|
self.pnet, self.rnet, self.onet = detect_face.create_mtcnn(sess, None)
|
|
|
|
self.threshold = args.threshold
|
|
self.det_minsize = 50
|
|
self.det_threshold = [0.4,0.6,0.6]
|
|
self.det_factor = 0.9
|
|
_vec = args.image_size.split(',')
|
|
assert len(_vec)==2
|
|
self.image_size = (int(_vec[0]), int(_vec[1]))
|
|
_vec = args.model.split(',')
|
|
assert len(_vec)==2
|
|
prefix = _vec[0]
|
|
epoch = int(_vec[1])
|
|
print('loading',prefix, epoch)
|
|
self.model = edict()
|
|
self.model.ctx = mx.gpu(args.gpu)
|
|
self.model.sym, self.model.arg_params, self.model.aux_params = mx.model.load_checkpoint(prefix, epoch)
|
|
self.model.arg_params, self.model.aux_params = ch_dev(self.model.arg_params, self.model.aux_params, self.model.ctx)
|
|
all_layers = self.model.sym.get_internals()
|
|
self.model.sym = all_layers['fc1_output']
|
|
|
|
def get_aligned_face(self, img, force = False):
|
|
#print('before det', img.shape)
|
|
bounding_boxes, points = detect_face.detect_face(img, self.det_minsize, self.pnet, self.rnet, self.onet, self.det_threshold, self.det_factor)
|
|
#if bounding_boxes.shape[0]==0:
|
|
# fimg = np.copy(img)
|
|
# do_flip(fimg)
|
|
# bounding_boxes, points = detect_face.detect_face(fimg, self.det_minsize, self.pnet, self.rnet, self.onet, self.det_threshold, self.det_factor)
|
|
if bounding_boxes.shape[0]==0 and force:
|
|
print('force det', img.shape)
|
|
bounding_boxes, points = detect_face.detect_face(img, self.det_minsize, self.pnet, self.rnet, self.onet, [0.3, 0.3, 0.1], self.det_factor)
|
|
#bounding_boxes, points = detect_face.detect_face_force(img, None, self.pnet, self.rnet, self.onet)
|
|
#print('after det')
|
|
if bounding_boxes.shape[0]==0:
|
|
return None
|
|
bindex = 0
|
|
nrof_faces = bounding_boxes.shape[0]
|
|
det = bounding_boxes[:,0:4]
|
|
img_size = np.asarray(img.shape)[0:2]
|
|
if nrof_faces>1:
|
|
bounding_box_size = (det[:,2]-det[:,0])*(det[:,3]-det[:,1])
|
|
img_center = img_size / 2
|
|
offsets = np.vstack([ (det[:,0]+det[:,2])/2-img_center[1], (det[:,1]+det[:,3])/2-img_center[0] ])
|
|
offset_dist_squared = np.sum(np.power(offsets,2.0),0)
|
|
bindex = np.argmax(bounding_box_size-offset_dist_squared*2.0) # some extra weight on the centering
|
|
det = bounding_boxes[:,0:4]
|
|
det = det[bindex,:]
|
|
points = points[:, bindex]
|
|
landmark = points.reshape((2,5)).T
|
|
#points need to be transpose, points = points.reshape( (5,2) ).transpose()
|
|
det = np.squeeze(det)
|
|
bb = det
|
|
points = list(points.flatten())
|
|
assert(len(points)==10)
|
|
str_image_size = "%d,%d"%(self.image_size[0], self.image_size[1])
|
|
warped = face_preprocess.preprocess(img, bbox=bb, landmark = landmark, image_size=str_image_size)
|
|
warped = np.transpose(warped, (2,0,1))
|
|
print(warped.shape)
|
|
return warped
|
|
|
|
def get_all_faces(self, img):
|
|
str_image_size = "%d,%d"%(self.image_size[0], self.image_size[1])
|
|
bounding_boxes, points = detect_face.detect_face(img, self.det_minsize, self.pnet, self.rnet, self.onet, self.det_threshold, self.det_factor)
|
|
ret = []
|
|
for i in xrange(bounding_boxes.shape[0]):
|
|
bbox = bounding_boxes[i,0:4]
|
|
landmark = points[:, i].reshape((2,5)).T
|
|
aligned = face_preprocess.preprocess(img, bbox=bbox, landmark = landmark, image_size=str_image_size)
|
|
aligned = np.transpose(aligned, (2,0,1))
|
|
ret.append(aligned)
|
|
return ret
|
|
|
|
def get_feature_impl(self, face_img, norm):
|
|
embedding = None
|
|
for flipid in [0,1]:
|
|
_img = np.copy(face_img)
|
|
if flipid==1:
|
|
do_flip(_img)
|
|
#nimg = np.zeros(_img.shape, dtype=np.float32)
|
|
#nimg[:,ppatch[1]:ppatch[3],ppatch[0]:ppatch[2]] = _img[:, ppatch[1]:ppatch[3], ppatch[0]:ppatch[2]]
|
|
#_img = nimg
|
|
input_blob = np.expand_dims(_img, axis=0)
|
|
self.model.arg_params["data"] = mx.nd.array(input_blob, self.model.ctx)
|
|
self.model.arg_params["softmax_label"] = mx.nd.empty((1,), self.model.ctx)
|
|
exe = self.model.sym.bind(self.model.ctx, self.model.arg_params ,args_grad=None, grad_req="null", aux_states=self.model.aux_params)
|
|
exe.forward(is_train=False)
|
|
_embedding = exe.outputs[0].asnumpy()
|
|
#print(_embedding.shape)
|
|
if embedding is None:
|
|
embedding = _embedding
|
|
else:
|
|
embedding += _embedding
|
|
if norm:
|
|
embedding = sklearn.preprocessing.normalize(embedding)
|
|
return embedding
|
|
|
|
def get_feature(self, face_img, norm=True):
|
|
#aligned_face = self.get_aligned_face(img, force)
|
|
#if aligned_face is None:
|
|
# return None
|
|
return self.get_feature_impl(face_img, norm)
|
|
|
|
def is_same_id(self, source_img, target_img_list):
|
|
source_face = self.get_aligned_face(source_img, True)
|
|
print('source face', source_face.shape)
|
|
target_face_list = []
|
|
pp = 0
|
|
for img in target_img_list:
|
|
target_force = False
|
|
if pp==len(target_img_list)-1 and len(target_face_list)==0:
|
|
target_force = True
|
|
target_face = self.get_aligned_face(img, target_force)
|
|
if target_face is not None:
|
|
target_face_list.append(target_face)
|
|
pp+=1
|
|
print('target face', len(target_face_list))
|
|
source_feature = self.get_feature(source_face, True)
|
|
target_feature = None
|
|
for target_face in target_face_list:
|
|
_feature = self.get_feature(target_face, False)
|
|
if target_feature is None:
|
|
target_feature = _feature
|
|
else:
|
|
target_feature += _feature
|
|
target_feature = sklearn.preprocessing.normalize(target_feature)
|
|
#sim = np.dot(source_feature, target_feature.T)
|
|
diff = np.subtract(source_feature, target_feature)
|
|
dist = np.sum(np.square(diff),1)
|
|
print('dist', dist)
|
|
#print(sim, dist)
|
|
if dist<=self.threshold:
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def sim(self, source_img, target_img_list):
|
|
print('sim start')
|
|
source_face = self.get_aligned_face(source_img, True)
|
|
print('source face', source_face.shape)
|
|
target_face_list = []
|
|
pp = 0
|
|
for img in target_img_list:
|
|
target_force = False
|
|
if pp==len(target_img_list)-1 and len(target_face_list)==0:
|
|
target_force = True
|
|
target_face = self.get_aligned_face(img, target_force)
|
|
if target_face is not None:
|
|
target_face_list.append(target_face)
|
|
pp+=1
|
|
print('target face', len(target_face_list))
|
|
source_feature = self.get_feature(source_face, True)
|
|
target_feature = None
|
|
sim_list = []
|
|
for target_face in target_face_list:
|
|
_feature = self.get_feature(target_face, True)
|
|
_sim = np.dot(source_feature, _feature.T)
|
|
sim_list.append(_sim)
|
|
return np.max(sim_list)
|