Source code for olo.vision.detection

"""Vision wrapper over olo.vision.v1.Vision: appliance-side fiducial tag detection."""

from collections.abc import Iterator, Mapping
from typing import TYPE_CHECKING, Optional

import grpc
from olo_protos.vision.v1 import vision_pb2

from olo.core._payloads import payload_to_proto
from olo.core.images import Image, image_to_message
from olo.core.types import TopicInfo
from olo.errors import wrap_rpc_error
from olo.utils.common import resolve_topic_name
from olo.vision.tags import (
    CharucoBoard,
    DetectorParams,
    TagFrame,
    TagLayout,
    family_to_proto,
    intrinsics_to_proto,
    sizes_to_proto,
    tag_frame_from_proto,
)
from olo.vision.types import CameraIntrinsics

if TYPE_CHECKING:
    from olo.vision.calibration import Calibration

_DEFAULT_DETECT_TIMEOUT_SEC = 30.0
_IMAGE_TYPE = "sensor_msgs/msg/Image"
_IMAGE_TYPES = (_IMAGE_TYPE, "sensor_msgs/msg/CompressedImage")
_CAMERA_INFO_TYPE = "sensor_msgs/msg/CameraInfo"




class TagStream:
    """Context-managed sync iterator over a Vision.SubscribeTags server stream."""



    def __init__(self, stream) -> None:
        """Wrap a SubscribeTags server stream."""
        self._stream = stream
        self._closed = False


    def __iter__(self) -> Iterator[TagFrame]:
        """Iterate incoming tag frames."""
        return self

    def __next__(self) -> TagFrame:
        """Return the next tag frame."""
        if self._closed:
            raise StopIteration
        try:
            response = next(self._stream)
        except StopIteration:
            self._closed = True
            raise
        except grpc.RpcError as exc:
            self._closed = True
            raise wrap_rpc_error(exc) from exc
        return tag_frame_from_proto(response.frame)



    def close(self) -> None:
        """Cancel the stream and mark it closed."""
        if self._closed:
            return
        self._closed = True
        cancel = getattr(self._stream, "cancel", None)
        if callable(cancel):
            cancel()


    def __enter__(self) -> "TagStream":
        """Enter a context manager."""
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        """Close the stream on context exit."""
        self.close()





class TagDetector:
    """A tag family plus physical sizes, bound to the appliance detector.

    Holds configuration only; detection runs on the appliance. ``size_m`` is
    the default edge length and ``sizes`` overrides it per marker id. Without
    a size, tags are reported with corners but no pose. A ``target`` (a
    :class:`TagLayout` or :class:`CharucoBoard`) is additionally solved as one
    pose per frame, reported on :attr:`TagFrame.target_transform`.
    """



    def __init__(
        self,
        vision: "Vision",
        family: str,
        *,
        size_m: float | None = None,
        sizes: Mapping[int, float] | None = None,
        params: DetectorParams | None = None,
        target: TagLayout | CharucoBoard | None = None,
        namespace: str | None = None,
    ) -> None:
        """Validate the family and capture sizes; nothing is sent until a call is made."""
        self._vision = vision
        self.family = family
        self._family_proto = family_to_proto(family)
        self._sizes = sizes_to_proto(size_m, sizes)
        self._params = params.to_proto() if params is not None else None
        self._layout = target.to_proto() if isinstance(target, TagLayout) else None
        self._charuco = target.to_proto() if isinstance(target, CharucoBoard) else None
        if isinstance(target, CharucoBoard) and target.family != family:
            raise ValueError(f"Charuco board family {target.family!r} does not match detector family {family!r}")
        self._namespace = namespace




    def detect(
        self,
        image: Image,
        *,
        intrinsics: CameraIntrinsics | None = None,
        timeout: float = _DEFAULT_DETECT_TIMEOUT_SEC,
    ) -> TagFrame:
        """Detect tags in an image the caller already holds."""
        request = vision_pb2.DetectTagsRequest(
            image=payload_to_proto(_IMAGE_TYPE, image_to_message(image)),
            family=self._family_proto,
        )
        if intrinsics is not None:
            request.intrinsics.CopyFrom(intrinsics_to_proto(intrinsics))
        self._apply_common(request)
        try:
            response = self._vision._stub.DetectTags(request, timeout=timeout)
        except grpc.RpcError as exc:
            raise wrap_rpc_error(exc) from exc
        return tag_frame_from_proto(response.frame)




    def subscribe(
        self,
        image_topic: TopicInfo,
        camera_info_topic: TopicInfo | None = None,
        *,
        timeout: Optional[float] = None,
    ) -> TagStream:
        """Stream detections from a camera topic; poses need ``camera_info_topic`` and a size."""
        if image_topic.msg_type not in _IMAGE_TYPES:
            raise ValueError(f"image_topic must be one of {_IMAGE_TYPES}, got {image_topic.msg_type!r}")
        if camera_info_topic is not None and camera_info_topic.msg_type != _CAMERA_INFO_TYPE:
            raise ValueError(f"camera_info_topic must be {_CAMERA_INFO_TYPE}, got {camera_info_topic.msg_type!r}")
        request = vision_pb2.SubscribeTagsRequest(
            image_topic=resolve_topic_name(self._namespace, image_topic.name),
            image_message_type=image_topic.msg_type,
            camera_info_topic=(
                resolve_topic_name(self._namespace, camera_info_topic.name) if camera_info_topic is not None else ""
            ),
            family=self._family_proto,
        )
        self._apply_common(request)
        try:
            stream = self._vision._stub.SubscribeTags(request, timeout=timeout)
        except grpc.RpcError as exc:
            raise wrap_rpc_error(exc) from exc
        return TagStream(stream)


    def _apply_common(self, request) -> None:
        if self._sizes is not None:
            request.sizes.CopyFrom(self._sizes)
        if self._params is not None:
            request.params.CopyFrom(self._params)
        if self._layout is not None:
            request.layout.CopyFrom(self._layout)
        if self._charuco is not None:
            request.charuco.CopyFrom(self._charuco)





