File size: 4,025 Bytes
aac350d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
23d337e
aac350d
 
 
 
 
 
 
 
23d337e
 
 
aac350d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
23d337e
aac350d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
109
110
111
112
113
114
"""
Recognition service — runs recognition jobs against the reference
gallery.
"""

from __future__ import annotations

import time
from typing import Optional

from confidence.engine import ConfidenceEngine
from confidence.conflicts import ConflictDetector
from metrics.collector import MetricsCollector
from models.jobs import JobRequest
from models.providers import ProviderCapability
from normalization.merger import ReportMerger
from orchestrator.runner import Orchestrator
from pipeline import (
    InputValidator,
    ImagePreprocessor,
    ImageHasher,
    FeatureExtractor,
)
from storage.reference_store import ReferenceStore
from utils.logging import execution_context, new_execution_id


class RecognitionService:
    """Handles recognition-only jobs."""

    def __init__(
        self,
        orchestrator: Orchestrator,
        metrics: MetricsCollector,
        validator: InputValidator,
        preprocessor: ImagePreprocessor,
        hasher: ImageHasher,
        feature_extractor: FeatureExtractor,
        reference_store: ReferenceStore,
        confidence_engine: ConfidenceEngine,
        conflict_detector: ConflictDetector,
    ) -> None:
        self._orchestrator = orchestrator
        self._metrics = metrics
        self._validator = validator
        self._preprocessor = preprocessor
        self._hasher = hasher
        self._feature_extractor = feature_extractor
        self._gallery = reference_store
        self._merger = ReportMerger(confidence_engine, conflict_detector)

    async def recognize(self, request: JobRequest) -> dict:
        eid = new_execution_id()
        with execution_context(execution_id=eid, provider_id="recognition_service"):
            t0 = time.perf_counter()

            vr = self._validator.validate(
                image_url=request.image_url,
                image_base64=request.image_base64,
            )
            if not vr.valid:
                return {"success": False, "error": vr.error, "error_type": "ValidationError"}

            if vr.source == "url":
                pre = self._preprocessor.from_url(request.image_url)
            else:
                pre = self._preprocessor.from_bytes(vr.image_bytes, vr.source)

            img_hash = self._hasher.hash(pre.image)
            pipeline_output = self._feature_extractor.extract(
                pre.image, img_hash, pre.width, pre.height, pre.source,
                original_bytes=pre.original_bytes,
                original_format=pre.original_format,
            )

            # Attach gallery to pipeline output so recognition providers can read it
            pipeline_output.gallery = self._gallery.get_all()

            results = await self._orchestrator.run(
                pipeline_output=pipeline_output,
                capabilities=[ProviderCapability.RECOGNITION],
                provider_whitelist=request.providers or None,
                execution_id=eid,
            )

            elapsed = (time.perf_counter() - t0) * 1000.0
            report = self._merger.merge(
                results=results,
                image_hash=img_hash,
                job_id=eid,
                total_elapsed_ms=elapsed,
                kind="recognition",
            )
            self._metrics.timings.record("job.recognition", elapsed)
            self._metrics.counters.inc("jobs.recognition.completed")
            return {
                "success": True,
                "report": report.model_dump(),
                "elapsed_ms": round(elapsed, 3),
            }

    def list_known_persons(self) -> list[dict]:
        return self._gallery.list_persons()

    def add_known_person(self, name: str, embedding_bytes: bytes) -> dict:
        import numpy as np
        emb = np.frombuffer(embedding_bytes, dtype=np.float32)
        self._gallery.add(name, emb)
        return {"name": name, "added": True}

    def remove_known_person(self, name: str) -> dict:
        removed = self._gallery.remove(name)
        return {"name": name, "removed": removed}