Attempt to make the __call__/serialization work

bob/ip/tensorflow_extractor/FaceNet.py

@@ -9,7 +9,6 @@ from bob.io.image import to_matplotlib
 from bob.extension import rc
 import bob.extension.download
 import bob.io.base
-import multiprocessing
 
 logger = logging.getLogger(__name__)
 
@@ -46,7 +45,7 @@ def get_model_filenames(model_dir):
                 ckpt_file = step_str.groups()[0]
     return meta_file, ckpt_file
 
-_semaphore = multiprocessing.Semaphore()
+
 class FaceNet(object):
     """Wrapper for the free FaceNet variant:
     https://github.com/davidsandberg/facenet
@@ -86,19 +85,16 @@ class FaceNet(object):
         super(FaceNet, self).__init__()
         self.model_path = model_path
         self.image_size = image_size
-        self.layer_name = layer_name        
+        self.layer_name = layer_name
+        self.loaded = False
         self._clean_unpicklables()
 
-
     def _clean_unpicklables(self):
         self.session = None
         self.embeddings = None
         self.graph = None
         self.images_placeholder = None
-        self.embeddings = None
         self.phase_train_placeholder = None
-        self.session = None        
-
 
     def _check_feature(self, img):
         img = numpy.ascontiguousarray(img)
@@ -111,6 +107,9 @@ class FaceNet(object):
         return img[None, ...]
 
     def load_model(self):
+        self.graph = tf.Graph()
+        self.session = tf.compat.v1.Session(graph=self.graph)
+
         if self.model_path is None:
             self.model_path = self.get_modelpath()
         if not os.path.exists(self.model_path):
@@ -151,20 +150,19 @@ class FaceNet(object):
         self.embeddings = self.graph.get_tensor_by_name(self.layer_name)
         self.phase_train_placeholder = self.graph.get_tensor_by_name("phase_train:0")
         logger.info("Successfully loaded the model.")
+        self.loaded = True
 
     def __call__(self, img):
-        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)
+        # with _semaphore:
+        images = self._check_feature(img)
+        if not self.loaded:
+            self.load_model()
+
+        feed_dict = {
+            self.images_placeholder: images,
+            self.phase_train_placeholder: False,
+        }
+        features = self.session.run(self.embeddings, feed_dict=feed_dict)
         return features.flatten()
 
     @staticmethod
@@ -192,9 +190,18 @@ class FaceNet(object):
     def __setstate__(self, d):
         # Handling unpicklable objects
         self.__dict__ = d
+        self.load_model()
 
     def __getstate__(self):
         # Handling unpicklable objects
-        with _semaphore:
-            self._clean_unpicklables()
-        return self.__dict__
+        # with _semaphore:
+        # self._clean_unpicklables()
+        self.loaded = False
+        d = self.__dict__
+        d.pop("session") if "session" in d else None
+        d.pop("embeddings") if "embeddings" in d else None
+        d.pop("graph") if "graph" in d else None
+        d.pop("images_placeholder") if "images_placeholder" in d else None
+        d.pop("phase_train_placeholder") if "phase_train_placeholder" in d else None
+
+        return d