File size: 2,789 Bytes
0577537 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | # SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Object Classification: 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([])
INPUT_VIDEO = "test_video.mp4"
# Traffic-relevant COCO labels grouped into higher-level city categories.
CATEGORY_BY_LABEL = {
"person": "People",
"bicycle": "Vehicles",
"car": "Vehicles",
"motorcycle": "Vehicles",
"bus": "Vehicles",
"train": "Vehicles",
"truck": "Vehicles",
}
# Select the model: "yolo26n" (default) or "detr-resnet-50".
MODEL = "yolo26n"
# For CPU: change device=GPU to device=CPU.
# For NPU: change device=GPU to device=NPU (batch-size=1, nireq=4 recommended).
# gvawatermark displ-cfg:
# show-roi=... draws only the traffic-relevant classes (person + vehicles),
# so non-traffic detections such as handbag/backpack are not boxed.
# font-scale=1.5 enlarges the label text for better visualization.
pipeline_str = (
f"filesrc location={INPUT_VIDEO} ! decodebin3 ! "
"videoconvert ! "
f"gvadetect model={MODEL}_openvino_model/{MODEL}.xml "
"device=GPU "
"threshold=0.4 ! queue ! "
"gvawatermark "
"displ-cfg=show-roi=person:bicycle:car:motorcycle:bus:train:truck,font-scale=2.5 ! "
"videoconvert ! video/x-raw,format=I420 ! "
"openh264enc ! h264parse ! "
"mp4mux ! filesink name=sink location=output_dlstreamer.mp4"
)
pipeline = Gst.parse_launch(pipeline_str)
totals = {"People": 0, "Vehicles": 0}
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
counts = {"People": 0, "Vehicles": 0}
idx = 1
while True:
ok, od = rmeta.get_od_mtd(idx)
if not ok:
break
label = GLib.quark_to_string(od.get_obj_type())
category = CATEGORY_BY_LABEL.get(label)
if category is not None:
counts[category] += 1
totals[category] += 1
idx += 1
if counts["People"] or counts["Vehicles"]:
print(f"frame: People={counts['People']} Vehicles={counts['Vehicles']}",
flush=True)
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)
print(f"Summary: People={totals['People']} Vehicles={totals['Vehicles']}")
|