worker-safety-detection / dlstreamer_sample.py
vagheshpatel's picture
Update native RC1 samples and supported device documentation
7dec3f8 verified
Raw History Blame Contribute Delete
1.82 kB
# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Worker Safety Detection: DLStreamer pipeline sample.
# See README.md for prerequisites and inputs. Run: python3 dlstreamer_sample.py
import gi
gi.require_version("Gst", "1.0")
gi.require_version("GstAnalytics", "1.0")
from gi.repository import Gst, GLib, GstAnalytics
Gst.init([])
MODEL_XML = (
"models/worker_safety_gear_detection/"
"deployment/Detection/model/model.xml"
)
INPUT_VIDEO = "test_video.avi"
# For CPU: change device=GPU to device=CPU.
# For NPU: change device=GPU to device=NPU (batch-size=1, nireq=4 recommended).
pipeline_str = (
f"filesrc location={INPUT_VIDEO} ! decodebin3 ! "
f"videoconvert ! "
f"gvadetect model={MODEL_XML} device=GPU "
f"threshold=0.4 ! queue ! "
f"gvawatermark ! videoconvert ! video/x-raw,format=I420 ! "
f"openh264enc ! h264parse ! "
f"mp4mux ! filesink name=sink location=output_dlstreamer.mp4"
)
pipeline = Gst.parse_launch(pipeline_str)
def on_buffer(pad, info):
buf = info.get_buffer()
rmeta = GstAnalytics.buffer_get_analytics_relation_meta(buf)
if rmeta is None:
return Gst.PadProbeReturn.OK
idx = 1
while True:
ok, od = rmeta.get_od_mtd(idx)
if not ok:
break
label = GLib.quark_to_string(od.get_obj_type())
_, x, y, w, h, conf = od.get_location()
print(f" [PPE] {label} conf={conf:.2f}", flush=True)
idx += 1
return Gst.PadProbeReturn.OK
sink = pipeline.get_by_name("sink")
sink_pad = sink.get_static_pad("sink")
sink_pad.add_probe(Gst.PadProbeType.BUFFER, on_buffer)
pipeline.set_state(Gst.State.PLAYING)
bus = pipeline.get_bus()
bus.timed_pop_filtered(
Gst.CLOCK_TIME_NONE,
Gst.MessageType.EOS | Gst.MessageType.ERROR,
)
pipeline.set_state(Gst.State.NULL)