1 files changed, 50 insertions, 1 deletions
diff --git a/megapixels/app/processors/face_detector.py b/megapixels/app/processors/face_detector.py
index 0e194f7d..fbf91071 100644
--- a/megapixels/app/processors/face_detector.py
+++ b/megapixels/app/processors/face_detector.py
@@ -14,8 +14,57 @@ 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/TropComplique/mtcnn-pytorch/
+  # pip install mtcnn
+
+  dnn_size = (300, 300)
+
+  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
+
+    im = im_utils.resize(im, width=dnn_size[0], height=dnn_size[1])
+    dim = im.shape[:2][::-1]
+    dets = self.detector.detect_faces(im)
+    for det in dets:
+      rect = det['box']
+      #keypoints = det['keypoints']  # not using here. see 'face_landmarks.py'
+      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 DetectorMTCNN_TF:
   
+  # using TF for inference can cause GPU issues with other frameworks
   # https://github.com/ipazc/mtcnn
   # pip install mtcnn