| |
| |
| |
| |
| |
|
|
| 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" |
| REFERENCE_IMAGE = "face_a.jpg" |
| 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.4 |
|
|
|
|
| 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 |
|
|
|
|
| def run_pipeline(source_desc, on_frame): |
| |
| |
| pipeline = Gst.parse_launch( |
| f"{source_desc} ! 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") |
| sink.connect("new-sample", on_frame) |
| 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) |
|
|
|
|
| |
| ref = {"emb": None, "area": 0} |
|
|
|
|
| def on_reference(sink): |
| sample = sink.emit("pull-sample") |
| if sample is None: |
| return Gst.FlowReturn.OK |
| vf = VideoFrame(sample.get_buffer(), caps=sample.get_caps()) |
| for (x, y, w, h), emb in face_embeddings(vf): |
| if w * h > ref["area"]: |
| ref["area"] = w * h |
| ref["emb"] = emb |
| return Gst.FlowReturn.OK |
|
|
|
|
| run_pipeline(f"filesrc location={REFERENCE_IMAGE} ! jpegdec", on_reference) |
| if ref["emb"] is None: |
| raise SystemExit("No face detected in the reference image") |
| ref_emb = ref["emb"] |
|
|
| |
| 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 = [] |
| for (x, y, w, h), emb in face_embeddings(vf): |
| similarity = float(np.dot(emb, ref_emb)) |
| if similarity >= MATCH_THRESHOLD: |
| matches.append((x, y, w, h, similarity)) |
|
|
| 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, similarity in matches: |
| cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2) |
| cv2.putText(frame, f"MATCH {similarity:.2f}", (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 |
|
|
|
|
| run_pipeline(f"filesrc location={INPUT_VIDEO} ! decodebin3", on_video) |
| if writer["w"] is not None: |
| writer["w"].release() |
| print(f"Frames with a matched face: {match_frames}", flush=True) |
| print(f"Saved: {OUTPUT_VIDEO}", flush=True) |
|
|
