# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Object Counting: DLStreamer pipeline sample.
# See README.md for prerequisites and inputs. Run: python3 dlstreamer_sample.py

from collections import Counter

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"

# Only these classes are counted; the sample video contains pedestrians,
# a cyclist, and cars.
CLASSES_OF_INTEREST = {"person", "bicycle", "car"}

# 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).
pipeline_str = (
    f"filesrc location={INPUT_VIDEO} ! decodebin3 ! "
    "videoconvert ! "
    f"gvadetect model={MODEL}_openvino_model/{MODEL}.xml "
    "device=GPU "
    "threshold=0.4 ! queue ! "
    "gvawatermark ! videoconvert ! video/x-raw,format=I420 ! "
    "openh264enc ! h264parse ! "
    "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
    counts = Counter()
    idx = 1
    while True:
        ok, od = rmeta.get_od_mtd(idx)
        if not ok:
            break
        label = GLib.quark_to_string(od.get_obj_type())
        if label in CLASSES_OF_INTEREST:
            counts[label] += 1
        idx += 1
    if counts:
        summary = ", ".join(f"{n} {name}" for name, n in counts.items())
        print(f"Object counts: {summary}", flush=True)
    return Gst.PadProbeReturn.OK


sink = pipeline.get_by_name("sink")
sink.get_static_pad("sink").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)