JeVoisBase  1.21
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
Loading...
Searching...
No Matches
PyCoralDetect.py
Go to the documentation of this file.
1import pyjevois
2if pyjevois.pro: import libjevoispro as jevois
3else: import libjevois as jevois
4import cv2 as cv
5import numpy as np
6from PIL import Image
7from PIL import ImageDraw
8from pycoral.adapters import common
9from pycoral.adapters import detect
10from pycoral.utils.dataset import read_label_file
11from pycoral.utils.edgetpu import make_interpreter
12import time
13
14## Object detection using Coral Edge TPU
15#
16# More pre-trained models are available at https://coral.ai/models/
17#
18#
19# @author Laurent Itti
20#
21# @videomapping YUYV 320 264 30.0 YUYV 320 240 30.0 JeVois PyClassificationDNN
22# @email itti@usc.edu
23# @address 880 W 1st St Suite 807, Los Angeles CA 90012, USA
24# @copyright Copyright (C) 2020 by Laurent Itti
25# @mainurl http://jevois.org
26# @supporturl http://jevois.org
27# @otherurl http://jevois.org
28# @license GPL v3
29# @distribution Unrestricted
30# @restrictions None
31# @ingroup modules
33 # ####################################################################################################
34 ## Constructor
35 def __init__(self):
37 jevois.LFATAL("A Google Coral EdgeTPU is required for this module (PCIe M.2 2230 A+E or USB)")
38
39 self.threshold = 0.4 # Confidence threshold (0..1), higher for stricter confidence.
40 self.rgb = True # True if model expects RGB inputs, otherwise it expects BGR
41
42 # Select one of the models:
43 self.model = 'MobileDetSSD' # expects 320x320
44
45 # You should not have to edit anything beyond this point.
46 if (self.model == 'MobileDetSSD'):
47 classnames = 'coco_labels.txt'
48 modelname = 'ssdlite_mobiledet_coco_qat_postprocess_edgetpu.tflite'
49
50 # Load names of classes:
51 sdir = pyjevois.share + '/coral/detection/'
52 self.labels = read_label_file(sdir + classnames)
53
54 # Load network:
55 self.interpreter = make_interpreter(sdir + modelname)
56 #self.interpreter = make_interpreter(*modelname.split('@'))
57 self.interpreter.allocate_tensors()
58 self.timer = jevois.Timer('Coral classification', 10, jevois.LOG_DEBUG)
59
60 # ####################################################################################################
61 def stringToRGBA(self, str):
62 col = 0x80808080
63 alpha = 0xff
64 for c in str: col = ord(c) + ((col << 5) - col)
65 col = (col & 0xffffff) | (alpha << 24)
66 return col
67
68 # ####################################################################################################
69 def draw_objects(self, draw, objs, labels):
70 """Draws the bounding box and label for each object."""
71 for obj in objs:
72 bbox = obj.bbox
73 draw.rectangle([(bbox.xmin, bbox.ymin), (bbox.xmax, bbox.ymax)], outline='red')
74 draw.text((bbox.xmin+10, bbox.ymin+10), '%s: %.2f' % (labels.get(obj.id, obj.id), obj.score), fill='red')
75
76 # ####################################################################################################
77 ## JeVois main processing function
78 def process(self, inframe, outframe):
79 frame = inframe.getCvRGB() if self.rgb else inframe.getCvBGR()
80 self.timer.start()
81
82 h = frame.shape[0]
83 w = frame.shape[1]
84
85 # Set the input:
86 image = Image.fromarray(frame);
87 _, scale = common.set_resized_input(self.interpreter, image.size,
88 lambda size: image.resize(size, Image.LANCZOS))
89
90 # Run the model
91 start = time.perf_counter()
92 self.interpreter.invoke()
93 inference_time = time.perf_counter() - start
94
95 # Get detections with high enough scores:
96 # JEVOIS: this used to work but now fails with pycoral-2.0.0 / libedgetpu-1.6.0 / tflite 2.5.0 with error:
97 # AttributeError: 'Interpreter' object has no attribute '_get_full_signature_list'
98 #objs = detect.get_objects(self.interpreter, self.threshold, scale)
99
100 # Use modified function below instead:
101 objs = self.get_objects(self.interpreter, self.threshold, scale)
102
103 # Draw the detections:
104 image = image.convert('RGB')
105 self.draw_objects(ImageDraw.Draw(image), objs, self.labels)
106
107 # Back to OpenCV:
108 frame = np.array(image)
109
110 # Output to serial:
111 #for obj in objs:
112 # print(self.labels.get(obj.id, obj.id))
113 # print(' id: ', obj.id)
114 # print(' score: ', obj.score)
115 # print(' bbox: ', obj.bbox)
116
117 # Put efficiency information:
118 cv.putText(frame, 'JeVois Coral Detection - ' + self.model, (3, 15),
119 cv.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv.LINE_AA)
120
121 fps = self.timer.stop()
122 label = fps + ', %dms' % (inference_time * 1000.0)
123 cv.putText(frame, label, (3, h-5), cv.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv.LINE_AA)
124
125 # Send output frame to host:
126 if self.rgb: outframe.sendCvRGB(frame)
127 else: outframe.sendCv(frame)
128
129 # ####################################################################################################
130 ## Process function with GUI output
131 def processGUI(self, inframe, helper):
132 # Start a new display frame, gets its size and also whether mouse/keyboard are idle:
133 idle, winw, winh = helper.startFrame()
134
135 # Draw full-resolution input frame from camera:
136 x, y, w, h = helper.drawInputFrame("c", inframe, False, False)
137
138 # Get the next camera image at processing resolution (may block until it is captured):
139 frame = inframe.getCvRGBp() if self.rgb else inframe.getCvBGRp()
140
141 # Start measuring image processing time:
142 self.timer.start()
143
144 # Set the input:
145 image = Image.fromarray(frame);
146 _, scale = common.set_resized_input(self.interpreter, image.size,
147 lambda size: image.resize(size, Image.LANCZOS))
148 # Run the model
149 start = time.perf_counter()
150 self.interpreter.invoke()
151 inference_time = time.perf_counter() - start
152
153 # Get detections with high enough scores:
154 # JEVOIS: this used to work but now fails with pycoral-2.0.0 / libedgetpu-1.6.0 / tflite 2.5.0 with error:
155 # AttributeError: 'Interpreter' object has no attribute '_get_full_signature_list'
156 #objs = detect.get_objects(self.interpreter, self.threshold, scale)
157
158 # Use modified function below instead:
159 objs = self.get_objects(self.interpreter, self.threshold, scale)
160
161 # Draw the detections:
162 for obj in objs:
163 bbox = obj.bbox
164 label = self.labels.get(obj.id, obj.id)
165 col = self.stringToRGBA(label)
166 helper.drawRect(bbox.xmin, bbox.ymin, bbox.xmax, bbox.ymax, col, True)
167 helper.drawText(bbox.xmin+2, bbox.ymin+1, '%s: %.2f' % (label, obj.score), col)
168
169 # Output to serial:
170 #for obj in objs:
171 # print(self.labels.get(obj.id, obj.id))
172 # print(' id: ', obj.id)
173 # print(' score: ', obj.score)
174 # print(' bbox: ', obj.bbox)
175
176 # Write frames/s info from our timer:
177 fps = self.timer.stop()
178 helper.iinfo(inframe, fps, winw, winh);
179 helper.itext('JeVois-Pro Python Coral Detection - %s - %dms/inference' %
180 (self.model, inference_time * 1000.0))
181
182 # End of frame:
183 helper.endFrame()
184
185 # ####################################################################################################
186 ## Modified from https://github.com/google-coral/pycoral/blob/master/pycoral/adapters/detect.py
187 # to avoid calling interpreter._get_full_signature_list() which does not seem to exist anymore...
188 def get_objects(self, interpreter, score_threshold = -float('inf'), image_scale = (1.0, 1.0)):
189 """Gets results from a detection model as a list of detected objects.
190
191 Args:
192 interpreter: The ``tf.lite.Interpreter`` to query for results.
193 score_threshold (float): The score threshold for results. All returned
194 results have a score greater-than-or-equal-to this value.
195 image_scale (float, float): Scaling factor to apply to the bounding boxes as
196 (x-scale-factor, y-scale-factor), where each factor is from 0 to 1.0.
197
198 Returns:
199 A list of :obj:`Object` objects, which each contains the detected object's
200 id, score, and bounding box as :obj:`BBox`.
201 """
202 # If a model has signature, we use the signature output tensor names to parse
203 # the results. Otherwise, we parse the results based on some assumption of the
204 # output tensor order and size.
205 # pylint: disable=protected-access
206 # JEVOIS: this fails on pycoral 2.0.0/libedgetpu 1.6.0: signature_list = interpreter._get_full_signature_list()
207 # pylint: enable=protected-access
208 if False: ###signature_list:
209 if len(signature_list) > 1:
210 raise ValueError('Only support model with one signature.')
211 signature = signature_list[next(iter(signature_list))]
212 count = int(interpreter.tensor(signature['outputs']['output_0'])()[0])
213 scores = interpreter.tensor(signature['outputs']['output_1'])()[0]
214 class_ids = interpreter.tensor(signature['outputs']['output_2'])()[0]
215 boxes = interpreter.tensor(signature['outputs']['output_3'])()[0]
216 elif common.output_tensor(interpreter, 3).size == 1:
217 boxes = common.output_tensor(interpreter, 0)[0]
218 class_ids = common.output_tensor(interpreter, 1)[0]
219 scores = common.output_tensor(interpreter, 2)[0]
220 count = int(common.output_tensor(interpreter, 3)[0])
221 else:
222 scores = common.output_tensor(interpreter, 0)[0]
223 boxes = common.output_tensor(interpreter, 1)[0]
224 count = (int)(common.output_tensor(interpreter, 2)[0])
225 class_ids = common.output_tensor(interpreter, 3)[0]
226
227 width, height = common.input_size(interpreter)
228 image_scale_x, image_scale_y = image_scale
229 sx, sy = width / image_scale_x, height / image_scale_y
230
231 def make(i):
232 ymin, xmin, ymax, xmax = boxes[i]
233 return detect.Object(
234 id = int(class_ids[i]),
235 score = float(scores[i]),
236 bbox = detect.BBox(xmin=xmin, ymin=ymin, xmax=xmax, ymax=ymax).scale(sx, sy).map(int))
237
238 return [make(i) for i in range(count) if scores[i] >= score_threshold]
239
Object detection using Coral Edge TPU.
draw_objects(self, draw, objs, labels)
process(self, inframe, outframe)
JeVois main processing function.
processGUI(self, inframe, helper)
Process function with GUI output.
get_objects(self, interpreter, score_threshold=-float('inf'), image_scale=(1.0, 1.0))
Modified from https://github.com/google-coral/pycoral/blob/master/pycoral/adapters/detect....
__init__(self)
Constructor.
size_t getNumInstalledTPUs()