.

1 files changed, 103 insertions, 0 deletions

diff --git a/megapixels/app/processors/face_detector.py b/megapixels/app/processors/face_detector.py
new file mode 100644
index 00000000..02d068dc
--- /dev/null
+++ b/megapixels/app/processors/face_detector.py
@@ -0,0 +1,103 @@
+import os
+from os.path import join
+from pathlib import Path
+
+import cv2 as cv
+import numpy as np
+import dlib
+# import imutils
+
+from app.utils import im_utils, logger_utils
+from app.models.bbox import BBox
+from app.settings import app_cfg as cfg
+
+class DetectorDLIBCNN:
+
+  dnn_size = (300, 300)
+  pyramids = 0
+  conf_thresh = 0.85
+
+  def __init__(self, opt_gpu):
+    self.log = logger_utils.Logger.getLogger()
+    cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', '')
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(opt_gpu)
+    self.log.info('load model: {}'.format(cfg.DIR_MODELS_DLIB_CNN))
+    self.detector = dlib.cnn_face_detection_model_v1(cfg.DIR_MODELS_DLIB_CNN)
+    os.environ['CUDA_VISIBLE_DEVICES'] = cuda_visible_devices  # reset
+
+  def detect(self, im, opt_size=None, opt_conf_thresh=None, opt_pyramids=None):
+    rois = []
+    conf_thresh = self.conf_thresh if opt_conf_thresh is None else opt_conf_thresh
+    pyramids = self.pyramids if opt_pyramids is None else opt_pyramids
+    dnn_size = self.dnn_size if opt_size is None else opt_size
+    # resize image
+    im = im_utils.resize(im, width=dnn_size[0], height=dnn_size[1])
+    dim = im.shape[:2][::-1]
+    im = im_utils.bgr2rgb(im)  # convert to RGB for dlib
+    # run detector
+    mmod_rects = self.detector(im, 1)
+    # sort results
+    for mmod_rect in mmod_rects:
+      if mmod_rect.confidence > conf_thresh:
+        bbox = BBox.from_dlib_dim(mmod_rect.rect, dim)
+        rois.append(bbox)
+    return rois
+
+
+class DetectorDLIBHOG:
+
+  size = (320, 240)
+  pyramids = 0
+
+  def __init__(self):
+    self.detector = dlib.get_frontal_face_detector()
+
+  def detect(self, im, opt_size=None, opt_conf_thresh=None, opt_pyramids=0):
+    conf_thresh = self.conf_thresh if opt_conf_thresh is None else opt_conf_thresh
+    dnn_size = self.size if opt_size is None else opt_size
+    pyramids = self.pyramids if opt_pyramids is None else opt_pyramids
+    
+    im = im_utils.resize(im, width=opt_size[0], height=opt_size[1])
+    dim = im.shape[:2][::-1]
+    im = im_utils.bgr2rgb(im)  # ?
+    hog_results = self.detector.run(im, pyramids)
+    
+    rois = []
+    if len(hog_results[0]) > 0:
+      for rect, score, direction in zip(*hog_results):
+        if score > opt_conf_thresh:
+          bbox = BBox.from_dlib_dim(rect, dim)
+          rois.append(bbox)
+    return rois
+
+class DetectorCVDNN:
+
+  dnn_scale = 1.0  # fixed
+  dnn_mean = (104.0, 177.0, 123.0)  # fixed
+  dnn_crop = False  # crop or force resize
+  size = (300, 300)
+  conf_thresh = 0.85
+  
+  def __init__(self):
+    fp_prototxt = join(cfg.DIR_MODELS_CAFFE, 'face_detect', 'opencv_face_detector.prototxt')
+    fp_model = join(cfg.DIR_MODELS_CAFFE, 'face_detect', 'opencv_face_detector.caffemodel')
+    self.net = cv.dnn.readNet(fp_prototxt, fp_model)
+    self.net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
+    self.net.setPreferableTarget(cv.dnn.DNN_TARGET_CPU)
+
+  def detect(self, im, opt_size=None, opt_conf_thresh=None):
+    """Detects faces and returns (list) of (BBox)"""
+    conf_thresh = self.conf_thresh if opt_conf_thresh is None else opt_conf_thresh
+    dnn_size = self.size if opt_size is None else opt_size
+    im = cv.resize(im, dnn_size)
+    blob = cv.dnn.blobFromImage(im, self.dnn_scale, dnn_size, self.dnn_mean)
+    self.net.setInput(blob)
+    net_outputs = self.net.forward()
+
+    rois = []
+    for i in range(0, net_outputs.shape[2]):
+      conf = net_outputs[0, 0, i, 2]
+      if conf > opt_conf_thresh:
+        rect_norm = net_outputs[0, 0, i, 3:7]
+        rois.append(BBox(*rect_norm))
+    return rois
\ No newline at end of file