mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
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:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user