Merge pull request #1 from MacherLabs/dev

added mobilenet support

README.md

@@ -7,6 +7,7 @@ A face interface class implementing different face detection, alignment and reco
 * Face Detector (detector_method='dlib')
 * Dlib CNN Face Detector (detector_method='cnn')
 * OpenCV Face Detector (detector_method='opencv')
+* Mobilenet Face Detector (detector_method='mobilenet')
 * Dlib Face Recognition (recognition_method='dlib')
 * Dlib Facial Landmarks (predictor_model='small' for 5 face landmarks)
 
@@ -33,7 +34,6 @@ pip install --user git+ssh://[email protected]/macherlabs/facelib.git
 
     import face, cv2
     facedemo = face.Face(detector_method='dlib')
-
     image_url1 = 'test.png'
     image_url2 = 'test2.png'
 
@@ -43,3 +43,13 @@ pip install --user git+ssh://[email protected]/macherlabs/facelib.git
     if imgcv1 is not None and imgcv2 is not None:
         results = facedemo.compare(imgcv1, imgcv2)
         print results# facelib
+
+### Mobilenet usage ###
+
+    faceDetector = Face(detector_method='mobilenet',
+             detector_model='mobilenet_300_frozen_trt_inference_graph_face.pb', # Default mobilenet_512_frozen_inference_graph_face.pb
+             recognition_method=None ,
+             gpu_frac=0.3,
+             trt_enable=False,# converts graph to tensorrt optimized graph
+             precision='FP16' # Defalut 'FP32', use 'FP16'only if gpu support is available
+             )

face/__init__.py

@@ -1 +1 @@
-from face import Face
+from .face import Face

face/detect_face.py

@@ -2,6 +2,7 @@
 
 import cv2
 import os
+import numpy as np
 
 WORK_DIR = "/LFS/facelib" #os.path.dirname(os.path.abspath(__file__))
 MODEL_DIR = 'models'
@@ -106,7 +107,7 @@ def __init__(self, model_loc=None):
         self._detector = dlib.cnn_face_detection_model_v1(model_loc)
 
     def detect_raw(self, imgcv, **kwargs):
-        upsamples = kwargs.get('upsamples', 1)
+        upsamples = kwargs.get('upsamples', 0)
         return self._detector(imgcv, upsamples)
 
     def detect(self, imgcv, **kwargs):
@@ -155,6 +156,103 @@ def _format_result(self, result):
             out_list.append(formatted_res)
         return out_list
 
