File size: 6,434 Bytes
31de45b | 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | import cv2
import numpy as np
from PIL import Image
from ov_models import OVModel
DET_CONF = 0.6
SAME_PERSON = 0.55 # cosine sim above this = same identity (tune 0.4-0.6)
SAMPLE_FPS = 2.0 # detection sampling rate (identity discovery)
# ArcFace-style reference 5 points, scaled to a 128x128 aligned crop.
_REF_5PTS = np.array([
[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366],
[41.5493, 92.3655], [70.7299, 92.2041],
], dtype=np.float32) * (128.0 / 112.0)
_det = _lm = _reid = None
def _load():
global _det, _lm, _reid
if _det is None:
_det = OVModel("face-detection-retail-0004")
_lm = OVModel("landmarks-regression-retail-0009")
_reid = OVModel("face-reidentification-retail-0095")
return _det, _lm, _reid
def _detect(frame_bgr) -> list[list[int]]:
det, _, _ = _load()
h, w = frame_bgr.shape[:2]
out = det.infer(frame_bgr).reshape(-1, 7) # [_,_,conf,x1,y1,x2,y2] norm
boxes = []
for _, _, conf, x1, y1, x2, y2 in out:
if conf < DET_CONF:
continue
boxes.append([max(0, int(x1 * w)), max(0, int(y1 * h)),
min(w, int(x2 * w)), min(h, int(y2 * h))])
return boxes
def _aligned_embed(frame_bgr, box) -> np.ndarray | None:
"""Crop -> landmark-align to 128x128 -> re-id embedding (L2-normalized).
Alignment matters: re-id-0095 expects an aligned face; raw crops cluster
badly on tilted/profile shots, which would split or merge people.
"""
_, lm, reid = _load()
x1, y1, x2, y2 = box
crop = frame_bgr[y1:y2, x1:x2]
if crop.size == 0:
return None
# landmarks-regression-retail-0009 -> 10 values = 5 (x,y) in crop-relative
pts = lm.infer(crop).reshape(5, 2)
pts[:, 0] *= (x2 - x1)
pts[:, 1] *= (y2 - y1)
pts[:, 0] += x1
pts[:, 1] += y1
M, _ = cv2.estimateAffinePartial2D(pts.astype(np.float32), _REF_5PTS)
if M is None:
aligned = cv2.resize(crop, (128, 128)) # fallback: no alignment
else:
aligned = cv2.warpAffine(frame_bgr, M, (128, 128))
v = reid.infer(aligned).reshape(-1)
return v / (np.linalg.norm(v) + 1e-9)
def scan_video(video_path: str, max_seconds: float = 600.0):
"""Sparse pass. Returns:
embeds: list[np.ndarray] # one per detected face
crops: list[PIL.Image] # matching crops
motion: {frame_idx: float} # for key-event picking
meta: {"fps", "n_frames"}
"""
_load()
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
stride = max(1, int(round(fps / SAMPLE_FPS)))
embeds, crops, motion = [], [], {}
prev_small = None
idx = 0
while True:
ok, frame = cap.read()
if not ok or idx / fps > max_seconds:
break
if idx % stride == 0:
# motion vs previous SAMPLED frame (coarse but fine for events)
small = cv2.cvtColor(cv2.resize(frame, (160, 90)), cv2.COLOR_BGR2GRAY)
if prev_small is not None:
motion[idx] = float(np.mean(cv2.absdiff(small, prev_small)))
prev_small = small
for box in _detect(frame):
emb = _aligned_embed(frame, box)
if emb is None:
continue
embeds.append(emb)
x1, y1, x2, y2 = box
crops.append(Image.fromarray(
cv2.cvtColor(frame[y1:y2, x1:x2], cv2.COLOR_BGR2RGB)))
idx += 1
n_frames = idx
cap.release()
return embeds, crops, motion, {"fps": fps, "n_frames": n_frames}
def cluster_faces(embeds, crops):
"""Greedy-cluster all detected faces into unique identities. Returns:
identities: [{"centroid": emb, "crop": best PIL, "count": n}, ...]
Index in this list == identity id used everywhere else.
"""
identities = []
for emb, crop in zip(embeds, crops):
best, best_sim = None, SAME_PERSON
for ident in identities:
sim = float(np.dot(emb, ident["centroid"]))
if sim >= best_sim:
best, best_sim = ident, sim
if best is None:
identities.append({"centroid": emb, "crop": crop,
"crop_area": crop.size[0] * crop.size[1],
"count": 1})
else:
n = best["count"]
best["centroid"] = (best["centroid"] * n + emb) / (n + 1)
best["centroid"] /= (np.linalg.norm(best["centroid"]) + 1e-9)
best["count"] = n + 1
if crop.size[0] * crop.size[1] > best["crop_area"]:
best["crop"] = crop
best["crop_area"] = crop.size[0] * crop.size[1]
# drop singletons (likely false detections) unless that leaves nothing
strong = [i for i in identities if i["count"] >= 2]
return strong or identities
def annotate_event_frame(frame_bgr, identities, names) -> Image.Image:
"""For ONE key-event frame: detect faces, match each to a named identity,
burn the name above the box. This image is VLM-input only — the user never
sees it. This is how the VLM grounds names to the right person.
"""
img = frame_bgr.copy()
centroids = [i["centroid"] for i in identities]
for box in _detect(frame_bgr):
emb = _aligned_embed(frame_bgr, box)
if emb is None or not centroids:
continue
sims = [float(np.dot(emb, c)) for c in centroids]
iid = int(np.argmax(sims))
if sims[iid] < SAME_PERSON:
continue
name = names.get(iid, "").strip()
if not name:
continue
x1, y1, x2, y2 = box
scale = max(0.5, (x2 - x1) / 200.0)
thick = max(1, int(scale * 2))
(tw, th), _ = cv2.getTextSize(name, cv2.FONT_HERSHEY_DUPLEX, scale, thick)
tx, ty = x1, max(th + 6, y1 - 8)
cv2.rectangle(img, (tx - 4, ty - th - 6), (tx + tw + 4, ty + 4),
(0, 0, 0), -1)
cv2.putText(img, name, (tx, ty), cv2.FONT_HERSHEY_DUPLEX,
scale, (0, 255, 255), thick, cv2.LINE_AA)
return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
def frame_at(video_path: str, frame_idx: int):
cap = cv2.VideoCapture(video_path)
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
ok, frame = cap.read()
cap.release()
return frame if ok else None
|