summaryrefslogtreecommitdiff
path: root/megapixels/app/processors/face_detector.py
blob: 747e057b01a38d59b179e2cf9a53f20562dda6a6 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
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 DetectorHaar:

  im_size = (400, 400)
  cascade_name = types.HaarCascade.FRONTAL

  def __init__(self, cascade=types.HaarCascade.FRONTAL):
    self.log = logger_utils.Logger.getLogger()

  def detect(self, im, scale_factor=1.05, overlaps=5):
    return
    

class DetectorDLIBCNN:

  dnn_size = (300, 300)
  pyramids = 0
  conf_thresh = 0.85

  def __init__(self, opt_gpu=0):
    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