Source code for ovkit.recognize.base

"""Base class shared by task adapters.

An *adapter* turns a generic :class:`~ovkit.core.backend.Backend` plus an input
image into a task-specific :class:`~ovkit.core.results.Results`. The adapter owns
the pre/post-processing; the backend only moves tensors.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np

from ..core.backend import Backend
from ..core.errors import multi_input_message
from ..core.results import Results




class BaseAdapter:
    """Common configuration + interface for every task adapter."""

    task: str = "base"

    def __init__(
        self,
        *,
        imgsz: int = 640,
        preprocess: dict[str, Any] | None = None,
        postprocess: dict[str, Any] | None = None,
        names: dict[int, str] | None = None,
    ) -> None:
        self.imgsz = imgsz
        self.pre = preprocess or {}
        self.post = postprocess or {}
        self.names = names or {}



    def run(self, backend: Backend, image: np.ndarray, **kwargs: Any) -> Results:
        """Run the full pre/infer/post pipeline for a single image."""
        raise NotImplementedError


    # -- shared preprocessing helpers --------------------------------------

    def _scale_mean_std(self) -> tuple[float, np.ndarray, np.ndarray]:
        scale = float(self.pre.get("scale", 255.0))
        mean = np.asarray(self.pre.get("mean", [0.0, 0.0, 0.0]), dtype=np.float32)
        std = np.asarray(self.pre.get("std", [1.0, 1.0, 1.0]), dtype=np.float32)
        return scale, mean, std



    def preprocess_square(self, image: np.ndarray, rgb: bool = True) -> np.ndarray:
        """Resize to ``imgsz`` square, normalize, and return an NCHW float tensor."""
        return self.preprocess(image, (self.imgsz, self.imgsz), rgb=rgb)




    @staticmethod
    def labels_beside(model_path: str | None) -> dict[int, str] | None:
        """Class names from a ``labels.txt`` next to the model, if present.

        Keeping labels as data (one name per line, mirrored with the model)
        avoids hardcoding a thousand ImageNet strings in the package and works
        for any dataset a model was trained on.
        """
        if not model_path:
            return None
        path = Path(model_path).parent / "labels.txt"
        if not path.is_file():
            return None
        try:
            lines = [ln.strip() for ln in path.read_text(encoding="utf-8").splitlines()]
        except OSError:
            return None
        names = [ln for ln in lines if ln]
        return {i: n for i, n in enumerate(names)} if names else None




    def model_input_hw(self, backend: Backend) -> tuple[int, int]:
        """Return the model's static spatial ``(h, w)`` (last two dims), or ``imgsz``.

        Works for 4-D ``[N,C,H,W]`` and higher-rank inputs (e.g. video clips
        ``[N,C,T,H,W]``); only the trailing two dims are taken as ``H, W``.
        Raises a clear error for multi-input models, which a single image
        cannot drive automatically.
        """
        if len(backend.inputs) > 1:
            names = []
            for inp in backend.inputs:
                try:
                    names.append(inp.get_any_name())
                except RuntimeError:
                    names.append("?")
            raise ValueError(multi_input_message("This model", names))
        shape = backend.input_shape  # full shape, -1 for dynamic dims
        # NHWC (TF-converted OMZ models): channels last -> spatial dims are 1, 2.
        if len(shape) == 4 and shape[-1] in (1, 3) and shape[1] not in (1, 3):
            if shape[1] > 0 and shape[2] > 0:
                return int(shape[1]), int(shape[2])
        if len(shape) >= 2 and shape[-1] and shape[-1] > 0 and shape[-2] and shape[-2] > 0:
            return int(shape[-2]), int(shape[-1])
        return self.imgsz, self.imgsz




    def preprocess(
        self,
        image: np.ndarray,
        size: tuple[int, int],
        rgb: bool = True,
        scale: float | None = None,
    ) -> np.ndarray:
        """Resize to ``(h, w)``, normalize, and return an NCHW float tensor.

        ``scale`` overrides the manifest/default divisor (e.g. ``1.0`` to keep
        raw ``[0, 255]`` input, ``255.0`` to map to ``[0, 1]``).
        """
        from ..image import ops

        h, w = size
        img = ops.resize(image, (w, h))
        if rgb:
            img = ops.bgr_to_rgb(img)
        if img.ndim == 2:  # a grayscale crop
            img = img[:, :, None]

        # One pass from uint8 HWC to float32 NCHW. The transpose is a view, so
        # ``ascontiguousarray`` does the cast and the reorder in a single copy.
        # Written out as five statements this allocated a full-size float32
        # array five times a frame (cast, divide, subtract, divide, contiguous)
        # and preprocessing came to two thirds of the time a frame cost.
        arr = np.ascontiguousarray(img.transpose(2, 0, 1)[None], dtype=np.float32)

        default_scale, mean, std = self._scale_mean_std()
        scale = default_scale if scale is None else float(scale)
        if scale != 1.0:
            arr /= scale
        # ``arr`` is NCHW now, so the per-channel vectors broadcast on axis 1.
        if np.any(mean != 0.0):
            arr -= mean.reshape(1, -1, 1, 1)
        if np.any(std != 1.0):
            arr /= std.reshape(1, -1, 1, 1)
        return arr