summaryrefslogtreecommitdiff
path: root/megapixels/app/processors/person_detector.py
blob: 6daa8c40778c77ae675c5c0f4c0c91e9a0b2bbef (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
import sys
import os
from os.path import join
from pathlib import Path

import cv2 as cv
import numpy as np
import imutils
import operator

from app.utils import im_utils, logger_utils
from app.models.bbox import BBox
from app.settings import app_cfg as cfg
from app.settings import types


class DetectorCVDNN:

  # MobileNet SSD
  dnn_scale = 0.007843  # fixed
  dnn_mean = (127.5, 127.5, 127.5)  # fixed
  dnn_crop = False  # crop or force resize
  blob_size = (300, 300)
  conf = 0.95

  # detect
  CLASSES = ["background", "aeroplane", "bicycle", "bird", "boat",
    "bottle", "bus", "car", "cat", "chair", "cow", "diningtable",
    "dog", "horse", "motorbike", "person", "pottedplant", "sheep",
    "sofa", "train", "tvmonitor"]
  
  def __init__(self):
    self.log = logger_utils.Logger.getLogger()
    fp_prototxt = join(cfg.DIR_MODELS_CAFFE, 'mobilenet_ssd', 'MobileNetSSD_deploy.prototxt')
    fp_model = join(cfg.DIR_MODELS_CAFFE, 'mobilenet_ssd', 'MobileNetSSD_deploy.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, conf=None, largest=False, pyramids=None, zone=False, blob_size=None):
    """Detects bodies and returns (list) of (BBox)"""
    conf = self.conf if conf is None else conf
    blob_size = self.blob_size if blob_size is None else blob_size
    im = cv.resize(im, blob_size)
    dim = im.shape[:2][::-1]
    blob = cv.dnn.blobFromImage(im, self.dnn_scale, dim, self.dnn_mean)
    self.net.setInput(blob)
    net_outputs = self.net.forward()

    bboxes = []
    for i in range(0, net_outputs.shape[2]):
      det_conf = float(net_outputs[0, 0, i, 2])
      bounds = np.array(net_outputs[0, 0, i, 3:7])  # bug: ensure all x,y within 1.0 ?
      if det_conf > conf and np.all(bounds < 1):
        idx = int(net_outputs[0, 0, i, 1])
        if self.CLASSES[idx] == "person":
          rect_norm = net_outputs[0, 0, i, 3:7]
          bboxes.append(BBox(*rect_norm))

    if largest and len(bboxes) > 1:
      # only keep largest
      bboxes.sort(key=operator.attrgetter('area'), reverse=True)
      bboxes = [bboxes[0]]

    return bboxes