fix face detectors

1 files changed, 12 insertions, 3 deletions

diff --git a/megapixels/app/processors/face_detector.py b/megapixels/app/processors/face_detector.py
index 75ba54d4..a805a474 100644
--- a/megapixels/app/processors/face_detector.py
+++ b/megapixels/app/processors/face_detector.py
@@ -119,7 +119,7 @@ class DetectorDLIBHOG:
     self.log = logger_utils.Logger.getLogger()
     self.detector = dlib.get_frontal_face_detector()
 
-  def detect(self, im, size=None, conf_thresh=None, pyramids=0, largest=False):
+  def detect(self, im, size=None, conf_thresh=None, pyramids=0, largest=False, zone=False):
     conf_thresh = self.conf_thresh if conf_thresh is None else conf_thresh
     dnn_size = self.size if size is None else size
     pyramids = self.pyramids if pyramids is None else pyramids
@@ -136,8 +136,13 @@ class DetectorDLIBHOG:
           bbox = BBox.from_dlib_dim(rect, dim)
           bboxes.append(bbox)
     
+    # filter to keep on faces inside zone
+    if zone:
+      bboxes = [b for b in bboxes if b.cx > zone[0] and b.cx < 1.0 - zone[0] \
+        and b.cy > zone[1] and b.cy < 1.0 - zone[1]]
+
+    # filter to keep only largest face
     if largest and len(bboxes) > 1:
-      # only keep largest
       bboxes.sort(key=operator.attrgetter('area'), reverse=True)
       bboxes = [bboxes[0]]
 
@@ -159,7 +164,7 @@ class DetectorCVDNN:
     self.net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
     self.net.setPreferableTarget(cv.dnn.DNN_TARGET_CPU)
 
-  def detect(self, im, size=None, conf_thresh=None, largest=False, pyramids=None):
+  def detect(self, im, size=None, conf_thresh=None, largest=False, pyramids=None, zone=False):
     """Detects faces and returns (list) of (BBox)"""
     conf_thresh = self.conf_thresh if conf_thresh is None else conf_thresh
     dnn_size = self.size if size is None else size
@@ -175,6 +180,10 @@ class DetectorCVDNN:
         rect_norm = net_outputs[0, 0, i, 3:7]
         bboxes.append(BBox(*rect_norm))
     
+    if zone:
+      bboxes = [b for b in bboxes if b.cx > zone[0] and b.cx < 1.0 - zone[0] \
+        and b.cy > zone[1] and b.cy < 1.0 - zone[1]]
+
     if largest and len(bboxes) > 1:
       # only keep largest
       bboxes.sort(key=operator.attrgetter('area'), reverse=True)