FaceRecog/face_recognition.py at main · Lixtt/FaceRecog · GitHub

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
126
127
128
129
130
131
132
# -*-coding: utf-8 -*-

import facenet
import tensorflow as tf
import align.detect_face as detect_face
import numpy as np


class facenetEmbedding:
    def __init__(self, model_path):
        tf.compat.v1.disable_eager_execution()  # 保证sess.run()能够正常运行
        self.sess = tf.compat.v1.InteractiveSession()
        self.sess.run(tf.compat.v1.global_variables_initializer())
        # Load the model
        facenet.load_model(model_path)
        # Get input and output tensors
        self.images_placeholder = tf.compat.v1.get_default_graph().get_tensor_by_name(
            "input:0"
        )
        self.tf_embeddings = tf.compat.v1.get_default_graph().get_tensor_by_name(
            "embeddings:0"
        )
        self.phase_train_placeholder = (
            tf.compat.v1.get_default_graph().get_tensor_by_name("phase_train:0")
        )

    def get_embedding(self, images):
        feed_dict = {
            self.images_placeholder: images,
            self.phase_train_placeholder: False,
        }
        embedding = self.sess.run(self.tf_embeddings, feed_dict=feed_dict)
        return embedding

    def free(self):
        self.sess.close()


class Facedetection:
    def __init__(self):
        self.minsize = 30  # minimum size of face
        self.threshold = [0.6, 0.7, 0.7]  # three steps's threshold
        self.factor = 0.709  # scale factor
        print("Creating networks and loading parameters")
        with tf.Graph().as_default():
            # gpu_memory_fraction = 1.0
            # 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))
            sess = tf.compat.v1.Session()
            with sess.as_default():
                self.pnet, self.rnet, self.onet = detect_face.create_mtcnn(sess, None)

    def detect_face(self, image, fixed=None):
        """
        mtcnn人脸检测，
        PS：人脸检测获得bboxes并不一定是正方形的矩形框，参数fixed指定等宽或者等高的bboxes
        :param image:
        :param fixed:
        :return:
        """
        bboxes, landmarks = detect_face.detect_face(
            image,
            self.minsize,
            self.pnet,
            self.rnet,
            self.onet,
            self.threshold,
            self.factor,
        )
        landmarks_list = []
        landmarks = np.transpose(landmarks)
        bboxes = bboxes.astype(int)
        bboxes = [b[:4] for b in bboxes]
        for landmark in landmarks:
            face_landmarks = [[landmark[j], landmark[j + 5]] for j in range(5)]
            landmarks_list.append(face_landmarks)
        if fixed is not None:
            bboxes, landmarks_list = self.get_square_bboxes(
                bboxes, landmarks_list, fixed
            )
        return bboxes, landmarks_list

    def get_square_bboxes(self, bboxes, landmarks, fixed="height"):
        """
        获得等宽或者等高的bboxes
        :param bboxes:
        :param landmarks:
        :param fixed: width or height
        :return:
        """
        new_bboxes = []
        for bbox in bboxes:
            x1, y1, x2, y2 = bbox
            w = x2 - x1
            h = y2 - y1
            center_x, center_y = (int((x1 + x2) / 2), int((y1 + y2) / 2))
            if fixed == "height":
                dd = h / 2
            elif fixed == "width":
                dd = w / 2
            x11 = int(center_x - dd)
            y11 = int(center_y - dd)
            x22 = int(center_x + dd)
            y22 = int(center_y + dd)
            new_bbox = (x11, y11, x22, y22)
            new_bboxes.append(new_bbox)
        return new_bboxes, landmarks


def detection_face(img):
    minsize = 20  # minimum size of face
    threshold = [0.6, 0.7, 0.7]  # three steps's threshold
    factor = 0.709  # scale factor
    print("Creating networks and loading parameters")
    with tf.Graph().as_default():
        # gpu_memory_fraction = 1.0
        # 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))
        sess = tf.compat.v1.Session()
        with sess.as_default():
            pnet, rnet, onet = detect_face.create_mtcnn(sess, None)
            bboxes, landmarks = detect_face.detect_face(
                img, minsize, pnet, rnet, onet, threshold, factor
            )
    landmarks = np.transpose(landmarks)
    bboxes = bboxes.astype(int)
    bboxes = [b[:4] for b in bboxes]
    landmarks_list = []
    for landmark in landmarks:
        face_landmarks = [[landmark[j], landmark[j + 5]] for j in range(5)]
        landmarks_list.append(face_landmarks)
    return bboxes, landmarks_list