Merge branch 'master' of github.com:adamhrv/megapixels_dev

1 files changed, 73 insertions, 4 deletions

diff --git a/megapixels/app/processors/face_detector.py b/megapixels/app/processors/face_detector.py
index 0e194f7d..7b5310c5 100644
--- a/megapixels/app/processors/face_detector.py
+++ b/megapixels/app/processors/face_detector.py
@@ -14,9 +14,17 @@ from app.settings import app_cfg as cfg
 from app.settings import types
 
 
-class DetectorMTCNN:
+class DetectorMTCNN_CVDNN:
+
+  # https://github.com/CongWeilin/mtcnn-caffe
+
+  def __init__(self):
+    pass
+    
+
+class DetectorMTCNN_PT:
   
-  # https://github.com/ipazc/mtcnn
+  # https://github.com/TropComplique/mtcnn-pytorch/
   # pip install mtcnn
 
   dnn_size = (300, 300)
@@ -54,6 +62,64 @@ class DetectorMTCNN:
 
     return bboxes
 
+class DetectorMTCNN_TF:
+  
+  # using TF for inference can cause GPU issues with other frameworks
+  # https://github.com/ipazc/mtcnn
+  # pip install mtcnn
+
+  dnn_size = (300, 300)
+  conf_thresh = 0.9
+
+  def __init__(self, size=(400,400), gpu=0):
+    self.log = logger_utils.Logger.getLogger()
+    device_cur = os.getenv('CUDA_VISIBLE_DEVICES', '')
+    self.log.info(f'Change CUDA_VISIBLE_DEVICES from "{device_cur}" to "{gpu}"')
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
+    from mtcnn.mtcnn import MTCNN
+    self.detector = MTCNN()
+    os.environ['CUDA_VISIBLE_DEVICES'] = device_cur  # reset
+
+  def detect(self, im, size=(400,400), conf_thresh=None, pyramids=None, largest=False, zone=None):
+    '''Detects face using MTCNN and returns (list) of BBox
+    :param im: (numpy.ndarray) image
+    :returns list of BBox
+    '''
+
+    bboxes = []
+    dnn_size = self.dnn_size if size is None else size
+    conf_thresh = self.conf_thresh if conf_thresh is None else conf_thresh
+
+    im = im_utils.resize(im, width=dnn_size[0], height=dnn_size[1])
+    dim = im.shape[:2][::-1]
+    dets = self.detector.detect_faces(im)
+    '''
+    {
+      'box': [4, 140, 14, 18], 
+      'confidence': 0.9588413834571838, 
+      'keypoints': {
+          'left_eye': (8, 147), 
+          'right_eye': (14, 146),
+          'nose': (12, 151),
+          'mouth_left': (9, 155),
+          'mouth_right': (14, 154)
+        }
+      }
+    '''
+    for det in dets:
+      rect = det['box']
+      conf = det['confidence']
+      if conf > conf_thresh:
+        bbox = BBox.from_xywh_dim(*rect, dim)
+        bboxes.append(bbox)
+
+    if largest and len(bboxes) > 1:
+      # only keep largest
+      bboxes.sort(key=operator.attrgetter('area'), reverse=True)
+      bboxes = [bboxes[0]]
+
+    return bboxes
+
 
 class DetectorHaar:
 
@@ -173,8 +239,11 @@ class DetectorCVDNN:
 
     bboxes = []
     for i in range(0, net_outputs.shape[2]):
-      conf = net_outputs[0, 0, i, 2]
-      if conf > conf_thresh:
+      conf = float(net_outputs[0, 0, i, 2])
+      # BUG: this face detector creates ghost face detections in stage-left from nose-bottom neck
+      # temp fix is to elminate ROI extending outside of frame
+      bounds = np.array(net_outputs[0, 0, i, 3:7])
+      if conf > conf_thresh and np.all(bounds < 1):
         rect_norm = net_outputs[0, 0, i, 3:7]
         bboxes.append(BBox(*rect_norm))