Source code for olo.ai.types

"""Shared value types for the AI module."""

from dataclasses import dataclass, field

import numpy as np
from olo_protos.ai.v1 import ai_pb2

from olo.vision import BoundingBox




@dataclass(frozen=True, slots=True)
class TensorInfo:
    """Declared model input/output; dynamic dimensions are -1."""

    name: str
    dtype: np.dtype | None
    shape: tuple[int, ...]





@dataclass(frozen=True, slots=True)
class ModelInfo:
    """Metadata for a loaded ONNX model session."""

    model_id: str
    name: str
    inputs: tuple[TensorInfo, ...]
    outputs: tuple[TensorInfo, ...]
    metadata: dict[str, str] = field(default_factory=dict)





@dataclass(frozen=True, slots=True)
class Detection:
    """Standard object-detection result returned by detector adapters.

    Fields map 1:1 onto the olo.ai.v1.Detection wire message. The box is an
    olo.vision.BoundingBox, so detections feed directly into depth helpers
    such as sample_depth_region and deproject_region.
    """

    label: str
    class_id: int
    score: float
    box: BoundingBox



def detection_from_proto(msg: ai_pb2.Detection) -> Detection:
    """Convert an olo.ai.v1.Detection wire message to a Detection."""
    return Detection(
        label=msg.label,
        class_id=msg.class_id,
        score=msg.score,
        box=BoundingBox.from_xyxy(
            msg.box.minimum.x, msg.box.minimum.y, msg.box.maximum.x, msg.box.maximum.y
        ),
    )