JeVoisBase  1.18
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
PyCoralSegment.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 pycoral.adapters import common
8 from pycoral.adapters import segment
9 from pycoral.utils.edgetpu import make_interpreter
10 import time
11 
12 ## Semantic segmentation using Coral Edge TPU
13 #
14 # More pre-trained models are available at https://coral.ai/models/
15 #
16 #
17 # @author Laurent Itti
18 #
19 # @videomapping YUYV 320 264 30.0 YUYV 320 240 30.0 JeVois PyCoralSegment
20 # @videomapping JVUI 0 0 30.0 CropScale=RGB24@512x288:YUYV 1920 1080 30.0 JeVois PyCoralSegment
21 # @email itti@usc.edu
22 # @address 880 W 1st St Suite 807, Los Angeles CA 90012, USA
23 # @copyright Copyright (C) 2020 by Laurent Itti
24 # @mainurl http://jevois.org
25 # @supporturl http://jevois.org
26 # @otherurl http://jevois.org
27 # @license GPL v3
28 # @distribution Unrestricted
29 # @restrictions None
30 # @ingroup modules
32  # ####################################################################################################
33  ## Constructor
34  def __init__(self):
35  self.rgb = True # True if model expects RGB inputs, otherwise it expects BGR
36  self.keepaspect = True # Keep aspect ratio using zero padding
37  alpha = 128 # Transparency alpha values for processGUI, higher is less transparent
38  tidx = 0 # Class index of transparent background
39 
40  # Select one of the models:
41  self.model = 'UNet128' # expects 128x128
42  #self.model = 'MobileNetV2DeepLabV3' # expects 513x513
43 
44  # You should not have to edit anything beyond this point.
45  if (self.model == 'MobileNetV2DeepLabV3'):
46  modelname = 'deeplabv3_mnv2_dm05_pascal_quant_edgetpu.tflite'
47  elif (self.model == 'UNet128'):
48  modelname = 'keras_post_training_unet_mv2_128_quant_edgetpu.tflite'
49  tidx = 1
50 
51  # Load network:
52  sdir = pyjevois.share + '/coral/segmentation/'
53  self.interpreter = make_interpreter(sdir + modelname)
54  #self.interpreter = make_interpreter(*modelname.split('@'))
55  self.interpreter.allocate_tensors()
56  self.timer = jevois.Timer('Coral segmentation', 10, jevois.LOG_DEBUG)
58  self.cmapRGBA = self.create_pascal_label_colormapRGBA(alpha, tidx)
59 
60  # ####################################################################################################
62  """Creates a label colormap used in PASCAL VOC segmentation benchmark.
63  Returns:
64  A Colormap for visualizing segmentation results.
65  """
66  colormap = np.zeros((256, 3), dtype=int)
67  indices = np.arange(256, dtype=int)
68 
69  for shift in reversed(range(8)):
70  for channel in range(3):
71  colormap[:, channel] |= ((indices >> channel) & 1) << shift
72  indices >>= 3
73 
74  return colormap.astype(np.uint8)
75 
76  # ####################################################################################################
77  def create_pascal_label_colormapRGBA(self, alpha, tidx):
78  """Creates a label colormap used in PASCAL VOC segmentation benchmark.
79  Returns:
80  A Colormap for visualizing segmentation results.
81  """
82  colormap = np.zeros((256, 4), dtype=int)
83  indices = np.arange(256, dtype=int)
84 
85  for shift in reversed(range(8)):
86  for channel in range(3):
87  colormap[:, channel] |= ((indices >> channel) & 1) << shift
88  indices >>= 3
89 
90  colormap[:, 3] = alpha
91  colormap[tidx, 3] = 0 # force fully transparent for entry tidx
92  return colormap.astype(np.uint8)
93 
94  # ####################################################################################################
95  ## JeVois main processing function
96  def process(self, inframe, outframe):
97  frame = inframe.getCvRGB() if self.rgb else inframe.getCvBGR()
98  self.timer.start()
99 
100  h = frame.shape[0]
101  w = frame.shape[1]
102 
103  # Set the input:
104  width, height = common.input_size(self.interpreter)
105  img = Image.fromarray(frame);
106  if self.keepaspect:
107  resized_img, _ = common.set_resized_input(self.interpreter, img.size,
108  lambda size: img.resize(size, Image.ANTIALIAS))
109  else:
110  resized_img = img.resize((width, height), Image.ANTIALIAS)
111  common.set_input(self.interpreter, resized_img)
112 
113  # Run the model
114  start = time.perf_counter()
115  self.interpreter.invoke()
116  inference_time = time.perf_counter() - start
117 
118  # Draw segmentation results:
119  result = segment.get_output(self.interpreter)
120  if len(result.shape) == 3: result = np.argmax(result, axis=-1)
121 
122  # If keep_aspect_ratio, we need to remove the padding area.
123  new_width, new_height = resized_img.size
124  result = result[:new_height, :new_width]
125  mask_img = Image.fromarray(self.cmapRGB[result])
126 
127  # Concat resized input image and processed segmentation results.
128  output_img = Image.new('RGB', (2 * img.width, img.height))
129  output_img.paste(img, (0, 0))
130  output_img.paste(mask_img.resize(img.size), (img.width, 0))
131 
132  # Back to opencv:
133  outcv = np.array(output_img)
134 
135  # Put efficiency information.
136  cv.putText(outcv, 'JeVois Coral Segmentation - ' + self.model, (3, 15),
137  cv.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv.LINE_AA)
138 
139  fps = self.timer.stop()
140  label = fps + ', %dms' % (inference_time * 1000.0)
141  cv.putText(outcv, label, (3, h-5), cv.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv.LINE_AA)
142 
143  # Send output frame to host:
144  if self.rgb: outframe.sendCvRGB(outcv)
145  else: outframe.sendCv(outcv)
146 
147  # ###################################################################################################
148  ## Process function with GUI output
149  def processGUI(self, inframe, helper):
150  # Start a new display frame, gets its size and also whether mouse/keyboard are idle:
151  idle, winw, winh = helper.startFrame()
152 
153  # Draw full-resolution input frame from camera:
154  x, y, w, h = helper.drawInputFrame("c", inframe, False, False)
155 
156  # Get the next camera image at processing resolution (may block until it is captured):
157  frame = inframe.getCvRGBp() if self.rgb else inframe.getCvBGRp()
158  iw, ih = frame.shape[1], frame.shape[0]
159 
160  # Start measuring image processing time:
161  self.timer.start()
162 
163  # Set the input:
164  width, height = common.input_size(self.interpreter)
165  img = Image.fromarray(frame);
166  if self.keepaspect:
167  resized_img, _ = common.set_resized_input(self.interpreter, img.size,
168  lambda size: img.resize(size, Image.ANTIALIAS))
169  else:
170  resized_img = img.resize((width, height), Image.ANTIALIAS)
171  common.set_input(self.interpreter, resized_img)
172 
173  # Run the model:
174  start = time.perf_counter()
175  self.interpreter.invoke()
176  inference_time = time.perf_counter() - start
177 
178  # Draw segmentation results:
179  result = segment.get_output(self.interpreter)
180  if len(result.shape) == 3: result = np.argmax(result, axis=-1)
181 
182  # If keep_aspect_ratio, we need to remove the padding area:
183  new_width, new_height = resized_img.size
184  result = result[:new_height, :new_width]
185  mask = self.cmapRGBA[result]
186 
187  # Draw the mask on top of our image, OpenGL will do the alpha blending:
188  helper.drawImage("m", mask, self.rgb, False, True)
189 
190  # Put efficiency information:
191  helper.itext('JeVois-Pro Python Coral Segmentation - %s - %dms/inference' %
192  (self.model, inference_time * 1000.0), 0, -1)
193 
194  # Write frames/s info from our timer:
195  fps = self.timer.stop()
196  helper.iinfo(inframe, fps, winw, winh);
197 
198  # End of frame:
199  helper.endFrame()
200 
PyCoralSegment.PyCoralSegment.process
def process(self, inframe, outframe)
JeVois main processing function.
Definition: PyCoralSegment.py:96
PyCoralSegment.PyCoralSegment.cmapRGB
cmapRGB
Definition: PyCoralSegment.py:57
PyCoralSegment.PyCoralSegment.create_pascal_label_colormap
def create_pascal_label_colormap(self)
Definition: PyCoralSegment.py:61
PyCoralSegment.PyCoralSegment.rgb
rgb
Definition: PyCoralSegment.py:35
PyCoralSegment.PyCoralSegment.model
model
Definition: PyCoralSegment.py:41
PyCoralSegment.PyCoralSegment.interpreter
interpreter
Definition: PyCoralSegment.py:53
PyCoralSegment.PyCoralSegment.cmapRGBA
cmapRGBA
Definition: PyCoralSegment.py:58
PyCoralSegment.PyCoralSegment.keepaspect
keepaspect
Definition: PyCoralSegment.py:36
PyCoralSegment.PyCoralSegment
Semantic segmentation using Coral Edge TPU.
Definition: PyCoralSegment.py:31
PyCoralSegment.PyCoralSegment.create_pascal_label_colormapRGBA
def create_pascal_label_colormapRGBA(self, alpha, tidx)
Definition: PyCoralSegment.py:77
PyCoralSegment.PyCoralSegment.__init__
def __init__(self)
Constructor.
Definition: PyCoralSegment.py:34
PyCoralSegment.PyCoralSegment.timer
timer
Definition: PyCoralSegment.py:56
PyCoralSegment.PyCoralSegment.processGUI
def processGUI(self, inframe, helper)
Process function with GUI output.
Definition: PyCoralSegment.py:149
jevois::Timer