dlstreamer_sample.py · Intel/object-classification at main

object-classification / dlstreamer_sample.py

Update native RC1 samples and supported device documentation

0577537 verified about 10 hours ago

2.79 kB

# 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']}")

	# 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']}")