feat(perception): add frozen yolox provider

src/k1link/perception/detector.py

@@ -0,0 +1,156 @@
+"""Frozen raw-KB4 YOLOX provider for class-agnostic object proposals."""
+
+from __future__ import annotations
+
+import time
+from collections import Counter
+from collections.abc import Callable
+from dataclasses import dataclass
+from threading import Lock
+from typing import Final
+
+import numpy as np
+from numpy.typing import NDArray
+
+from k1link.compute.yolox_object_detector import (
+    FROZEN_YOLOX_CONFIG,
+    YOLOX_MODEL_ID,
+    YOLOX_MODEL_VERSION,
+    FrozenYoloxConfig,
+    ImageResizer,
+    InferenceBackend,
+    YoloxDetection,
+    postprocess_yolox,
+    preprocess_raw_kb4,
+)
+
+from .contracts import BoundingRegion2D, ObjectProposal2D
+from .providers import SourcePacket
+
+FROZEN_YOLOX_PROVIDER_ID: Final = "triton-yolox-s-raw-kb4/v1"
+FROZEN_YOLOX_MODEL_ID: Final = f"{YOLOX_MODEL_ID}:{YOLOX_MODEL_VERSION}"
+FROZEN_YOLOX_PREPROCESS_ID: Final = "raw-kb4-valid-fov-letterbox/v1"
+
+
+class DetectorProviderError(RuntimeError):
+    """The detector input, frozen inference or proposal output is incompatible."""
+
+
+@dataclass(frozen=True, slots=True)
+class DetectorProviderSnapshot:
+    input_frames: int
+    completed_frames: int
+    failed_frames: int
+    zero_proposal_frames: int
+    proposal_count: int
+    rejected: tuple[tuple[str, int], ...]
+    core_duration_ns: int
+
+
+class FrozenYoloxDetectorProvider:
+    """One image payload produces one frozen inference request and proposal tuple."""
+
+    provider_id: str = FROZEN_YOLOX_PROVIDER_ID
+
+    def __init__(
+        self,
+        *,
+        mask: NDArray[np.bool_],
+        backend: InferenceBackend,
+        resizer: ImageResizer | None = None,
+        config: FrozenYoloxConfig = FROZEN_YOLOX_CONFIG,
+        clock_ns: Callable[[], int] = time.perf_counter_ns,
+    ) -> None:
+        if mask.shape != (600, 800) or mask.dtype != np.bool_ or not np.any(mask):
+            raise DetectorProviderError("frozen valid-FOV mask is incompatible")
+        self.mask = np.asarray(mask, dtype=np.bool_)
+        self.backend = backend
+        self.resizer = resizer
+        self.config = config
+        self._clock_ns = clock_ns
+        self._lock = Lock()
+        self._input_frames = 0
+        self._completed_frames = 0
+        self._failed_frames = 0
+        self._zero_proposal_frames = 0
+        self._proposal_count = 0
+        self._rejected: Counter[str] = Counter()
+        self._core_duration_ns = 0
+
+    def detect(self, packet: SourcePacket) -> tuple[ObjectProposal2D, ...]:
+        payload = packet.image_payload
+        with self._lock:
+            self._input_frames += 1
+        started_ns = int(self._clock_ns())
+        try:
+            if not isinstance(payload, np.ndarray):
+                raise DetectorProviderError("detector requires a decoded BGR image payload")
+            image = np.asarray(payload)
+            if image.dtype != np.uint8:
+                raise DetectorProviderError("decoded BGR image must be uint8")
+            tensor = preprocess_raw_kb4(
+                image,
+                self.mask,
+                config=self.config,
+                resizer=self.resizer,
+            )
+            output = self.backend.infer(tensor)
+            postprocessed = postprocess_yolox(output, self.mask, config=self.config)
+            proposals = proposals_from_detections(packet, postprocessed.detections)
+        except Exception:
+            with self._lock:
+                self._failed_frames += 1
+                self._core_duration_ns += max(0, int(self._clock_ns()) - started_ns)
+            raise
+        with self._lock:
+            self._completed_frames += 1
+            self._proposal_count += len(proposals)
+            self._zero_proposal_frames += not proposals
+            self._rejected.update(dict(postprocessed.rejected))
+            self._core_duration_ns += max(0, int(self._clock_ns()) - started_ns)
+        return proposals
+
+    def snapshot(self) -> DetectorProviderSnapshot:
+        with self._lock:
+            return DetectorProviderSnapshot(
+                input_frames=self._input_frames,
+                completed_frames=self._completed_frames,
+                failed_frames=self._failed_frames,
+                zero_proposal_frames=self._zero_proposal_frames,
+                proposal_count=self._proposal_count,
+                rejected=tuple(sorted(self._rejected.items())),
+                core_duration_ns=self._core_duration_ns,
+            )
+
+
+def proposals_from_detections(
+    packet: SourcePacket,
+    detections: tuple[YoloxDetection, ...],
+) -> tuple[ObjectProposal2D, ...]:
+    envelope = packet.envelope
+    return tuple(
+        ObjectProposal2D(
+            proposal_id=f"proposal-{envelope.sequence}-{index}",
+            source_id=envelope.source_id,
+            frame_id=envelope.frame_id,
+            region=BoundingRegion2D(*detection.bbox_xyxy),
+            objectness=detection.score,
+            provider_id=FROZEN_YOLOX_PROVIDER_ID,
+            model_id=FROZEN_YOLOX_MODEL_ID,
+            preprocess_id=FROZEN_YOLOX_PREPROCESS_ID,
+            semantic_hint=detection.label,
+            provider_tracklet=None,
+        )
+        for index, detection in enumerate(detections)
+    )
+
+
+__all__ = [
+    "FROZEN_YOLOX_MODEL_ID",
+    "FROZEN_YOLOX_PREPROCESS_ID",
+    "FROZEN_YOLOX_PROVIDER_ID",
+    "DetectorProviderError",
+    "DetectorProviderSnapshot",
+    "FrozenYoloxDetectorProvider",
+    "proposals_from_detections",
+]