initial

facelib/MTCExtractor.py

@@ -0,0 +1,66 @@
+import numpy as np
+import os
+import cv2
+
+from pathlib import Path
+
+from .mtcnn import *
+
+class MTCExtractor(object):   
+    def __init__(self, keras, tf, tf_session):
+        self.scale_to = 1920
+        self.keras = keras
+        self.tf = tf
+        self.tf_session = tf_session
+        
+        self.min_face_size = self.scale_to * 0.042
+        self.thresh1 = 0.7
+        self.thresh2 = 0.85
+        self.thresh3 = 0.6
+        self.scale_factor = 0.95
+        
+        '''
+        self.min_face_size = self.scale_to * 0.042
+        self.thresh1 = 7
+        self.thresh2 = 85
+        self.thresh3 = 6
+        self.scale_factor = 0.95
+        '''
+
+    def __enter__(self):
+        with self.tf.variable_scope('pnet2'):
+            data = self.tf.placeholder(self.tf.float32, (None,None,None,3), 'input')
+            pnet2 = PNet(self.tf, {'data':data})
+            pnet2.load(str(Path(__file__).parent/'det1.npy'), self.tf_session)
+        with self.tf.variable_scope('rnet2'):
+            data = self.tf.placeholder(self.tf.float32, (None,24,24,3), 'input')
+            rnet2 = RNet(self.tf, {'data':data})
+            rnet2.load(str(Path(__file__).parent/'det2.npy'), self.tf_session)
+        with self.tf.variable_scope('onet2'):
+            data = self.tf.placeholder(self.tf.float32, (None,48,48,3), 'input')
+            onet2 = ONet(self.tf, {'data':data})
+            onet2.load(str(Path(__file__).parent/'det3.npy'), self.tf_session)
+
+        self.pnet_fun = self.keras.backend.function([pnet2.layers['data']],[pnet2.layers['conv4-2'], pnet2.layers['prob1']])
+        self.rnet_fun = self.keras.backend.function([rnet2.layers['data']],[rnet2.layers['conv5-2'], rnet2.layers['prob1']])
+        self.onet_fun = self.keras.backend.function([onet2.layers['data']],[onet2.layers['conv6-2'], onet2.layers['conv6-3'], onet2.layers['prob1']])
+
+        faces, pnts = detect_face ( np.zeros ( (self.scale_to, self.scale_to, 3)), self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor ) 
+        return self
+        
+    def __exit__(self, exc_type=None, exc_value=None, traceback=None):
+        return False #pass exception between __enter__ and __exit__ to outter level
+        
+    def extract_from_bgr (self, input_image):
+        input_image = input_image[:,:,::-1].copy()
+        (h, w, ch) = input_image.shape
+
+
+        input_scale = self.scale_to / (w if w > h else h)
+        input_image = cv2.resize (input_image, ( int(w*input_scale), int(h*input_scale) ), interpolation=cv2.INTER_LINEAR)
+
+        detected_faces, pnts = detect_face ( input_image, self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor )
+        detected_faces = [ ( int(face[0]/input_scale), int(face[1]/input_scale), int(face[2]/input_scale), int(face[3]/input_scale)) for face in detected_faces ]
+          
+        return detected_faces
+