File size: 4,266 Bytes
a82c780 | 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | # SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Heatmap Generation: 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
import numpy as np
Gst.init([])
# Import cv2 after Gst.init to avoid GStreamer re-initialization conflicts.
import cv2
INPUT_VIDEO = "test_video.mp4"
HEATMAP_ALPHA = 0.55
MATCH_DIST = 80.0 # max px between frames to treat detections as the same object
MAX_SPEED = 20.0 # px/frame that maps to full red
def render_speed_heatmap(frame, speed_sum, count, alpha):
"""Color traffic by average speed: blue = slow, yellow = medium,
red = fast. Only regions where vehicles were seen are tinted, so
empty background keeps its original color."""
avg = np.zeros_like(speed_sum)
seen = count > 0
avg[seen] = speed_sum[seen] / count[seen]
avg = cv2.GaussianBlur(avg, (0, 0), sigmaX=15)
presence = cv2.GaussianBlur(seen.astype(np.float32), (0, 0), sigmaX=15)
norm = np.clip(avg / MAX_SPEED, 0, 1) # 0 = slow (blue), 1 = fast (red)
color = cv2.applyColorMap((norm * 255).astype(np.uint8), cv2.COLORMAP_JET)
weight = (np.clip(presence, 0, 1) * alpha)[..., np.newaxis]
overlay = frame.astype(np.float32) * (1 - weight) + color.astype(np.float32) * weight
return overlay.astype(np.uint8)
# 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 ! video/x-raw,format=BGR ! "
f"gvadetect model={MODEL}_openvino_model/{MODEL}.xml "
"device=GPU "
"threshold=0.4 ! queue ! "
"appsink name=sink emit-signals=false sync=false"
)
pipeline = Gst.parse_launch(pipeline_str)
sink = pipeline.get_by_name("sink")
pipeline.set_state(Gst.State.PLAYING)
speed_sum = None
count = None
prev_centroids = []
writer = None
frame_idx = 0
total_dets = 0
while True:
sample = sink.emit("pull-sample")
if sample is None:
break
buf = sample.get_buffer()
caps = sample.get_caps().get_structure(0)
width = caps.get_value("width")
height = caps.get_value("height")
if speed_sum is None:
speed_sum = np.zeros((height, width), dtype=np.float32)
count = np.zeros((height, width), dtype=np.float32)
ok, mapinfo = buf.map(Gst.MapFlags.READ)
if not ok:
continue
frame = np.ndarray((height, width, 3), dtype=np.uint8,
buffer=mapinfo.data).copy()
buf.unmap(mapinfo)
frame_idx += 1
rmeta = GstAnalytics.buffer_get_analytics_relation_meta(buf)
cur_centroids = []
det_count = 0
if rmeta is not None:
idx = 1
while True:
ok_od, od = rmeta.get_od_mtd(idx)
if not ok_od:
break
_, x, y, w, h, _ = od.get_location()
x1, y1 = max(0, int(x)), max(0, int(y))
x2, y2 = min(width, int(x + w)), min(height, int(y + h))
cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
cur_centroids.append((cx, cy))
# Speed = displacement from the nearest detection last frame.
speed = 0.0
if prev_centroids:
d = min(np.hypot(cx - px, cy - py) for px, py in prev_centroids)
if d <= MATCH_DIST:
speed = d
speed_sum[y1:y2, x1:x2] += speed
count[y1:y2, x1:x2] += 1.0
det_count += 1
idx += 1
prev_centroids = cur_centroids
total_dets += det_count
overlay = render_speed_heatmap(frame, speed_sum, count, HEATMAP_ALPHA)
if writer is None:
writer = cv2.VideoWriter(
"output_dlstreamer.mp4", cv2.VideoWriter_fourcc(*"mp4v"),
30.0, (width, height))
writer.write(overlay)
print(f"Frame {frame_idx}: detections={det_count}", flush=True)
pipeline.set_state(Gst.State.NULL)
if writer:
writer.release()
print(f"Processed {frame_idx} frames, {total_dets} total detections", flush=True)
|