File size: 3,177 Bytes
38b9516
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
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