JeVoisBase  1.18
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
PyCoralDetect.py
Go to the documentation of this file.
1 import pyjevois
2 if pyjevois.pro: import libjevoispro as jevois
3 else: import libjevois as jevois
4 import cv2 as cv
5 import numpy as np
6 from PIL import Image
7 from PIL import ImageDraw
8 from pycoral.adapters import common
9 from pycoral.adapters import detect
10 from pycoral.utils.dataset import read_label_file
11 from pycoral.utils.edgetpu import make_interpreter
12 import 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):
36  self.threshold = 0.4 # Confidence threshold (0..1), higher for stricter confidence.
37  self.rgb = True # True if model expects RGB inputs, otherwise it expects BGR
38 
39  # Select one of the models:
40  self.model = 'MobileDetSSD' # expects 320x320
41 
42  # You should not have to edit anything beyond this point.
43  if (self.model == 'MobileDetSSD'):
44  classnames = 'coco_labels.txt'
45  modelname = 'ssdlite_mobiledet_coco_qat_postprocess_edgetpu.tflite'
46 
47  # Load names of classes:
48  sdir = pyjevois.share + '/coral/detection/'
49  self.labels = read_label_file(sdir + classnames)
50 
51  # Load network:
52  self.interpreter = make_interpreter(sdir + modelname)
53  #self.interpreter = make_interpreter(*modelname.split('@'))
54  self.interpreter.allocate_tensors()
55  self.timer = jevois.Timer('Coral classification', 10, jevois.LOG_DEBUG)
56 
57  # ####################################################################################################
58  def stringToRGBA(self, str):
59  col = 0x80808080
60  alpha = 0xff
61  for c in str: col = ord(c) + ((col << 5) - col)
62  col = (col & 0xffffff) | (alpha << 24)
63  return col
64 
65  # ####################################################################################################
66  def draw_objects(self, draw, objs, labels):
67  """Draws the bounding box and label for each object."""
68  for obj in objs:
69  bbox = obj.bbox
70  draw.rectangle([(bbox.xmin, bbox.ymin), (bbox.xmax, bbox.ymax)], outline='red')
71  draw.text((bbox.xmin+10, bbox.ymin+10), '%s: %.2f' % (labels.get(obj.id, obj.id), obj.score), fill='red')
72 
73  # ####################################################################################################
74  ## JeVois main processing function
75  def process(self, inframe, outframe):
76  frame = inframe.getCvRGB() if self.rgb else inframe.getCvBGR()
77  self.timer.start()
78 
79  h = frame.shape[0]
80  w = frame.shape[1]
81 
82  # Set the input:
83  image = Image.fromarray(frame);
84  _, scale = common.set_resized_input(self.interpreter, image.size,
85  lambda size: image.resize(size, Image.ANTIALIAS))
86 
87  # Run the model
88  start = time.perf_counter()
89  self.interpreter.invoke()
90  inference_time = time.perf_counter() - start
91 
92  # Get detections with high enough scores:
93  objs = detect.get_objects(self.interpreter, self.threshold, scale)
94 
95  # Draw the detections:
96  image = image.convert('RGB')
97  self.draw_objects(ImageDraw.Draw(image), objs, self.labels)
98 
99  # Back to OpenCV:
100  frame = np.array(image)
101 
102  # Output to serial:
103  #for obj in objs:
104  # print(self.labels.get(obj.id, obj.id))
105  # print(' id: ', obj.id)
106  # print(' score: ', obj.score)
107  # print(' bbox: ', obj.bbox)
108 
109  # Put efficiency information:
110  cv.putText(frame, 'JeVois Coral Detection - ' + self.model, (3, 15),
111  cv.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv.LINE_AA)
112 
113  fps = self.timer.stop()
114  label = fps + ', %dms' % (inference_time * 1000.0)
115  cv.putText(frame, label, (3, h-5), cv.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv.LINE_AA)
116 
117  # Send output frame to host:
118  if self.rgb: outframe.sendCvRGB(frame)
119  else: outframe.sendCv(frame)
120 
121  # ####################################################################################################
122  ## Process function with GUI output
123  def processGUI(self, inframe, helper):
124  # Start a new display frame, gets its size and also whether mouse/keyboard are idle:
125  idle, winw, winh = helper.startFrame()
126 
127  # Draw full-resolution input frame from camera:
128  x, y, w, h = helper.drawInputFrame("c", inframe, False, False)
129 
130  # Get the next camera image at processing resolution (may block until it is captured):
131  frame = inframe.getCvRGBp() if self.rgb else inframe.getCvBGRp()
132 
133  # Start measuring image processing time:
134  self.timer.start()
135 
136  # Set the input:
137  image = Image.fromarray(frame);
138  _, scale = common.set_resized_input(self.interpreter, image.size,
139  lambda size: image.resize(size, Image.ANTIALIAS))
140  # Run the model
141  start = time.perf_counter()
142  self.interpreter.invoke()
143  inference_time = time.perf_counter() - start
144 
145  # Get detections with high enough scores:
146  objs = detect.get_objects(self.interpreter, self.threshold, scale)
147 
148  # Draw the detections:
149  for obj in objs:
150  bbox = obj.bbox
151  label = self.labels.get(obj.id, obj.id)
152  col = self.stringToRGBA(label)
153  helper.drawRect(bbox.xmin, bbox.ymin, bbox.xmax, bbox.ymax, col, True)
154  helper.drawText(bbox.xmin+2, bbox.ymin+1, '%s: %.2f' % (label, obj.score), col)
155 
156  # Output to serial:
157  #for obj in objs:
158  # print(self.labels.get(obj.id, obj.id))
159  # print(' id: ', obj.id)
160  # print(' score: ', obj.score)
161  # print(' bbox: ', obj.bbox)
162 
163  # Write frames/s info from our timer:
164  fps = self.timer.stop()
165  helper.iinfo(inframe, fps, winw, winh);
166  helper.itext('JeVois-Pro Python Coral Detection - %s - %dms/inference' %
167  (self.model, inference_time * 1000.0), 0, -1)
168 
169  # End of frame:
170  helper.endFrame()
PyCoralDetect.PyCoralDetect.__init__
def __init__(self)
Constructor.
Definition: PyCoralDetect.py:35
PyCoralDetect.PyCoralDetect.draw_objects
def draw_objects(self, draw, objs, labels)
Definition: PyCoralDetect.py:66
PyCoralDetect.PyCoralDetect.process
def process(self, inframe, outframe)
JeVois main processing function.
Definition: PyCoralDetect.py:75
PyCoralDetect.PyCoralDetect.timer
timer
Definition: PyCoralDetect.py:55
PyCoralDetect.PyCoralDetect.model
model
Definition: PyCoralDetect.py:40
PyCoralDetect.PyCoralDetect.threshold
threshold
Definition: PyCoralDetect.py:36
PyCoralDetect.PyCoralDetect
Object detection using Coral Edge TPU.
Definition: PyCoralDetect.py:32
PyCoralDetect.PyCoralDetect.stringToRGBA
def stringToRGBA(self, str)
Definition: PyCoralDetect.py:58
PyCoralDetect.PyCoralDetect.rgb
rgb
Definition: PyCoralDetect.py:37
PyCoralDetect.PyCoralDetect.interpreter
interpreter
Definition: PyCoralDetect.py:52
PyCoralDetect.PyCoralDetect.processGUI
def processGUI(self, inframe, helper)
Process function with GUI output.
Definition: PyCoralDetect.py:123
PyCoralDetect.PyCoralDetect.labels
labels
Definition: PyCoralDetect.py:49
jevois::Timer