Source code for zoloto.marker

from functools import lru_cache
from typing import Any, Dict, List, Optional, Tuple

from cached_property import cached_property
from cv2 import aruco
from numpy import arctan2, array, linalg, ndarray

from .calibration import CalibrationParameters
from .coords import Coordinates, Orientation, Spherical, ThreeDCoordinates
from .exceptions import MissingCalibrationsError


[docs]class Marker:
[docs] def __init__( self, marker_id: int, corners: List[ndarray], size: int, calibration_params: Optional[CalibrationParameters] = None, precalculated_vectors: Optional[Tuple[ndarray, ndarray]] = None, ): self.__id = marker_id self.__pixel_corners = corners self.__size = size self.__camera_calibration_params = calibration_params self.__precalculated_vectors = precalculated_vectors
@property # noqa: A003 def id(self) -> int: return self.__id @property def size(self) -> int: return self.__size def _is_eager(self) -> bool: return self.__precalculated_vectors is not None @property def pixel_corners(self) -> List[Coordinates]: return [Coordinates(*coords) for coords in self.__pixel_corners] @cached_property def pixel_centre(self) -> Coordinates: tl, _, br, _ = self.__pixel_corners return Coordinates( x=tl[0] + (self.__size / 2) - 1, y=br[1] - (self.__size / 2), ) @cached_property def distance(self) -> int: return int(linalg.norm(self._tvec)) @property def orientation(self) -> Orientation: return Orientation(*self._rvec) @cached_property def spherical(self) -> Spherical: x, y, z = self._tvec return Spherical(rot_x=arctan2(y, z), rot_y=arctan2(x, z), dist=self.distance) @property def cartesian(self) -> ThreeDCoordinates: return ThreeDCoordinates(*self._tvec) @lru_cache(maxsize=None) def _get_pose_vectors(self) -> Tuple[ndarray, ndarray]: # Check if the Marker is eager. # We cannot call _is_eager, else mypy will think we are returning Optional if self.__precalculated_vectors is not None: return self.__precalculated_vectors if self.__camera_calibration_params is None: raise MissingCalibrationsError() rvec, tvec, _ = aruco.estimatePoseSingleMarkers( [self.__pixel_corners], self.__size, self.__camera_calibration_params.camera_matrix, self.__camera_calibration_params.distance_coefficients, ) return rvec[0][0], tvec[0][0] @property def _rvec(self) -> ndarray: rvec, _ = self._get_pose_vectors() return rvec @property def _tvec(self) -> ndarray: _, tvec = self._get_pose_vectors() return tvec
[docs] def as_dict(self) -> Dict[str, Any]: marker_dict = { "id": self.id, "size": self.size, "pixel_corners": self.__pixel_corners.tolist(), # type: ignore } try: marker_dict.update({"rvec": list(self._rvec), "tvec": list(self._tvec)}) except MissingCalibrationsError: pass return marker_dict
[docs] @classmethod def from_dict(cls, marker_dict: Dict[str, Any]) -> "Marker": marker_args = [ marker_dict["id"], array(marker_dict["pixel_corners"]), marker_dict["size"], None, ] if "rvec" in marker_dict and "tvec" in marker_dict: marker_args.append((array(marker_dict["rvec"]), array(marker_dict["tvec"]))) return cls(*marker_args)