class VisionHandle:
    """Namespace-scoped Vision handle; relative topic names resolve against the robot."""



    def __init__(self, vision: "Vision", namespace: str) -> None:
        """Bind to a vision client and resolved robot namespace."""
        self._vision = vision
        self.namespace = namespace




    def tags(
        self,
        family: str,
        *,
        size_m: float | None = None,
        sizes: Mapping[int, float] | None = None,
        params: DetectorParams | None = None,
        target: TagLayout | CharucoBoard | None = None,
    ) -> TagDetector:
        """Return a detector for one tag family whose topics resolve against this robot."""
        return TagDetector(
            self._vision, family, size_m=size_m, sizes=sizes, params=params, target=target, namespace=self.namespace
        )




    def charuco(self, board: CharucoBoard, *, params: DetectorParams | None = None) -> TagDetector:
        """Return a detector that solves a ChArUco board pose and reports its chessboard corners."""
        return self.tags(board.family, params=params, target=board)




    def detect_tags(
        self,
        image: Image,
        family: str,
        *,
        intrinsics: CameraIntrinsics | None = None,
        size_m: float | None = None,
        sizes: Mapping[int, float] | None = None,
        params: DetectorParams | None = None,
        timeout: float = _DEFAULT_DETECT_TIMEOUT_SEC,
    ) -> TagFrame:
        """Detect tags of one family in an image."""
        return self.tags(family, size_m=size_m, sizes=sizes, params=params).detect(
            image, intrinsics=intrinsics, timeout=timeout
        )




    def subscribe_tags(
        self,
        image_topic: TopicInfo,
        family: str,
        *,
        camera_info_topic: TopicInfo | None = None,
        size_m: float | None = None,
        sizes: Mapping[int, float] | None = None,
        params: DetectorParams | None = None,
        timeout: Optional[float] = None,
    ) -> TagStream:
        """Stream detections of one family from a camera topic on this robot."""
        return self.tags(family, size_m=size_m, sizes=sizes, params=params).subscribe(
            image_topic, camera_info_topic, timeout=timeout
        )






class Vision(VisionHandle):
    """Sync wrapper for the Vision gRPC service (topics are taken as fully qualified)."""



    def __init__(self, session) -> None:
        """Bind to a channel session."""
        super().__init__(self, namespace="")
        self._session = session


    @property
    def _stub(self):
        return self._session.vision_stub

    @property
    def calibration(self) -> "Calibration":
        """Return the calibration solvers (hand-eye, intrinsics)."""
        from olo.vision.calibration import Calibration

        return Calibration(self)



    def handle(self, namespace: Optional[str] = None) -> VisionHandle:
        """Return a namespace-scoped handle; ``None`` resolves the appliance default robot."""
        return VisionHandle(self, self._session.resolve_namespace(namespace))




__all__ = ["TagDetector", "TagStream", "Vision", "VisionHandle"]