FaceDetector.py 6.76 KB
Newer Older
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
1
2
3
4
5
6
7
8
9
10
11
#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
# @author: Tiago de Freitas Pereira <tiago.pereira@idiap.ch>

import numpy
import os
import bob.core

logger = bob.core.log.setup("bob.ip.mtcnn")
bob.core.log.set_verbosity_level(logger, 3)
import os
12
os.environ['GLOG_minloglevel'] = '2'
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
13
14
15
16
17
18
19
import caffe
import bob.ip.base

from bob.ip.facedetect import BoundingBox
from .legacy import detect_face
from .utils import bob_to_dlib_image_convertion, rectangle2bounding_box2

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
20

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
21
22
class FaceDetector(object):
    """
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
23
24
25
    Detects face and 5 landmarks using the MTCNN
    (https://github.com/kpzhang93/MTCNN_face_detection_alignment) from the
    paper.
26

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
27
28
29
    Zhang, Kaipeng, et al. "Joint face detection and alignment using multitask
    cascaded convolutional networks." IEEE Signal Processing Letters 23.10
    (2016): 1499-1503.
30

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
31
32
33
34
35
36
    """

    def __init__(self):
        """
        Load the caffe models
        """
37

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
38
39
40
41
42
43
44
        caffe_base_path = FaceDetector.get_mtcnn_model_path()

        # Default value from the example
        self.minsize = 20
        self.threshold = [0.6, 0.7, 0.7]
        self.factor = 0.709
        self.fastresize = False
45

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
46
47
        # Loading the models
        caffe.set_mode_cpu()
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
48
49
50
51
52
53
54
55
56
        self.p_net = caffe.Net(
            os.path.join(caffe_base_path, "det1.prototxt"),
            os.path.join(caffe_base_path, "det1.caffemodel"), caffe.TEST)
        self.r_net = caffe.Net(
            os.path.join(caffe_base_path, "det2.prototxt"),
            os.path.join(caffe_base_path, "det2.caffemodel"), caffe.TEST)
        self.o_net = caffe.Net(
            os.path.join(caffe_base_path, "det3.prototxt"),
            os.path.join(caffe_base_path, "det3.caffemodel"), caffe.TEST)
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
57
58
59
60
61
62
63

    def _convert_list_to_landmarks(self, points):
        """
        Convert the list to 10 landmarks to a dictionary with the points
        """

        landmarks = []
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
64
65
        possible_landmarks = ['reye', 'leye',
                              'nose', 'mouthleft', 'mouthright']
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
66
        for i in range(points.shape[0]):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
67
68
69
70
71
            landmark = dict()
            for offset, p in enumerate(possible_landmarks):
                landmark[p] = (int(points[i][offset + 5]),
                               int(points[i][offset]))
            landmarks.append(landmark)
72
73

        return landmarks
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
74

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
75
    def detect_all_faces(self, image, return_bob_bb=True):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
76
        """
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
        Detect all the faces with its respective landmarks, if any, in a
        COLORED image

        Parameters
        ----------
        image : numpy.array
            The color image [c, w, h]
        return_bob_bb : bool, optional
            If true, will return faces wrapped using
            :any:`bob.ip.facedetect.BoundingBox`.

        Returns
        -------
        object
            Returns two lists; the first on contains the bounding boxes with
            the detected faces and the second one contains list with the faces
            landmarks. The CNN returns 5 facial landmarks (leye, reye, nose,
            mouthleft, mouthright). If there's no face, `None` will be returned

        Raises
        ------
        ValueError
            When image.ndim is not 3.
100

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
101
102
        """
        assert image is not None
103

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
104
        if len(image.shape) != 3:
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
105
            raise ValueError("Only color images is supported")
106

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
107
108
        bb, landmarks = detect_face(bob_to_dlib_image_convertion(
            image), self.minsize, self.p_net, self.r_net, self.o_net, self.threshold, self.fastresize, self.factor)
109

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
110
111
112
        # If there's no face, return none
        if len(bb) == 0:
            return None, None
