| |
| |
| |
| |
| |
|
|
| 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/face-detection-adas-0001/FP16/face-detection-adas-0001.xml" |
| REID_MODEL = "intel/face-reidentification-retail-0095/FP16/face-reidentification-retail-0095.xml" |
| |
| DEVICE = "GPU" |
| DET_THRESHOLD = 0.6 |
| MATCH_THRESHOLD = 0.5 |
|
|
| |
| |
| gallery = [] |
| next_id = 1 |
|
|
|
|
| def recognize(embedding): |
| global next_id |
| best_index, best_sim = -1, 0.0 |
| for index, (_, gallery_emb) in enumerate(gallery): |
| sim = float(np.dot(embedding, gallery_emb)) |
| if sim > best_sim: |
| best_sim, best_index = sim, index |
| if best_sim >= MATCH_THRESHOLD: |
| person_id, gallery_emb = gallery[best_index] |
| |
| updated = 0.9 * gallery_emb + 0.1 * embedding |
| gallery[best_index] = (person_id, updated / np.linalg.norm(updated)) |
| return person_id |
| person_id = next_id |
| next_id += 1 |
| gallery.append((person_id, embedding)) |
| print(f"Enrolled ID {person_id}", flush=True) |
| return person_id |
|
|
|
|
| def face_embeddings(video_frame): |
| """Yield ((x, y, w, h), normalized_embedding) for each classified face.""" |
| for region in video_frame.regions(): |
| rect = region.rect() |
| emb = None |
| for tensor in region.tensors(): |
| if tensor.is_detection(): |
| continue |
| data = np.array(tensor.data(), dtype=np.float32) |
| if data.size >= 256: |
| emb = data[:256] |
| if emb is None: |
| continue |
| emb = emb / (np.linalg.norm(emb) + 1e-9) |
| yield (int(rect.x), int(rect.y), int(rect.w), int(rect.h)), emb |
|
|
|
|
| |
| |
| pipeline = Gst.parse_launch( |
| f"filesrc location={INPUT_VIDEO} ! decodebin3 ! " |
| "videoconvert ! video/x-raw,format=BGR ! " |
| f"gvadetect model={DETECTION_MODEL} device={DEVICE} " |
| f"threshold={DET_THRESHOLD} ! queue ! " |
| f"gvaclassify model={REID_MODEL} device={DEVICE} ! queue ! " |
| "appsink name=sink emit-signals=true sync=false max-buffers=4 drop=false" |
| ) |
| sink = pipeline.get_by_name("sink") |
|
|
| writer = {"w": None} |
|
|
|
|
| def on_video(sink): |
| sample = sink.emit("pull-sample") |
| if sample is None: |
| return Gst.FlowReturn.OK |
| vf = VideoFrame(sample.get_buffer(), caps=sample.get_caps()) |
| labeled = [] |
| for (x, y, w, h), emb in face_embeddings(vf): |
| labeled.append((x, y, w, h, recognize(emb))) |
|
|
| 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, person_id in labeled: |
| cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2) |
| cv2.putText(frame, f"ID {person_id}", (x, max(15, y - 8)), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) |
| writer["w"].write(frame) |
| return Gst.FlowReturn.OK |
|
|
|
|
| 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"Total identities recognized: {len(gallery)}", flush=True) |
| print(f"Saved: {OUTPUT_VIDEO}", flush=True) |
|
|
