FaceNet.py 6.81 KB
Newer Older
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
1
2
3
4
5
6
7
8
from __future__ import division
import os
import re
import logging
import numpy
import tensorflow as tf
from bob.ip.color import gray_to_rgb
from bob.io.image import to_matplotlib
9
from bob.extension import rc
10
11
import bob.extension.download
import bob.io.base
12
import multiprocessing
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
13
14
15

logger = logging.getLogger(__name__)

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
16
17
FACENET_MODELPATH_KEY = "bob.ip.tensorflow_extractor.facenet_modelpath"

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
18
19
20
21
22
23
24
25
26
27
28
29

def prewhiten(img):
    mean = numpy.mean(img)
    std = numpy.std(img)
    std_adj = numpy.maximum(std, 1.0 / numpy.sqrt(img.size))
    y = numpy.multiply(numpy.subtract(img, mean), 1 / std_adj)
    return y


def get_model_filenames(model_dir):
    # code from https://github.com/davidsandberg/facenet
    files = os.listdir(model_dir)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
30
    meta_files = [s for s in files if s.endswith(".meta")]
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
31
    if len(meta_files) == 0:
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
32
        raise ValueError("No meta file found in the model directory (%s)" % model_dir)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
33
34
    elif len(meta_files) > 1:
        raise ValueError(
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
35
36
37
            "There should not be more than one meta file in the model "
            "directory (%s)" % model_dir
        )
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
38
39
40
    meta_file = meta_files[0]
    max_step = -1
    for f in files:
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
41
        step_str = re.match(r"(^model-[\w\- ]+.ckpt-(\d+))", f)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
42
43
44
45
46
47
48
        if step_str is not None and len(step_str.groups()) >= 2:
            step = int(step_str.groups()[1])
            if step > max_step:
                max_step = step
                ckpt_file = step_str.groups()[0]
    return meta_file, ckpt_file

49
_semaphore = multiprocessing.Semaphore()
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
50
51
class FaceNet(object):
    """Wrapper for the free FaceNet variant:
52
53
54
55
56
57
58
59
    https://github.com/davidsandberg/facenet

    To use this class as a bob.bio.base extractor::

        from bob.bio.base.extractor import Extractor
        class FaceNetExtractor(FaceNet, Extractor):
            pass
        extractor = FaceNetExtractor()
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

    And for a preprocessor you can use::

        from bob.bio.face.preprocessor import FaceCrop
        # This is the size of the image that this model expects
        CROPPED_IMAGE_HEIGHT = 160
        CROPPED_IMAGE_WIDTH = 160
        # eye positions for frontal images
        RIGHT_EYE_POS = (46, 53)
        LEFT_EYE_POS = (46, 107)
        # Crops the face using eye annotations
        preprocessor = FaceCrop(
            cropped_image_size=(CROPPED_IMAGE_HEIGHT, CROPPED_IMAGE_WIDTH),
            cropped_positions={'leye': LEFT_EYE_POS, 'reye': RIGHT_EYE_POS},
            color_channel='rgb'
        )

77
    """
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
78

79
    def __init__(
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
80
81
82
83
84
85
        self,
        model_path=rc[FACENET_MODELPATH_KEY],
        image_size=160,
        layer_name="embeddings:0",
        **kwargs
    ):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
86
87
88
        super(FaceNet, self).__init__()
        self.model_path = model_path
        self.image_size = image_size
89
90
91
92
93
        self.layer_name = layer_name        
        self._clean_unpicklables()


    def _clean_unpicklables(self):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
94
95
        self.session = None
        self.embeddings = None
96
97
98
99
100
101
        self.graph = None
        self.images_placeholder = None
        self.embeddings = None
        self.phase_train_placeholder = None
        self.session = None        

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

    def _check_feature(self, img):
        img = numpy.ascontiguousarray(img)
        if img.ndim == 2:
            img = gray_to_rgb(img)
        assert img.shape[-1] == self.image_size
        assert img.shape[-2] == self.image_size
        img = to_matplotlib(img)
        img = prewhiten(img)
        return img[None, ...]

    def load_model(self):
        if self.model_path is None:
            self.model_path = self.get_modelpath()
        if not os.path.exists(self.model_path):
117
            bob.io.base.create_directories_safe(FaceNet.get_modelpath())
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
118
            zip_file = os.path.join(FaceNet.get_modelpath(), "20170512-110547.zip")
119
            urls = [
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
120
121
                # This link only works in Idiap CI to save bandwidth.
                "http://www.idiap.ch/private/wheels/gitlab/"
122
                "facenet_model2_20170512-110547.zip",
123
124
125
                # this link to dropbox would work for everybody
                "https://www.dropbox.com/s/"
                "k7bhxe58q7d48g7/facenet_model2_20170512-110547.zip?dl=1",
126
            ]
127
            bob.extension.download.download_and_unzip(urls, zip_file)
128

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
129
130
        # code from https://github.com/davidsandberg/facenet
        model_exp = os.path.expanduser(self.model_path)
131
        with self.graph.as_default():
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
132
133
134
135
            if os.path.isfile(model_exp):
                logger.info("Model filename: %s" % model_exp)
                with tf.compat.v1.gfile.FastGFile(model_exp, "rb") as f:
                    graph_def = tf.compat.v1.GraphDef()
136
                    graph_def.ParseFromString(f.read())
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
137
                    tf.import_graph_def(graph_def, name="")
138
            else:
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
139
                logger.info("Model directory: %s" % model_exp)
140
141
                meta_file, ckpt_file = get_model_filenames(model_exp)

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
142
143
                logger.info("Metagraph file: %s" % meta_file)
                logger.info("Checkpoint file: %s" % ckpt_file)
144

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
145
146
147
148
                saver = tf.compat.v1.train.import_meta_graph(
                    os.path.join(model_exp, meta_file)
                )
                saver.restore(self.session, os.path.join(model_exp, ckpt_file))
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
149
150
        # Get input and output tensors
        self.images_placeholder = self.graph.get_tensor_by_name("input:0")
151
        self.embeddings = self.graph.get_tensor_by_name(self.layer_name)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
152
        self.phase_train_placeholder = self.graph.get_tensor_by_name("phase_train:0")
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
153
154
155
        logger.info("Successfully loaded the model.")

    def __call__(self, img):
156
157
158
159
160
161
162
163
164
165
166
167
        with _semaphore:
            images = self._check_feature(img)
            if self.session is None:
                self.graph = tf.Graph()
                self.session = tf.compat.v1.Session(graph=self.graph)
            if self.embeddings is None:
                self.load_model()
            feed_dict = {
                self.images_placeholder: images,
                self.phase_train_placeholder: False,
            }
            features = self.session.run(self.embeddings, feed_dict=feed_dict)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
168
169
170
171
        return features.flatten()

    @staticmethod
    def get_modelpath():
172
        """
173
        Get default model path.
174
175

        First we try the to search this path via Global Configuration System.
176
177
        If we can not find it, we set the path in the directory
        `<project>/data`
178
        """
179

180
        # Priority to the RC path
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
181
        model_path = rc["bob.ip.tensorflow_extractor.facenet_modelpath"]
182
183
184

        if model_path is None:
            import pkg_resources
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
185

186
            model_path = pkg_resources.resource_filename(
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
187
188
                __name__, "data/FaceNet/20170512-110547"
            )
189
190

        return model_path
191
192
193
194
195
196
197
198
199
200

    def __setstate__(self, d):
        # Handling unpicklable objects
        self.__dict__ = d

    def __getstate__(self):
        # Handling unpicklable objects
        with _semaphore:
            self._clean_unpicklables()
        return self.__dict__