| |
| |
| |
| |
| |
|
|
| 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 |
|
|
| INPUT_VIDEO = "test_video.mp4" |
| HEATMAP_ALPHA = 0.55 |
| MATCH_DIST = 80.0 |
| MAX_SPEED = 20.0 |
|
|
|
|
| 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) |
| 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) |
|
|
|
|
| |
| MODEL = "yolo26n" |
|
|
| |
| |
| 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 = 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) |
|
|
