From a557a8598a41366e73077b6ddcb1d374bb742d8b Mon Sep 17 00:00:00 2001 From: Aleix Conchillo Flaque Date: Fri, 16 Aug 2019 22:20:54 -0700 Subject: [PATCH] python: added jetson.utils.VideoLoader class --- python/python/jetson/utils/__init__.py | 2 + python/python/jetson/utils/video_loader.py | 155 +++++++++++++++++++++ 2 files changed, 157 insertions(+) create mode 100644 python/python/jetson/utils/video_loader.py diff --git a/python/python/jetson/utils/__init__.py b/python/python/jetson/utils/__init__.py index 2a938316..8a753647 100644 --- a/python/python/jetson/utils/__init__.py +++ b/python/python/jetson/utils/__init__.py @@ -3,4 +3,6 @@ from jetson_utils_python import * +from .video_loader import * + VERSION = '1.0.0' diff --git a/python/python/jetson/utils/video_loader.py b/python/python/jetson/utils/video_loader.py new file mode 100644 index 00000000..d19b0c70 --- /dev/null +++ b/python/python/jetson/utils/video_loader.py @@ -0,0 +1,155 @@ +# +# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +# + +import queue + +import gi +gi.require_version("Gst", "1.0") +from gi.repository import Gst + +Gst.init(None) + +class VideoLoader: + + def __init__(self, logger, width=1280, height=720, framerate="5/1"): + self.__logger = logger + self.__width = width + self.__height = height + self.__framerate = framerate + + self.__pipeline = None + self.__queue = queue.Queue() + + def load(self, filename): + if self.__pipeline: + raise Exception("Already processing a file") + self.__pipeline = Gst.Pipeline.new("video-loader") + + bus = self.__pipeline.get_bus() + bus.add_signal_watch() + bus.connect("message", self.__on_message) + + self.__build_pipeline(filename) + self.__pipeline.set_state(Gst.State.PLAYING) + + + def stop(self): + if self.__pipeline: + self.__pipeline.set_state(Gst.State.NULL) + self.__pipeline = None + + def is_loading(self): + return self.__pipeline != None + + def sample_next(self): + return self.__queue.get() + + def sample_done(self): + self.__queue.task_done() + + def __build_pipeline(self, filename): + source = Gst.ElementFactory.make("filesrc") + source.set_property("location", filename) + self.__pipeline.add(source) + + decodebin = Gst.ElementFactory.make("decodebin") + decodebin.connect("pad-added", self.__on_pad_added) + self.__pipeline.add(decodebin) + + source.link(decodebin) + + def __on_message(self, bus, message): + t = message.type + if t == Gst.MessageType.EOS: + self.stop() + elif t == Gst.MessageType.ERROR: + self.stop() + err, debug = message.parse_error() + self.__logger.error("Error: %s (%s)" % (err, debug)) + + def __on_pad_added(self, decodebin, pad): + caps_str = pad.query_caps(None).to_string() + if caps_str.startswith("video"): + self.__logger.info("Detected video stream, processing...") + elif caps_str.startswith("audio"): + self.__logger.info("Detected audio stream, ignoring.") + return + + queue = Gst.ElementFactory.make("queue") + self.__pipeline.add(queue) + + # rate first so we dn't process as many frames. + videorate = Gst.ElementFactory.make("videorate") + self.__pipeline.add(videorate) + + caps = "video/x-raw,framerate=%s" % self.__framerate + ratecaps = Gst.ElementFactory.make("capsfilter") + ratecaps.set_property("caps", Gst.Caps(caps)) + self.__pipeline.add(ratecaps) + + # scale image in GPU if needed. + videoscale = Gst.ElementFactory.make("nvvidconv") + self.__pipeline.add(videoscale) + + caps = "video/x-raw(memory:NVMM),format=NV12" + if self.__width > 0: + caps += ",width=%d" % self.__width + if self.__height > 0: + caps += ",height=%d" % self.__height + + scalecaps = Gst.ElementFactory.make("capsfilter") + scalecaps.set_property("caps", Gst.Caps(caps)) + self.__pipeline.add(scalecaps) + + # copy to main memory. + videoconvert = Gst.ElementFactory.make("nvvidconv") + self.__pipeline.add(videoconvert) + + convcaps = Gst.ElementFactory.make("capsfilter") + convcaps.set_property("caps", Gst.Caps("video/x-raw")) + self.__pipeline.add(convcaps) + + # finally, add the sample sink. + appsink = Gst.ElementFactory.make("appsink") + appsink.set_property("emit-signals", True) + appsink.set_property("max-buffers", 10) + appsink.set_property("drop", True) + appsink.set_property("sync", True) + appsink.connect("new-sample", self.__on_new_sample) + self.__pipeline.add(appsink) + + # link everything together. + sinkpad = queue.get_static_pad("sink") + pad.link(sinkpad) + + queue.link(videorate) + videorate.link(ratecaps) + ratecaps.link(videoscale) + videoscale.link(scalecaps) + scalecaps.link(videoconvert) + videoconvert.link(convcaps) + convcaps.link(appsink) + + def __on_new_sample(self, appsink): + sample = appsink.emit("pull-sample") + self.__queue.put(sample) + return Gst.FlowReturn.OK