-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathface.py
125 lines (92 loc) · 4.21 KB
/
face.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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
import pickle
import os
import cv2
import numpy as np
import tensorflow as tf
from scipy import misc
import align.detect_face
import facenet
gpu_memory_fraction = 0.3
facenet_model_checkpoint = os.path.dirname(__file__) + "/../model_checkpoints/20170512-110547"
classifier_model = os.path.dirname(__file__) + "/../model_checkpoints/my_classifier_1.pkl"
debug = False
class Face:
def __init__(self):
self.name = None
self.bounding_box = None
self.image = None
self.container_image = None
self.embedding = None
class Recognition:
def __init__(self):
self.detect = Detection()
self.encoder = Encoder()
self.identifier = Identifier()
def add_identity(self, image, person_name):
faces = self.detect.find_faces(image)
if len(faces) == 1:
face = faces[0]
face.name = person_name
face.embedding = self.encoder.generate_embedding(face)
return faces
def identify(self, image):
faces = self.detect.find_faces(image)
for i, face in enumerate(faces):
if debug:
cv2.imshow("Face: " + str(i), face.image)
face.embedding = self.encoder.generate_embedding(face)
face.name = self.identifier.identify(face)
return faces
class Identifier:
def __init__(self):
with open(classifier_model, 'rb') as infile:
self.model, self.class_names = pickle.load(infile)
def identify(self, face):
if face.embedding is not None:
predictions = self.model.predict_proba([face.embedding])
best_class_indices = np.argmax(predictions, axis=1)
return self.class_names[best_class_indices[0]]
class Encoder:
def __init__(self):
self.sess = tf.Session()
with self.sess.as_default():
facenet.load_model(facenet_model_checkpoint)
def generate_embedding(self, face):
images_placeholder = tf.get_default_graph().get_tensor_by_name("input:0")
embeddings = tf.get_default_graph().get_tensor_by_name("embeddings:0")
phase_train_placeholder = tf.get_default_graph().get_tensor_by_name("phase_train:0")
prewhiten_face = facenet.prewhiten(face.image)
feed_dict = {images_placeholder: [prewhiten_face], phase_train_placeholder: False}
return self.sess.run(embeddings, feed_dict=feed_dict)[0]
class Detection:
minsize = 20
threshold = [0.6, 0.7, 0.7]
factor = 0.709 # scale factor
def __init__(self, face_crop_size=160, face_crop_margin=32):
self.pnet, self.rnet, self.onet = self._setup_mtcnn()
self.face_crop_size = face_crop_size
self.face_crop_margin = face_crop_margin
def _setup_mtcnn(self):
with tf.Graph().as_default():
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_memory_fraction)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False))
with sess.as_default():
return align.detect_face.create_mtcnn(sess, None)
def find_faces(self, image):
faces = []
bounding_boxes, _ = align.detect_face.detect_face(image, self.minsize,
self.pnet, self.rnet, self.onet,
self.threshold, self.factor)
for bb in bounding_boxes:
face = Face()
face.container_image = image
face.bounding_box = np.zeros(4, dtype=np.int32)
img_size = np.asarray(image.shape)[0:2]
face.bounding_box[0] = np.maximum(bb[0] - self.face_crop_margin / 2, 0)
face.bounding_box[1] = np.maximum(bb[1] - self.face_crop_margin / 2, 0)
face.bounding_box[2] = np.minimum(bb[2] + self.face_crop_margin / 2, img_size[1])
face.bounding_box[3] = np.minimum(bb[3] + self.face_crop_margin / 2, img_size[0])
cropped = image[face.bounding_box[1]:face.bounding_box[3], face.bounding_box[0]:face.bounding_box[2], :]
face.image = misc.imresize(cropped, (self.face_crop_size, self.face_crop_size), interp='bilinear')
faces.append(face)
return faces