JeVoisBase  1.21
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
Loading...
Searching...
No Matches
lpd_yunet.py
Go to the documentation of this file.
1from itertools import product
2
3import numpy as np
4import cv2 as cv
5
6class 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
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
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_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_postprocess(outputBlob)
60
61 return results
62
63 def _postprocess(self, blob):
64 # Decode
65 dets = self._decode_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
__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
setBackend(self, backendId)
Definition lpd_yunet.py:31
setTarget(self, targetId)
Definition lpd_yunet.py:35
_postprocess(self, blob)
Definition lpd_yunet.py:63
infer(self, image)
Definition lpd_yunet.py:47
setInputSize(self, inputSize)
Definition lpd_yunet.py:39
_decode(self, blob)
Definition lpd_yunet.py:111
_preprocess(self, image)
Definition lpd_yunet.py:44