| 1 | # Copyright 2019 Google LLC
|
|---|
| 2 | #
|
|---|
| 3 | # Licensed under the Apache License, Version 2.0 (the 'License');
|
|---|
| 4 | # you may not use this file except in compliance with the License.
|
|---|
| 5 | # You may obtain a copy of the License at
|
|---|
| 6 | #
|
|---|
| 7 | # https://www.apache.org/licenses/LICENSE-2.0
|
|---|
| 8 | #
|
|---|
| 9 | # Unless required by applicable law or agreed to in writing, software
|
|---|
| 10 | # distributed under the License is distributed on an 'AS IS' BASIS,
|
|---|
| 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|---|
| 12 | # See the License for the specific language governing permissions and
|
|---|
| 13 | # limitations under the License.
|
|---|
| 14 |
|
|---|
| 15 | import sys
|
|---|
| 16 | import svgwrite
|
|---|
| 17 | import threading
|
|---|
| 18 |
|
|---|
| 19 | import gi
|
|---|
| 20 | gi.require_version('Gst', '1.0')
|
|---|
| 21 | gi.require_version('GstBase', '1.0')
|
|---|
| 22 | gi.require_version('Gtk', '3.0')
|
|---|
| 23 | from gi.repository import GLib, GObject, Gst, GstBase, Gtk
|
|---|
| 24 |
|
|---|
| 25 | GObject.threads_init()
|
|---|
| 26 | Gst.init(None)
|
|---|
| 27 |
|
|---|
| 28 | class GstPipeline:
|
|---|
| 29 | def __init__(self, pipeline, user_function, src_size):
|
|---|
| 30 | self.user_function = user_function
|
|---|
| 31 | self.running = False
|
|---|
| 32 | self.gstbuffer = None
|
|---|
| 33 | self.sink_size = None
|
|---|
| 34 | self.src_size = src_size
|
|---|
| 35 | self.box = None
|
|---|
| 36 | self.condition = threading.Condition()
|
|---|
| 37 |
|
|---|
| 38 | self.pipeline = Gst.parse_launch(pipeline)
|
|---|
| 39 | self.overlay = self.pipeline.get_by_name('overlay')
|
|---|
| 40 | self.overlaysink = self.pipeline.get_by_name('overlaysink')
|
|---|
| 41 | appsink = self.pipeline.get_by_name('appsink')
|
|---|
| 42 | appsink.connect('new-sample', self.on_new_sample)
|
|---|
| 43 |
|
|---|
| 44 | # Set up a pipeline bus watch to catch errors.
|
|---|
| 45 | bus = self.pipeline.get_bus()
|
|---|
| 46 | bus.add_signal_watch()
|
|---|
| 47 | bus.connect('message', self.on_bus_message)
|
|---|
| 48 |
|
|---|
| 49 | # Set up a full screen window on Coral, no-op otherwise.
|
|---|
| 50 | self.setup_window()
|
|---|
| 51 |
|
|---|
| 52 | def run(self):
|
|---|
| 53 | # Start inference worker.
|
|---|
| 54 | self.running = True
|
|---|
| 55 | worker = threading.Thread(target=self.inference_loop)
|
|---|
| 56 | worker.start()
|
|---|
| 57 |
|
|---|
| 58 | # Run pipeline.
|
|---|
| 59 | self.pipeline.set_state(Gst.State.PLAYING)
|
|---|
| 60 | try:
|
|---|
| 61 | Gtk.main()
|
|---|
| 62 | except:
|
|---|
| 63 | pass
|
|---|
| 64 |
|
|---|
| 65 | # Clean up.
|
|---|
| 66 | self.pipeline.set_state(Gst.State.NULL)
|
|---|
| 67 | while GLib.MainContext.default().iteration(False):
|
|---|
| 68 | pass
|
|---|
| 69 | with self.condition:
|
|---|
| 70 | self.running = False
|
|---|
| 71 | self.condition.notify_all()
|
|---|
| 72 | worker.join()
|
|---|
| 73 |
|
|---|
| 74 | def on_bus_message(self, bus, message):
|
|---|
| 75 | t = message.type
|
|---|
| 76 | if t == Gst.MessageType.EOS:
|
|---|
| 77 | Gtk.main_quit()
|
|---|
| 78 | elif t == Gst.MessageType.WARNING:
|
|---|
| 79 | err, debug = message.parse_warning()
|
|---|
| 80 | sys.stderr.write('Warning: %s: %s\n' % (err, debug))
|
|---|
| 81 | elif t == Gst.MessageType.ERROR:
|
|---|
| 82 | err, debug = message.parse_error()
|
|---|
| 83 | sys.stderr.write('Error: %s: %s\n' % (err, debug))
|
|---|
| 84 | Gtk.main_quit()
|
|---|
| 85 | return True
|
|---|
| 86 |
|
|---|
| 87 | def on_new_sample(self, sink):
|
|---|
| 88 | sample = sink.emit('pull-sample')
|
|---|
| 89 | if not self.sink_size:
|
|---|
| 90 | s = sample.get_caps().get_structure(0)
|
|---|
| 91 | self.sink_size = (s.get_value('width'), s.get_value('height'))
|
|---|
| 92 | with self.condition:
|
|---|
| 93 | self.gstbuffer = sample.get_buffer()
|
|---|
| 94 | self.condition.notify_all()
|
|---|
| 95 | return Gst.FlowReturn.OK
|
|---|
| 96 |
|
|---|
| 97 | def get_box(self):
|
|---|
| 98 | if not self.box:
|
|---|
| 99 | glbox = self.pipeline.get_by_name('glbox')
|
|---|
| 100 | if glbox:
|
|---|
| 101 | glbox = glbox.get_by_name('filter')
|
|---|
| 102 | box = self.pipeline.get_by_name('box')
|
|---|
| 103 | assert glbox or box
|
|---|
| 104 | assert self.sink_size
|
|---|
| 105 | if glbox:
|
|---|
| 106 | self.box = (glbox.get_property('x'), glbox.get_property('y'),
|
|---|
| 107 | glbox.get_property('width'), glbox.get_property('height'))
|
|---|
| 108 | else:
|
|---|
| 109 | self.box = (-box.get_property('left'), -box.get_property('top'),
|
|---|
| 110 | self.sink_size[0] + box.get_property('left') + box.get_property('right'),
|
|---|
| 111 | self.sink_size[1] + box.get_property('top') + box.get_property('bottom'))
|
|---|
| 112 | return self.box
|
|---|
| 113 |
|
|---|
| 114 | def inference_loop(self):
|
|---|
| 115 | while True:
|
|---|
| 116 | with self.condition:
|
|---|
| 117 | while not self.gstbuffer and self.running:
|
|---|
| 118 | self.condition.wait()
|
|---|
| 119 | if not self.running:
|
|---|
| 120 | break
|
|---|
| 121 | gstbuffer = self.gstbuffer
|
|---|
| 122 | self.gstbuffer = None
|
|---|
| 123 |
|
|---|
| 124 | # Passing Gst.Buffer as input tensor avoids 2 copies of it:
|
|---|
| 125 | # * Python bindings copies the data when mapping gstbuffer
|
|---|
| 126 | # * Numpy copies the data when creating ndarray.
|
|---|
| 127 | # This requires a recent version of the python3-edgetpu package. If this
|
|---|
| 128 | # raises an exception please make sure dependencies are up to date.
|
|---|
| 129 | input_tensor = gstbuffer
|
|---|
| 130 | svg = self.user_function(input_tensor, self.src_size, self.get_box())
|
|---|
| 131 | if svg:
|
|---|
| 132 | if self.overlay:
|
|---|
| 133 | self.overlay.set_property('data', svg)
|
|---|
| 134 | if self.overlaysink:
|
|---|
| 135 | self.overlaysink.set_property('svg', svg)
|
|---|
| 136 |
|
|---|
| 137 | def setup_window(self):
|
|---|
| 138 | # Only set up our own window if we have Coral overlay sink in the pipeline.
|
|---|
| 139 | if not self.overlaysink:
|
|---|
| 140 | return
|
|---|
| 141 |
|
|---|
| 142 | gi.require_version('GstGL', '1.0')
|
|---|
| 143 | gi.require_version('GstVideo', '1.0')
|
|---|
| 144 | from gi.repository import GstGL, GstVideo
|
|---|
| 145 |
|
|---|
| 146 | # Needed to commit the wayland sub-surface.
|
|---|
| 147 | def on_gl_draw(sink, widget):
|
|---|
| 148 | widget.queue_draw()
|
|---|
| 149 |
|
|---|
| 150 | # Needed to account for window chrome etc.
|
|---|
| 151 | def on_widget_configure(widget, event, overlaysink):
|
|---|
| 152 | allocation = widget.get_allocation()
|
|---|
| 153 | overlaysink.set_render_rectangle(allocation.x, allocation.y,
|
|---|
| 154 | allocation.width, allocation.height)
|
|---|
| 155 | return False
|
|---|
| 156 |
|
|---|
| 157 | window = Gtk.Window(Gtk.WindowType.TOPLEVEL)
|
|---|
| 158 | window.fullscreen()
|
|---|
| 159 |
|
|---|
| 160 | drawing_area = Gtk.DrawingArea()
|
|---|
| 161 | window.add(drawing_area)
|
|---|
| 162 | drawing_area.realize()
|
|---|
| 163 |
|
|---|
| 164 | self.overlaysink.connect('drawn', on_gl_draw, drawing_area)
|
|---|
| 165 |
|
|---|
| 166 | # Wayland window handle.
|
|---|
| 167 | wl_handle = self.overlaysink.get_wayland_window_handle(drawing_area)
|
|---|
| 168 | self.overlaysink.set_window_handle(wl_handle)
|
|---|
| 169 |
|
|---|
| 170 | # Wayland display context wrapped as a GStreamer context.
|
|---|
| 171 | wl_display = self.overlaysink.get_default_wayland_display_context()
|
|---|
| 172 | self.overlaysink.set_context(wl_display)
|
|---|
| 173 |
|
|---|
| 174 | drawing_area.connect('configure-event', on_widget_configure, self.overlaysink)
|
|---|
| 175 | window.connect('delete-event', Gtk.main_quit)
|
|---|
| 176 | window.show_all()
|
|---|
| 177 |
|
|---|
| 178 | # The appsink pipeline branch must use the same GL display as the screen
|
|---|
| 179 | # rendering so they get the same GL context. This isn't automatically handled
|
|---|
| 180 | # by GStreamer as we're the ones setting an external display handle.
|
|---|
| 181 | def on_bus_message_sync(bus, message, overlaysink):
|
|---|
| 182 | if message.type == Gst.MessageType.NEED_CONTEXT:
|
|---|
| 183 | _, context_type = message.parse_context_type()
|
|---|
| 184 | if context_type == GstGL.GL_DISPLAY_CONTEXT_TYPE:
|
|---|
| 185 | sinkelement = overlaysink.get_by_interface(GstVideo.VideoOverlay)
|
|---|
| 186 | gl_context = sinkelement.get_property('context')
|
|---|
| 187 | if gl_context:
|
|---|
| 188 | display_context = Gst.Context.new(GstGL.GL_DISPLAY_CONTEXT_TYPE, True)
|
|---|
| 189 | GstGL.context_set_gl_display(display_context, gl_context.get_display())
|
|---|
| 190 | message.src.set_context(display_context)
|
|---|
| 191 | return Gst.BusSyncReply.PASS
|
|---|
| 192 |
|
|---|
| 193 | bus = self.pipeline.get_bus()
|
|---|
| 194 | bus.set_sync_handler(on_bus_message_sync, self.overlaysink)
|
|---|
| 195 |
|
|---|
| 196 | def detectCoralDevBoard():
|
|---|
| 197 | try:
|
|---|
| 198 | if 'MX8MQ' in open('/sys/firmware/devicetree/base/model').read():
|
|---|
| 199 | print('Detected Edge TPU dev board.')
|
|---|
| 200 | return True
|
|---|
| 201 | except: pass
|
|---|
| 202 | return False
|
|---|
| 203 |
|
|---|
| 204 | def run_pipeline(user_function,
|
|---|
| 205 | src_size,
|
|---|
| 206 | appsink_size,
|
|---|
| 207 | videosrc='/dev/video1',
|
|---|
| 208 | videofmt='raw'):
|
|---|
| 209 | if videofmt == 'h264':
|
|---|
| 210 | SRC_CAPS = 'video/x-h264,width={width},height={height},framerate=30/1'
|
|---|
| 211 | elif videofmt == 'jpeg':
|
|---|
| 212 | SRC_CAPS = 'image/jpeg,width={width},height={height},framerate=30/1'
|
|---|
| 213 | else:
|
|---|
| 214 | SRC_CAPS = 'video/x-raw,width={width},height={height},framerate=30/1'
|
|---|
| 215 | PIPELINE = 'v4l2src device=%s ! {src_caps}'%videosrc
|
|---|
| 216 |
|
|---|
| 217 | if detectCoralDevBoard():
|
|---|
| 218 | scale_caps = None
|
|---|
| 219 | PIPELINE += """ ! decodebin ! glupload ! tee name=t
|
|---|
| 220 | t. ! queue ! glfilterbin filter=glbox name=glbox ! {sink_caps} ! {sink_element}
|
|---|
| 221 | t. ! queue ! glsvgoverlaysink name=overlaysink
|
|---|
| 222 | """
|
|---|
| 223 | else:
|
|---|
| 224 | scale = min(appsink_size[0] / src_size[0], appsink_size[1] / src_size[1])
|
|---|
| 225 | scale = tuple(int(x * scale) for x in src_size)
|
|---|
| 226 | scale_caps = 'video/x-raw,width={width},height={height}'.format(width=scale[0], height=scale[1])
|
|---|
| 227 | PIPELINE += """ ! tee name=t
|
|---|
| 228 | t. ! {leaky_q} ! videoconvert ! videoscale ! {scale_caps} ! videobox name=box autocrop=true
|
|---|
| 229 | ! {sink_caps} ! {sink_element}
|
|---|
| 230 | t. ! {leaky_q} ! videoconvert
|
|---|
| 231 | ! rsvgoverlay name=overlay ! videoconvert ! jpegenc ! tcpclientsink host=172.24.24.93 port=9001
|
|---|
| 232 | """
|
|---|
| 233 |
|
|---|
| 234 | SINK_ELEMENT = 'appsink name=appsink emit-signals=true max-buffers=1 drop=true'
|
|---|
| 235 | SINK_CAPS = 'video/x-raw,format=RGB,width={width},height={height}'
|
|---|
| 236 | LEAKY_Q = 'queue'
|
|---|
| 237 |
|
|---|
| 238 | src_caps = SRC_CAPS.format(width=src_size[0], height=src_size[1])
|
|---|
| 239 | sink_caps = SINK_CAPS.format(width=appsink_size[0], height=appsink_size[1])
|
|---|
| 240 | pipeline = PIPELINE.format(leaky_q=LEAKY_Q,
|
|---|
| 241 | src_caps=src_caps, sink_caps=sink_caps,
|
|---|
| 242 | sink_element=SINK_ELEMENT, scale_caps=scale_caps)
|
|---|
| 243 |
|
|---|
| 244 | print('Gstreamer pipeline:\n', pipeline)
|
|---|
| 245 |
|
|---|
| 246 | pipeline = GstPipeline(pipeline, user_function, src_size)
|
|---|
| 247 | pipeline.run()
|
|---|
