forked from deepinsight/insightface
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathface_embedding.py
94 lines (86 loc) · 2.96 KB
/
face_embedding.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
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 cv2
import sklearn
from sklearn.decomposition import PCA
from time import sleep
from easydict import EasyDict as edict
from mtcnn_detector import MtcnnDetector
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'src', 'common'))
import face_image
import face_preprocess
def do_flip(data):
for idx in range(data.shape[0]):
data[idx,:,:] = np.fliplr(data[idx,:,:])
class FaceModel:
def __init__(self, args):
self.args = args
model = edict()
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
image_size = (int(_vec[0]), int(_vec[1]))
self.image_size = image_size
_vec = args.model.split(',')
assert len(_vec)==2
prefix = _vec[0]
epoch = int(_vec[1])
print('loading',prefix, epoch)
ctx = mx.gpu(args.gpu)
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
all_layers = sym.get_internals()
sym = all_layers['fc1_output']
model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)
#model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,))])
model.bind(data_shapes=[('data', (1, 3, image_size[0], image_size[1]))])
model.set_params(arg_params, aux_params)
self.model = model
mtcnn_path = os.path.join(os.path.dirname(__file__), 'mtcnn-model')
detector = MtcnnDetector(model_folder=mtcnn_path, ctx=ctx, num_worker=1, accurate_landmark = True, threshold=[0.0,0.0,0.2])
self.detector = detector
def get_feature(self, face_img):
#face_img is bgr image
ret = self.detector.detect_face_limited(face_img, det_type = self.args.det)
if ret is None:
return None
bbox, points = ret
if bbox.shape[0]==0:
return None
bbox = bbox[0,0:4]
points = points[0,:].reshape((2,5)).T
#print(bbox)
#print(points)
nimg = face_preprocess.preprocess(face_img, bbox, points, image_size='112,112')
nimg = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
aligned = np.transpose(nimg, (2,0,1))
#print(nimg.shape)
embedding = None
for flipid in [0,1]:
if flipid==1:
if self.args.flip==0:
break
do_flip(aligned)
input_blob = np.expand_dims(aligned, axis=0)
data = mx.nd.array(input_blob)
db = mx.io.DataBatch(data=(data,))
self.model.forward(db, is_train=False)
_embedding = self.model.get_outputs()[0].asnumpy()
#print(_embedding.shape)
if embedding is None:
embedding = _embedding
else:
embedding += _embedding
embedding = sklearn.preprocessing.normalize(embedding).flatten()
return embedding