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

import gi

gi.require_version("Gst", "1.0")
from gi.repository import Gst

Gst.init([])

import numpy as np
import cv2
from gstgva import VideoFrame

INPUT_VIDEO = "test_video.mp4"
OUTPUT_VIDEO = "output_dlstreamer.mp4"
DETECTION_MODEL = "intel/person-detection-retail-0013/FP16/person-detection-retail-0013.xml"
ATTRIBUTES_MODEL = "intel/person-attributes-recognition-crossroad-0230/FP16/person-attributes-recognition-crossroad-0230.xml"
# For CPU: change "GPU" to "CPU". For NPU: change "GPU" to "NPU".
DEVICE = "GPU"
DET_THRESHOLD = 0.6
ATTR_THRESHOLD = 0.5

# person-attributes-recognition-crossroad-0230 emits eight binary appearance
# attributes (output layer "453"), in this order:
ATTRIBUTE_NAMES = [
    "is_male", "has_bag", "has_backpack", "has_hat",
    "has_longsleeves", "has_longpants", "has_longhair", "has_coat_jacket",
]

# The appearance being searched for. Set an attribute to True to require it or
# False to require its absence; omit attributes you do not care about.
QUERY = {"has_longhair": True}


def person_attributes(video_frame):
    """Yield ((x, y, w, h), {attribute: score}) for each classified person."""
    for region in video_frame.regions():
        rect = region.rect()
        scores = None
        for tensor in region.tensors():
            if tensor.is_detection():
                continue
            data = np.array(tensor.data(), dtype=np.float32)
            # The attributes vector is the length-8 output ("453"); the model
            # also emits two length-2 colour points that are ignored here.
            if data.size == len(ATTRIBUTE_NAMES):
                scores = {name: float(data[i])
                          for i, name in enumerate(ATTRIBUTE_NAMES)}
        if scores is None:
            continue
        yield (int(rect.x), int(rect.y), int(rect.w), int(rect.h)), scores


def matches_query(scores):
    return all((scores[name] >= ATTR_THRESHOLD) == wanted
               for name, wanted in QUERY.items())


query_text = ", ".join(k if v else f"no {k}" for k, v in QUERY.items())
writer = {"w": None}
match_frames = 0


def on_video(sink):
    global match_frames
    sample = sink.emit("pull-sample")
    if sample is None:
        return Gst.FlowReturn.OK
    vf = VideoFrame(sample.get_buffer(), caps=sample.get_caps())
    matches = [(x, y, w, h) for (x, y, w, h), scores in person_attributes(vf)
               if matches_query(scores)]

    with vf.data() as mat:
        frame = mat.copy()

    if writer["w"] is None:
        frame_h, frame_w = frame.shape[:2]
        structure = sample.get_caps().get_structure(0)
        ok_fr, fps_n, fps_d = structure.get_fraction("framerate")
        fps = fps_n / fps_d if ok_fr and fps_d else 12
        writer["w"] = cv2.VideoWriter(
            OUTPUT_VIDEO, cv2.VideoWriter_fourcc(*"mp4v"), fps, (frame_w, frame_h))

    for x, y, w, h in matches:
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
        cv2.putText(frame, f"MATCH: {query_text}", (x, max(15, y - 8)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
    if matches:
        match_frames += 1
    writer["w"].write(frame)
    return Gst.FlowReturn.OK


# Convert to BGR before inference so gvaclassify's attribute tensors survive to
# the appsink (a later format-changing videoconvert would strip them).
pipeline = Gst.parse_launch(
    f"filesrc location={INPUT_VIDEO} ! decodebin3 ! videoconvert ! "
    f"video/x-raw,format=BGR ! "
    f"gvadetect model={DETECTION_MODEL} device={DEVICE} "
    f"threshold={DET_THRESHOLD} ! queue ! "
    f"gvaclassify model={ATTRIBUTES_MODEL} device={DEVICE} ! queue ! "
    "appsink name=sink emit-signals=true sync=false max-buffers=4 drop=false"
)
sink = pipeline.get_by_name("sink")
sink.connect("new-sample", on_video)
pipeline.set_state(Gst.State.PLAYING)
pipeline.get_bus().timed_pop_filtered(
    Gst.CLOCK_TIME_NONE, Gst.MessageType.EOS | Gst.MessageType.ERROR)
pipeline.set_state(Gst.State.NULL)
if writer["w"] is not None:
    writer["w"].release()
print(f"Frames with a person matching [{query_text}]: {match_frames}", flush=True)
print(f"Saved: {OUTPUT_VIDEO}", flush=True)