diff options
Diffstat (limited to 'megapixels/app/processors/person_detector.py')
| -rw-r--r-- | megapixels/app/processors/person_detector.py | 65 |
1 files changed, 65 insertions, 0 deletions
diff --git a/megapixels/app/processors/person_detector.py b/megapixels/app/processors/person_detector.py new file mode 100644 index 00000000..6daa8c40 --- /dev/null +++ b/megapixels/app/processors/person_detector.py @@ -0,0 +1,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
\ No newline at end of file |
