| """ |
| OpenCV DNN face detector (Caffe SSD). |
| |
| Uses cores.vision for image operations. Model download logic is |
| self-contained; no external service required. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import urllib.request |
|
|
| import cv2 |
| import numpy as np |
|
|
| from config.settings import Settings, settings as _default_settings, MODELS_DIR |
| from cores.vision import BBox |
| from pipeline.feature_extraction import PipelineOutput |
| from providers.base import BaseProvider, ProviderCapability |
|
|
|
|
| class DNNDetector(BaseProvider): |
| name = "dnn" |
| capability = ProviderCapability.DETECTION |
|
|
| PROTOTXT_PATH = MODELS_DIR / "deploy.prototxt" |
| CAFFEMODEL_PATH = MODELS_DIR / "res10_300x300_ssd_iter_140000.caffemodel" |
| PROTOTXT_URL = ( |
| "https://raw.githubusercontent.com/opencv/opencv_3rdparty/" |
| "dnn_samples_face_detector_20170830/deploy.prototxt" |
| ) |
| CAFFEMODEL_URL = ( |
| "https://raw.githubusercontent.com/opencv/opencv_3rdparty/" |
| "dnn_samples_face_detector_20170830/res10_300x300_ssd_iter_140000.caffemodel" |
| ) |
|
|
| def __init__(self, settings: Settings | None = None) -> None: |
| super().__init__(settings=settings or _default_settings) |
| self._net = None |
| self._init_error: str | None = None |
| try: |
| self._ensure_models_downloaded() |
| self._net = cv2.dnn.readNetFromCaffe( |
| str(self.PROTOTXT_PATH), str(self.CAFFEMODEL_PATH) |
| ) |
| try: |
| self._net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA) |
| self._net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA) |
| except Exception: |
| pass |
| except Exception as e: |
| self._init_error = str(e) |
|
|
| def is_available(self) -> bool: |
| return self._net is not None and self._init_error is None |
|
|
| def _run(self, pipeline_output: PipelineOutput) -> tuple[dict, dict]: |
| if self._net is None: |
| raise RuntimeError(f"DNN net not loaded: {self._init_error}") |
|
|
| img: np.ndarray = pipeline_output.image |
| h, w = img.shape[:2] |
| blob = cv2.dnn.blobFromImage( |
| cv2.resize(img, (300, 300)), 1.0, (300, 300), (104.0, 177.0, 123.0), |
| ) |
| self._net.setInput(blob) |
| detections = self._net.forward() |
|
|
| threshold = self._settings.dnn_confidence_threshold |
| boxes_data: list[dict] = [] |
| confidences: list[float] = [] |
| raw_detections: list[dict] = [] |
|
|
| for i in range(detections.shape[2]): |
| confidence = float(detections[0, 0, i, 2]) |
| if confidence < threshold: |
| continue |
| x1 = max(0, min(int(detections[0, 0, i, 3] * w), w - 1)) |
| y1 = max(0, min(int(detections[0, 0, i, 4] * h), h - 1)) |
| x2 = max(0, min(int(detections[0, 0, i, 5] * w), w)) |
| y2 = max(0, min(int(detections[0, 0, i, 6] * h), h)) |
| bw, bh = x2 - x1, y2 - y1 |
| if bw <= 0 or bh <= 0: |
| continue |
| boxes_data.append(BBox(x1, y1, bw, bh).to_dict()) |
| confidences.append(confidence) |
| raw_detections.append({"index": i, "confidence": confidence, "box": [x1, y1, x2, y2]}) |
|
|
| raw = { |
| "model": "res10_300x300_ssd_iter_140000", |
| "threshold": threshold, |
| "detections": raw_detections, |
| "num_faces": len(boxes_data), |
| "image_size": {"width": w, "height": h}, |
| } |
| normalized = { |
| "boxes": boxes_data, |
| "num_faces": len(boxes_data), |
| "confidences": confidences, |
| "landmarks": None, |
| } |
| return raw, normalized |
|
|
| def _ensure_models_downloaded(self) -> None: |
| if not self.PROTOTXT_PATH.exists(): |
| urllib.request.urlretrieve(self.PROTOTXT_URL, self.PROTOTXT_PATH) |
| if not self.CAFFEMODEL_PATH.exists(): |
| urllib.request.urlretrieve(self.CAFFEMODEL_URL, self.CAFFEMODEL_PATH) |
|
|