+class FaceDetectorMobilenet(object):
+    def __init__(self, model_name='mobilenet_512_frozen_inference_graph_face.pb',trt_enable=False,precision ='FP32',gpu_frac=0.3):
+        import tensorflow as tf
+        import tensorflow.contrib.tensorrt as trt
+        if model_name==None:
+            model_name='mobilenet_512_frozen_inference_graph_face.pb'
+        model_loc = os.path.join(WORK_DIR, MODEL_DIR, model_name)
+        self.detection_graph = tf.Graph()
+
+        with self.detection_graph.as_default():
+            od_graph_def = tf.GraphDef()
+            with tf.gfile.GFile(model_loc, 'rb') as fid:
+                serialized_graph = fid.read()
+                od_graph_def.ParseFromString(serialized_graph)
+                # Convert to tensorrt graph if asked
+                if trt_enable == True:
+                    trt_graph_def=trt.create_inference_graph(input_graph_def= od_graph_def,
+                                                max_batch_size=1,
+                                                max_workspace_size_bytes=1<<20,
+                                                precision_mode=precision,
+                                                minimum_segment_size=5,
+                                                maximum_cached_engines=5,
+                                                outputs=['detection_boxes','detection_scores','detection_classes','num_detections'])
+
+                    tf.import_graph_def(trt_graph_def, name='')
+                else:
+                    tf.import_graph_def(od_graph_def, name='')
+
+        with self.detection_graph.as_default():
+            config = tf.ConfigProto()
+            # if trt_enable == True:
+            #     config.gpu_options.allow_growth = True
+            # else:
+            config.gpu_options.per_process_gpu_memory_fraction = gpu_frac
+            self.sess = tf.Session(graph=self.detection_graph, config=config)
+            self.windowNotSet = True
+
+
+    def run(self, image):
+        """image: bgr image
+        return (boxes, scores, classes, num_detections)
+        """
+        image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        # the array based representation of the image will be used later in order to prepare the
+        # result image with boxes and labels on it.
+        # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
+        image_np_expanded = np.expand_dims(image_np, axis=0)
+        image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
+        # Each box represents a part of the image where a particular object was detected.
+        boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
+        # Each score represent how level of confidence for each of the objects.
+        # Score is shown on the result image, together with the class label.
+        scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
+        classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
+        num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')
+        # Actual detection.
+        #start_time = time.time()
+        (boxes, scores, classes, num_detections) = self.sess.run(
+            [boxes, scores, classes, num_detections],
+            feed_dict={image_tensor: image_np_expanded})
+        #elapsed_time = time.time() - start_time
+        #print('inference time cost: {}'.format(elapsed_time))
+
+        return (boxes, scores, classes, num_detections)
+
+    def detect_raw(self, imgcv, **kwargs):
+        return self.run(imgcv)
+
+    def detect(self, imgcv, **kwargs):
+        threshold = kwargs.get('threshold', 0.7)
+        faces = self.detect_raw(imgcv, **kwargs)
+        im_height,im_width,_ = imgcv.shape
+        return self._format_result(faces,threshold,im_width,im_height)
+
+    # Format the results
+    def _format_result(self, result,threshold,im_width,im_height):
+        out_list = []
+        boxes,scores,classes,num_detections = result
+        indexes = np.squeeze(np.argwhere(scores[0]>threshold),axis=1)
+        #print("indexes",indexes)
+        for index_face in indexes:
+            box = boxes[0][index_face]
+            ymin, xmin, ymax, xmax = box[0],box[1],box[2],box[3]
+            (left, right, top, bottom) = (xmin * im_width, xmax * im_width,
+                                  ymin * im_height, ymax * im_height)
+            prob = scores[0][index_face]
+            if prob> threshold:
+                formatted_res = dict()
+                formatted_res["class"] = 'face'
+                formatted_res["prob"] = prob
+                formatted_res["box"] = {
+                    "topleft": {'x': int(left), 'y': int(top)},
+                    "bottomright": {'x': int(right), 'y': int(bottom)}
+                    }
+                out_list.append(formatted_res)
+        #Return the result
+        return out_list
 
 if __name__ == '__main__':
     import sys

face/face.py

@@ -1,6 +1,9 @@
 #!/usr/bin/env python
 
-import dlib
+try:
+    import dlib
+except:
+    pass
 import os
 
 WORK_DIR = "/LFS/facelib"
@@ -36,11 +39,12 @@ class Face(object):
     Methods:
         detect()
             imgcv = grayscale or colored numpy image
-            kwargs = optional argumentts for some algorithms e.g. upsamples=1
+            kwargs = optional argumentts for some algorithms e.g. upsamples=1 for dlib,
+            threshold =0.7 for mobilenet
     """
     def __init__(self, detector_method='opencv', detector_model=None,
                  predictor_model='small', recognition_method='dlib',
-                 recognition_model=None):
+                 recognition_model=None,trt_enable=False,precision ='FP32',gpu_frac=0.3):
         if detector_method == 'dlib':
             from .detect_face import FaceDetectorDlib
             self._detector = FaceDetectorDlib()
@@ -50,6 +54,9 @@ def __init__(self, detector_method='opencv', detector_model=None,
         elif detector_method == 'yolo':
             from .detect_face import FaceDetectorYolo
             self._detector = FaceDetectorYolo()
+        elif detector_method == 'mobilenet':
+            from .detect_face import FaceDetectorMobilenet
+            self._detector = FaceDetectorMobilenet(model_name=detector_model,trt_enable=trt_enable,precision =precision,gpu_frac=gpu_frac)
         else:
             from .detect_face import FaceDetectorOpenCV
             self._detector = FaceDetectorOpenCV(detector_model)

setup.py

@@ -18,6 +18,6 @@
         'face': ['models/*'],
     },
 
-    install_requires=['numpy', 'dlib'],
+    install_requires=['numpy'],
     zip_safe=False
     )