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