diff options
Diffstat (limited to 'megapixels/app/processors/face_detector.py')
| -rw-r--r-- | megapixels/app/processors/face_detector.py | 23 |
1 files changed, 18 insertions, 5 deletions
diff --git a/megapixels/app/processors/face_detector.py b/megapixels/app/processors/face_detector.py index 3a90c557..a805a474 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) @@ -115,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 @@ -132,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]] @@ -155,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 @@ -171,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) |