gps-denied-onboard/src/gps_denied_onboard/helpers/feature_extractor.py

"""`FeatureExtractor` — shared image → :class:`KeypointSet` helper (AZ-343 scope expansion).

L1 helper analogous to :mod:`gps_denied_onboard.helpers.lightglue_runtime`
and :mod:`gps_denied_onboard.helpers.ransac_filter`. Produces a
:class:`gps_denied_onboard._types.matching.KeypointSet` (the same
DTO that :class:`LightGlueRuntime.match` consumes) from a raw BGR
image.

Why a shared helper:

- C2.5 :class:`InlierCountReRanker` (AZ-343) consumes one
  :class:`FeatureExtractor` instance to extract features from each
  per-frame nav-camera image AND from each candidate tile's JPEG
  bytes. The same instance MUST produce comparable feature sets
  for both inputs — otherwise the LightGlue inlier count would
  collapse to noise.
- A future C3 backbone that wants to share keypoints with C2.5
  (rather than re-extracting them) can read the same handle from
  the composition root, mirroring the
  :class:`LightGlueRuntime` ownership pattern (R14 fix).

Concrete impls:

- :class:`OpenCvOrbExtractor`: CPU, deterministic, placeholder used
  by tests and by the airborne binary until the C7
  :class:`InferenceRuntime`-backed DISK / ALIKED extractor lands.
  ORB returns binary (``uint8``) descriptors of 32 bytes; we
  convert to ``float32`` per the
  :class:`gps_denied_onboard._types.matching.KeypointSet` contract.
- Future: TensorRT-backed DISK / ALIKED extractor; consumes
  :class:`InferenceRuntime` from C7.

This helper is intentionally L1 — it imports only ``numpy`` and
``cv2`` plus the L1 :class:`KeypointSet` DTO. Concrete strategies
that need GPU backbones live in their own modules and accept the
:class:`InferenceRuntime` via constructor injection.
"""

from __future__ import annotations

from typing import Protocol, runtime_checkable

import cv2
import numpy as np

from gps_denied_onboard._types.matching import KeypointSet

__all__ = [
    "FeatureExtractor",
    "FeatureExtractorError",
    "OpenCvOrbExtractor",
]


# ORB descriptors are 32 bytes (256 bits). LightGlue's KeypointSet
# requires float32 descriptors so we widen ORB's uint8 output. This is
# a placeholder choice; production will swap in DISK/ALIKED (128-d
# float32) via the C7 InferenceRuntime path.
_ORB_DESCRIPTOR_BYTES = 32
_ORB_FLOAT_DESCRIPTOR_DIM = _ORB_DESCRIPTOR_BYTES * 8  # 256-d float32


class FeatureExtractorError(RuntimeError):
    """Raised on extractor construction or per-image failure."""


@runtime_checkable
class FeatureExtractor(Protocol):
    """Image → :class:`KeypointSet` Protocol.

    Implementations are constructor-injected by the composition root
    and shared across consumers (e.g., C2.5 :class:`InlierCountReRanker`
    uses one instance for both query frames and tile pixels).

    Invariants:

    - :meth:`extract` returns a :class:`KeypointSet` whose
      ``descriptors.shape[1] == self.descriptor_dim()``.
    - ``keypoints`` is shape ``(N, 2)`` ``float32`` pixel coordinates.
    - ``descriptors`` is shape ``(N, descriptor_dim)`` ``float32``.
    - Empty inputs (zero keypoints detected) return an empty-but-shaped
      :class:`KeypointSet` (``N == 0``) rather than raising — the
      C2.5 strategy treats zero-feature candidates as drop events.
    - Deterministic for fixed inputs (no internal RNG state).
    """

    def extract(self, image_bgr: np.ndarray) -> KeypointSet:
        """Detect keypoints + compute descriptors on a single image."""
        ...

    def descriptor_dim(self) -> int:
        """Return the dim of every descriptor row produced by :meth:`extract`."""
        ...


class OpenCvOrbExtractor:
    """CPU :class:`FeatureExtractor` backed by ``cv2.ORB_create``.

    Placeholder implementation: ORB is fast (~5 ms / 480p image on a
    modern CPU) and stable enough to exercise the C2.5 strategy's
    orchestration logic, but its uint8 binary descriptors are NOT a
    drop-in for LightGlue-trained DISK/ALIKED features. Production
    deployments MUST replace this extractor with a deep-learning
    backbone before flight (tracked under the future C2.5
    backbone-extractor task).

    The ``nfeatures`` constructor arg caps the number of keypoints
    per image; default 1024 mirrors typical DISK / ALIKED budgets.
    """

    def __init__(self, *, nfeatures: int = 1024) -> None:
        if nfeatures < 1:
            raise FeatureExtractorError(
                f"OpenCvOrbExtractor.nfeatures must be >= 1; got {nfeatures}"
            )
        self._nfeatures: int = nfeatures
        # ORB itself is created lazily so test environments without
        # a working OpenCV install can still import this module.
        # Cached on first call to amortise the per-image cost.
        self._orb: cv2.ORB | None = None

    def descriptor_dim(self) -> int:
        return _ORB_FLOAT_DESCRIPTOR_DIM

    def _get_orb(self) -> cv2.ORB:
        if self._orb is None:
            self._orb = cv2.ORB_create(nfeatures=self._nfeatures)
        return self._orb

    def extract(self, image_bgr: np.ndarray) -> KeypointSet:
        if image_bgr.ndim == 3:
            gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
        elif image_bgr.ndim == 2:
            gray = image_bgr
        else:
            raise FeatureExtractorError(
                "image_bgr must be 2-D (gray) or 3-D (BGR); "
                f"got ndim={image_bgr.ndim} shape={image_bgr.shape}"
            )
        if gray.dtype != np.uint8:
            gray = gray.astype(np.uint8)
        try:
            keypoints_cv, descriptors_uint8 = self._get_orb().detectAndCompute(
                gray, mask=None
            )
        except cv2.error as exc:
            raise FeatureExtractorError(f"cv2.ORB.detectAndCompute failed: {exc}") from exc
        if descriptors_uint8 is None or len(keypoints_cv) == 0:
            keypoints = np.zeros((0, 2), dtype=np.float32)
            descriptors = np.zeros((0, _ORB_FLOAT_DESCRIPTOR_DIM), dtype=np.float32)
            return KeypointSet(keypoints=keypoints, descriptors=descriptors)
        keypoints = np.array(
            [(kp.pt[0], kp.pt[1]) for kp in keypoints_cv], dtype=np.float32
        )
        # Expand each 32-byte ORB descriptor to a 256-d float32 vector
        # of bit indicators (0/1). Matches the contract that
        # ``KeypointSet.descriptors`` is float32.
        bits = np.unpackbits(descriptors_uint8, axis=1).astype(np.float32)
        return KeypointSet(keypoints=keypoints, descriptors=bits)