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
|
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):
|
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
|
import mediapipe as mp
|
||||||
|
|
||||||
mp_face = mp.solutions.face_detection
|
mp_face = mp.solutions.face_detection
|
||||||
results = None
|
ih, iw = img_array.shape[:2]
|
||||||
|
all_faces = []
|
||||||
|
|
||||||
for model_sel in [0, 1]:
|
for model_sel in [0, 1]:
|
||||||
detector = mp_face.FaceDetection(
|
detector = mp_face.FaceDetection(
|
||||||
model_selection=model_sel,
|
model_selection=model_sel,
|
||||||
@@ -43,24 +79,16 @@ def _detect_with_solutions(img_array, min_confidence):
|
|||||||
)
|
)
|
||||||
results = detector.process(img_array)
|
results = detector.process(img_array)
|
||||||
detector.close()
|
detector.close()
|
||||||
if results.detections:
|
for detection in (results.detections or []):
|
||||||
break
|
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 []
|
return _nms_faces(all_faces)
|
||||||
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
|
|
||||||
|
|
||||||
|
|
||||||
def _detect_with_tasks(img_array, min_confidence):
|
def _detect_with_tasks(img_array, min_confidence):
|
||||||
|
|||||||
Reference in New Issue
Block a user