feat(ai): dual-model face detection with NMS deduplication

Run both short-range and full-range MediaPipe models and merge results,
then apply non-maximum suppression to remove duplicate bounding boxes.
Fixes missed faces in group photos where the single-model loop exited
early after the first positive detection.
This commit is contained in:
Siddharth Kumar Sah
2026-04-15 23:11:36 +08:00
parent ea98e9b5de
commit 3c1bf0a77d
+47 -19
View File
@@ -30,12 +30,48 @@ def _ensure_face_detect_model():
return _LOCAL_MODEL_PATH
def _iou(a, b):
"""Compute intersection-over-union between two face boxes."""
ax2, ay2 = a["x"] + a["w"], a["y"] + a["h"]
bx2, by2 = b["x"] + b["w"], b["y"] + b["h"]
inter_w = max(0, min(ax2, bx2) - max(a["x"], b["x"]))
inter_h = max(0, min(ay2, by2) - max(a["y"], b["y"]))
inter = inter_w * inter_h
union = a["w"] * a["h"] + b["w"] * b["h"] - inter
return inter / union if union > 0 else 0.0
def _nms_faces(faces, iou_threshold=0.4):
"""Remove duplicate detections using greedy non-maximum suppression."""
if len(faces) <= 1:
return faces
kept = []
used = [False] * len(faces)
for i in range(len(faces)):
if used[i]:
continue
kept.append(faces[i])
used[i] = True
for j in range(i + 1, len(faces)):
if not used[j] and _iou(faces[i], faces[j]) >= iou_threshold:
used[j] = True
return kept
def _detect_with_solutions(img_array, min_confidence):
"""Detect faces using legacy mp.solutions API (mediapipe < 0.10.30)."""
"""Detect faces using legacy mp.solutions API (mediapipe < 0.10.30).
Runs both short-range (model 0) and full-range (model 1) detectors and
merges the results. Previously the loop broke on the first model that
found any face, so group photos where model 0 caught only 1-2 large
faces would never have the remaining faces scanned by model 1.
"""
import mediapipe as mp
mp_face = mp.solutions.face_detection
results = None
ih, iw = img_array.shape[:2]
all_faces = []
for model_sel in [0, 1]:
detector = mp_face.FaceDetection(
model_selection=model_sel,
@@ -43,24 +79,16 @@ def _detect_with_solutions(img_array, min_confidence):
)
results = detector.process(img_array)
detector.close()
if results.detections:
break
for detection in (results.detections or []):
bbox = detection.location_data.relative_bounding_box
all_faces.append({
"x": int(bbox.xmin * iw),
"y": int(bbox.ymin * ih),
"w": int(bbox.width * iw),
"h": int(bbox.height * ih),
})
detections = results.detections or []
if not detections:
return []
ih, iw = img_array.shape[:2]
faces = []
for detection in detections:
bbox = detection.location_data.relative_bounding_box
faces.append({
"x": int(bbox.xmin * iw),
"y": int(bbox.ymin * ih),
"w": int(bbox.width * iw),
"h": int(bbox.height * ih),
})
return faces
return _nms_faces(all_faces)
def _detect_with_tasks(img_array, min_confidence):