FaceInsight_AI / src /inference /face_detector.py
vaisagan's picture
Upload src/inference/face_detector.py with huggingface_hub
38b9516 verified
"""
Face detection using OpenCV DNN (Res10 SSD).
The model is automatically downloaded to ~/.cache/face_detector/ on first use.
"""
from __future__ import annotations
import os
import urllib.request
from pathlib import Path
from typing import List, Tuple
import cv2
import numpy as np
CACHE_DIR = Path.home() / ".cache" / "face_detector"
PROTO_URL = "https://raw.githubusercontent.com/opencv/opencv/master/samples/dnn/face_detector/deploy.prototxt"
WEIGHTS_URL = "https://github.com/opencv/opencv_3rdparty/raw/dnn_samples_face_detector_20170830/res10_300x300_ssd_iter_140000.caffemodel"
PROTO_PATH = CACHE_DIR / "deploy.prototxt"
WEIGHTS_PATH = CACHE_DIR / "res10_300x300_ssd_iter_140000.caffemodel"
def _download_if_missing(url: str, path: Path) -> None:
if path.exists():
return
path.parent.mkdir(parents=True, exist_ok=True)
print(f"[face_detector] Downloading {path.name}…")
urllib.request.urlretrieve(url, path)
class FaceDetector:
"""Detects faces in an image and returns bounding boxes + crops."""
def __init__(self, confidence_threshold: float = 0.7) -> None:
self.threshold = confidence_threshold
_download_if_missing(PROTO_URL, PROTO_PATH)
_download_if_missing(WEIGHTS_URL, WEIGHTS_PATH)
self.net = cv2.dnn.readNetFromCaffe(
str(PROTO_PATH), str(WEIGHTS_PATH)
)
def detect(
self, image: np.ndarray
) -> List[Tuple[int, int, int, int]]:
"""
Args:
image: BGR numpy array (H, W, 3)
Returns:
List of (x1, y1, x2, y2) boxes, sorted left-to-right
"""
h, w = image.shape[:2]
blob = cv2.dnn.blobFromImage(
cv2.resize(image, (300, 300)), 1.0,
(300, 300), (104.0, 177.0, 123.0)
)
self.net.setInput(blob)
detections = self.net.forward()
boxes: List[Tuple[int, int, int, int]] = []
for i in range(detections.shape[2]):
conf = float(detections[0, 0, i, 2])
if conf < self.threshold:
continue
box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
x1, y1, x2, y2 = box.astype(int)
x1, y1 = max(0, x1), max(0, y1)
x2, y2 = min(w, x2), min(h, y2)
if x2 > x1 and y2 > y1:
boxes.append((x1, y1, x2, y2))
return sorted(boxes, key=lambda b: b[0])
def crop_faces(
self, image: np.ndarray, padding: float = 0.10
) -> Tuple[List[np.ndarray], List[Tuple[int, int, int, int]]]:
"""
Returns (crops, boxes) where crops are RGB face crops with optional padding.
"""
h, w = image.shape[:2]
boxes = self.detect(image)
crops = []
for x1, y1, x2, y2 in boxes:
pad_x = int((x2 - x1) * padding)
pad_y = int((y2 - y1) * padding)
cx1, cy1 = max(0, x1 - pad_x), max(0, y1 - pad_y)
cx2, cy2 = min(w, x2 + pad_x), min(h, y2 + pad_y)
crop = image[cy1:cy2, cx1:cx2]
crop_rgb = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
crops.append(crop_rgb)
return crops, boxes