diff options
Diffstat (limited to 'megapixels/app/processors/face_detector.py')
| -rw-r--r-- | megapixels/app/processors/face_detector.py | 48 |
1 files changed, 23 insertions, 25 deletions
diff --git a/megapixels/app/processors/face_detector.py b/megapixels/app/processors/face_detector.py index 593e9feb..3a90c557 100644 --- a/megapixels/app/processors/face_detector.py +++ b/megapixels/app/processors/face_detector.py @@ -24,15 +24,15 @@ class DetectorMTCNN: from mtcnn.mtcnn import MTCNN self.detector = MTCNN() - def detect(self, im, opt_size=(400,400), opt_conf_thresh=None, opt_pyramids=None, opt_largest=False): + def detect(self, im, size=(400,400), conf_thresh=None, pyramids=None, largest=False): '''Detects face using MTCNN and returns (list) of BBox :param im: (numpy.ndarray) image :returns list of BBox ''' bboxes = [] - #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 + #conf_thresh = self.conf_thresh if conf_thresh is None else conf_thresh + #pyramids = self.pyramids if pyramids is None else pyramids + 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] @@ -43,7 +43,7 @@ class DetectorMTCNN: bbox = BBox.from_xywh_dim(*rect, dim) bboxes.append(bbox) - if opt_largest and len(bboxes) > 1: + if largest and len(bboxes) > 1: # only keep largest bboxes.sort(key=operator.attrgetter('area'), reverse=True) bboxes = [bboxes[0]] @@ -70,34 +70,33 @@ class DetectorDLIBCNN: pyramids = 0 conf_thresh = 0.85 - def __init__(self, opt_gpu=0): + def __init__(self, gpu=0): import dlib self.log = logger_utils.Logger.getLogger() cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', '') - os.environ['CUDA_VISIBLE_DEVICES'] = str(opt_gpu) + os.environ['CUDA_VISIBLE_DEVICES'] = str(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, opt_largest=False): + def detect(self, im, size=None, conf_thresh=None, pyramids=None, largest=False): bboxes = [] - 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 + conf_thresh = self.conf_thresh if conf_thresh is None else conf_thresh + pyramids = self.pyramids if pyramids is None else pyramids + dnn_size = self.dnn_size if size is None else 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, opt_pyramids) + mmod_rects = self.detector(im, pyramids) # sort results for mmod_rect in mmod_rects: - self.log.debug('conf: {}, this: {}'.format(conf_thresh, mmod_rect.confidence)) if mmod_rect.confidence > conf_thresh: bbox = BBox.from_dlib_dim(mmod_rect.rect, dim) bboxes.append(bbox) - if opt_largest and len(bboxes) > 1: + if largest and len(bboxes) > 1: # only keep largest bboxes.sort(key=operator.attrgetter('area'), reverse=True) bboxes = [bboxes[0]] @@ -116,25 +115,24 @@ class DetectorDLIBHOG: self.log = logger_utils.Logger.getLogger() self.detector = dlib.get_frontal_face_detector() - def detect(self, im, opt_size=None, opt_conf_thresh=None, opt_pyramids=0, opt_largest=False): - 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 + def detect(self, im, size=None, conf_thresh=None, pyramids=0, largest=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 - im = im_utils.resize(im, width=opt_size[0], height=opt_size[1]) + im = im_utils.resize(im, width=dnn_size[0], height=dnn_size[1]) dim = im.shape[:2][::-1] im = im_utils.bgr2rgb(im) # ? hog_results = self.detector.run(im, pyramids) bboxes = [] if len(hog_results[0]) > 0: - self.log.debug(hog_results) for rect, score, direction in zip(*hog_results): if score > conf_thresh: bbox = BBox.from_dlib_dim(rect, dim) bboxes.append(bbox) - if opt_largest and len(bboxes) > 1: + if largest and len(bboxes) > 1: # only keep largest bboxes.sort(key=operator.attrgetter('area'), reverse=True) bboxes = [bboxes[0]] @@ -157,10 +155,10 @@ class DetectorCVDNN: 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, opt_largest=False, opt_pyramids=None): + def detect(self, im, size=None, conf_thresh=None, largest=False, pyramids=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 + conf_thresh = self.conf_thresh if conf_thresh is None else conf_thresh + dnn_size = self.size if size is None else size im = cv.resize(im, dnn_size) blob = cv.dnn.blobFromImage(im, self.dnn_scale, dnn_size, self.dnn_mean) self.net.setInput(blob) @@ -173,7 +171,7 @@ class DetectorCVDNN: rect_norm = net_outputs[0, 0, i, 3:7] bboxes.append(BBox(*rect_norm)) - if opt_largest and len(bboxes) > 1: + if largest and len(bboxes) > 1: # only keep largest bboxes.sort(key=operator.attrgetter('area'), reverse=True) bboxes = [bboxes[0]] |