File size: 3,974 Bytes
23d337e
 
 
892fa81
 
23d337e
 
 
 
 
 
 
 
 
 
892fa81
23d337e
892fa81
23d337e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
892fa81
23d337e
 
 
 
 
 
 
 
 
 
 
 
 
892fa81
 
 
 
23d337e
 
 
 
 
892fa81
23d337e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
"""
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  # CPU fallback silently
        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)