fix roi index, clean up pose, roi, records, vector

1 files changed, 6 insertions, 2 deletions

diff --git a/megapixels/app/processors/face_detector.py b/megapixels/app/processors/face_detector.py
index 3a90c557..75ba54d4 100644
--- a/megapixels/app/processors/face_detector.py
+++ b/megapixels/app/processors/face_detector.py
@@ -24,7 +24,7 @@ class DetectorMTCNN:
     from mtcnn.mtcnn import MTCNN
     self.detector = MTCNN()
 
-  def detect(self, im, size=(400,400), conf_thresh=None, pyramids=None, largest=False):
+  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
@@ -79,7 +79,7 @@ class DetectorDLIBCNN:
     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, size=None, conf_thresh=None, pyramids=None, largest=False):
+  def detect(self, im, size=None, conf_thresh=None, pyramids=None, largest=False, zone=None):
     bboxes = []
     conf_thresh = self.conf_thresh if conf_thresh is None else conf_thresh
     pyramids = self.pyramids if pyramids is None else pyramids
@@ -96,6 +96,10 @@ class DetectorDLIBCNN:
         bbox = BBox.from_dlib_dim(mmod_rect.rect, dim)
         bboxes.append(bbox)
 
+    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)