113

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
114
115
116
117
        if return_bob_bb:
            bb = rectangle2bounding_box2(bb)

        return bb, self._convert_list_to_landmarks(landmarks)
118

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
119
120
121
122
    def detect_single_face(self, image):
        """
        Returns the biggest face in a COLORED image, if any.

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
123
124
125
126
        Parameters
        ----------
        image : numpy.array
            numpy array with color image [c, w, h]
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
127

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
128
129
130
131
        Returns
        -------
        The face bounding box and its respective 5 landmarks (leye, reye, nose,
        mouthleft, mouthright). If there's no face, `None` will be returned
132

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
133
134
135
136
137
138
        """

        faces, landmarks = self.detect_all_faces(image)
        # Return None if
        if faces is None:
            return None, None
139

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
140
141
        index = numpy.argmax([(f.bottomright[0] - f.topleft[0])
                              * (f.bottomright[1] - f.topleft[1]) for f in faces])
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
142
143
144

        return faces[index], landmarks[index]

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
145
    def detect_crop_align(self, image, final_image_size=(160, 160)):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
146
147
148
        """
        Detects the biggest face and crop it based in the eyes location
        using py:class:`bob.ip.base.FaceEyesNorm`.
149

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
150
151
        Final eyes location was inspired here: https://gitlab.idiap.ch/bob/bob.bio.caffe_face/blob/master/bob/bio/caffe_face/config/preprocessor/vgg_preprocessor.py

152
        **Parameters**
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
153
154
155
156
157
           image: numpy array with color image [c, w, h]
           final_image_size: Image dimensions [w, h]

        **Returns**
           The cropped image. If there's no face, `None` will be returned
158

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
159
160
161
        """

        face, landmark = self.detect_single_face(image)
162

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
163
164
165
166
167
        if face is None:
            return None

        CROPPED_IMAGE_WIDTH = final_image_size[0]
        CROPPED_IMAGE_HEIGHT = final_image_size[1]
168

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
169
        # final image position w.r.t the image size
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
170
171
172
173
        RIGHT_EYE_POS = (CROPPED_IMAGE_HEIGHT / 3.44,
                         CROPPED_IMAGE_WIDTH / 3.02)
        LEFT_EYE_POS = (CROPPED_IMAGE_HEIGHT / 3.44,
                        CROPPED_IMAGE_WIDTH / 1.49)
174

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
175
176
        extractor = bob.ip.base.FaceEyesNorm(
            (CROPPED_IMAGE_HEIGHT, CROPPED_IMAGE_WIDTH), RIGHT_EYE_POS, LEFT_EYE_POS)
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
177
178
        return extractor(image, landmark['reye'], landmark['leye'])

179
    def detect_crop(self, image, final_image_size=(182, 182), margin=44):
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
180
181
        """
        Detects the biggest face and crop it
182
183

        **Parameters**
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
184
185
186
187
188
           image: numpy array with color image [c, w, h]
           final_image_size: Image dimensions [w, h]

        **Returns**
           The cropped image. If there's no face, `None` will be returned
189

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
190
191
192
        """

        face, landmark = self.detect_single_face(image)
193

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
194
195
196
        if face is None:
            return None

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
197
198
        top = numpy.uint(numpy.maximum(face.top - margin / 2, 0))
        left = numpy.uint(numpy.maximum(face.left - margin / 2, 0))
199

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
200
201
202
203
        bottom = numpy.uint(numpy.minimum(
            face.bottom + margin / 2, image.shape[1]))
        right = numpy.uint(numpy.minimum(
            face.right + margin / 2, image.shape[2]))
204
205
206
207
208
209

        cropped = image[:, top:bottom, left:right]

        dst = numpy.zeros(shape=(3, final_image_size[0], final_image_size[1]))
        bob.ip.base.scale(cropped, dst)
        return dst
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
210

Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
211
212
213
214
    @staticmethod
    def get_mtcnn_model_path():
        import pkg_resources
        return pkg_resources.resource_filename(__name__, 'data')