JeVoisBase  1.18
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
lpd_yunet.py
Go to the documentation of this file.
1 from itertools import product
2 
3 import numpy as np
4 import cv2 as cv
5 
6 class LPD_YuNet:
7  def __init__(self, modelPath, inputSize=[320, 240], confThreshold=0.8, nmsThreshold=0.3, topK=5000, keepTopK=750, backendId=0, targetId=0):
8  self.model_path = modelPath
9  self.input_size = np.array(inputSize)
10  self.confidence_threshold=confThreshold
11  self.nms_threshold = nmsThreshold
12  self.top_k = topK
13  self.keep_top_k = keepTopK
14  self.backend_id = backendId
15  self.target_id = targetId
16 
17  self.output_names = ['loc', 'conf', 'iou']
18  self.min_sizes = [[10, 16, 24], [32, 48], [64, 96], [128, 192, 256]]
19  self.steps = [8, 16, 32, 64]
20  self.variance = [0.1, 0.2]
21 
22  # load model
23  self.model = cv.dnn.readNet(self.model_path)
24  # generate anchors/priorboxes
25  self._priorGen()
26 
27  @property
28  def name(self):
29  return self.__class__.__name__
30 
31  def setBackend(self, backendId):
32  self.backend_id = backendId
33  self.model.setPreferableBackend(self.backend_id)
34 
35  def setTarget(self, targetId):
36  self.target_id = targetId
37  self.model.setPreferableTarget(self.target_id)
38 
39  def setInputSize(self, inputSize):
40  self.input_size = inputSize
41  # re-generate anchors/priorboxes
42  self._priorGen()
43 
44  def _preprocess(self, image):
45  return cv.dnn.blobFromImage(image)
46 
47  def infer(self, image):
48  assert image.shape[0] == self.input_size[1], '{} (height of input image) != {} (preset height)'.format(image.shape[0], self.input_size[1])
49  assert image.shape[1] == self.input_size[0], '{} (width of input image) != {} (preset width)'.format(image.shape[1], self.input_size[0])
50 
51  # Preprocess
52  inputBlob = self._preprocess(image)
53 
54  # Forward
55  self.model.setInput(inputBlob)
56  outputBlob = self.model.forward(self.output_names)
57 
58  # Postprocess
59  results = self._postprocess(outputBlob)
60 
61  return results
62 
63  def _postprocess(self, blob):
64  # Decode
65  dets = self._decode(blob)
66 
67  # NMS
68  keepIdx = cv.dnn.NMSBoxes(
69  bboxes=dets[:, 0:4].tolist(),
70  scores=dets[:, -1].tolist(),
71  score_threshold=self.confidence_threshold,
72  nms_threshold=self.nms_threshold,
73  top_k=self.top_k
74  ) # box_num x class_num
75  if len(keepIdx) > 0:
76  dets = dets[keepIdx]
77  return dets[:self.keep_top_k]
78  else:
79  return np.empty(shape=(0, 9))
80 
81  def _priorGen(self):
82  w, h = self.input_size
83  feature_map_2th = [int(int((h + 1) / 2) / 2),
84  int(int((w + 1) / 2) / 2)]
85  feature_map_3th = [int(feature_map_2th[0] / 2),
86  int(feature_map_2th[1] / 2)]
87  feature_map_4th = [int(feature_map_3th[0] / 2),
88  int(feature_map_3th[1] / 2)]
89  feature_map_5th = [int(feature_map_4th[0] / 2),
90  int(feature_map_4th[1] / 2)]
91  feature_map_6th = [int(feature_map_5th[0] / 2),
92  int(feature_map_5th[1] / 2)]
93 
94  feature_maps = [feature_map_3th, feature_map_4th,
95  feature_map_5th, feature_map_6th]
96 
97  priors = []
98  for k, f in enumerate(feature_maps):
99  min_sizes = self.min_sizes[k]
100  for i, j in product(range(f[0]), range(f[1])): # i->h, j->w
101  for min_size in min_sizes:
102  s_kx = min_size / w
103  s_ky = min_size / h
104 
105  cx = (j + 0.5) * self.steps[k] / w
106  cy = (i + 0.5) * self.steps[k] / h
107 
108  priors.append([cx, cy, s_kx, s_ky])
109  self.priors = np.array(priors, dtype=np.float32)
110 
111  def _decode(self, blob):
112  loc, conf, iou = blob
113  # get score
114  cls_scores = conf[:, 1]
115  iou_scores = iou[:, 0]
116  # clamp
117  _idx = np.where(iou_scores < 0.)
118  iou_scores[_idx] = 0.
119  _idx = np.where(iou_scores > 1.)
120  iou_scores[_idx] = 1.
121  scores = np.sqrt(cls_scores * iou_scores)
122  scores = scores[:, np.newaxis]
123 
124  scale = self.input_size
125 
126  # get four corner points for bounding box
127  bboxes = np.hstack((
128  (self.priors[:, 0:2] + loc[:, 4: 6] * self.variance[0] * self.priors[:, 2:4]) * scale,
129  (self.priors[:, 0:2] + loc[:, 6: 8] * self.variance[0] * self.priors[:, 2:4]) * scale,
130  (self.priors[:, 0:2] + loc[:, 10:12] * self.variance[0] * self.priors[:, 2:4]) * scale,
131  (self.priors[:, 0:2] + loc[:, 12:14] * self.variance[0] * self.priors[:, 2:4]) * scale
132  ))
133 
134  dets = np.hstack((bboxes, scores))
135  return dets
lpd_yunet.LPD_YuNet
Definition: lpd_yunet.py:6
lpd_yunet.LPD_YuNet.name
def name(self)
Definition: lpd_yunet.py:28
lpd_yunet.LPD_YuNet._postprocess
def _postprocess(self, blob)
Definition: lpd_yunet.py:63
lpd_yunet.LPD_YuNet._priorGen
def _priorGen(self)
Definition: lpd_yunet.py:81
demo.int
int
Definition: demo.py:37
lpd_yunet.LPD_YuNet.input_size
input_size
Definition: lpd_yunet.py:9
lpd_yunet.LPD_YuNet.model
model
Definition: lpd_yunet.py:23
lpd_yunet.LPD_YuNet.setTarget
def setTarget(self, targetId)
Definition: lpd_yunet.py:35
lpd_yunet.LPD_YuNet.model_path
model_path
Definition: lpd_yunet.py:8
lpd_yunet.LPD_YuNet.output_names
output_names
Definition: lpd_yunet.py:17
lpd_yunet.LPD_YuNet.priors
priors
Definition: lpd_yunet.py:109
lpd_yunet.LPD_YuNet.__init__
def __init__(self, modelPath, inputSize=[320, 240], confThreshold=0.8, nmsThreshold=0.3, topK=5000, keepTopK=750, backendId=0, targetId=0)
Definition: lpd_yunet.py:7
lpd_yunet.LPD_YuNet._preprocess
def _preprocess(self, image)
Definition: lpd_yunet.py:44
lpd_yunet.LPD_YuNet.infer
def infer(self, image)
Definition: lpd_yunet.py:47
lpd_yunet.LPD_YuNet.target_id
target_id
Definition: lpd_yunet.py:15
lpd_yunet.LPD_YuNet.keep_top_k
keep_top_k
Definition: lpd_yunet.py:13
lpd_yunet.LPD_YuNet.confidence_threshold
confidence_threshold
Definition: lpd_yunet.py:10
lpd_yunet.LPD_YuNet.top_k
top_k
Definition: lpd_yunet.py:12
lpd_yunet.LPD_YuNet._decode
def _decode(self, blob)
Definition: lpd_yunet.py:111
lpd_yunet.LPD_YuNet.backend_id
backend_id
Definition: lpd_yunet.py:14
lpd_yunet.LPD_YuNet.variance
variance
Definition: lpd_yunet.py:20
lpd_yunet.LPD_YuNet.steps
steps
Definition: lpd_yunet.py:19
lpd_yunet.LPD_YuNet.setBackend
def setBackend(self, backendId)
Definition: lpd_yunet.py:31
lpd_yunet.LPD_YuNet.nms_threshold
nms_threshold
Definition: lpd_yunet.py:11
lpd_yunet.LPD_YuNet.setInputSize
def setInputSize(self, inputSize)
Definition: lpd_yunet.py:39
lpd_yunet.LPD_YuNet.min_sizes
min_sizes
Definition: lpd_yunet.py:18