| |
| |
| |
| |
| |
|
|
| 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" |
| |
| DEVICE = "GPU" |
| DET_THRESHOLD = 0.6 |
| ATTR_THRESHOLD = 0.5 |
|
|
| |
| |
| ATTRIBUTE_NAMES = [ |
| "is_male", "has_bag", "has_backpack", "has_hat", |
| "has_longsleeves", "has_longpants", "has_longhair", "has_coat_jacket", |
| ] |
|
|
| |
| |
| 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) |
| |
| |
| 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 |
|
|
|
|
| |
| |
| 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) |
|
|
