from pathlib import Path

import cv2
from zoloto.cameras.camera import Camera
from zoloto.marker_type import MarkerType


class AxisAnnotatingCamera(Camera):
    def _annotate_frame(self, frame):
        """
        Annotate axis onto the image too.

        This does perform the detection twice, but for a demo this is fine.
        """
        super()._annotate_frame(frame)
        for marker in self.process_frame_eager(frame=frame):
            cv2.drawFrameAxes(
                frame,
                self.calibration_params.camera_matrix,
                self.calibration_params.distance_coefficients,
                marker._rvec,
                marker._tvec,
                self._marker_size
            )


with AxisAnnotatingCamera(0, marker_type=MarkerType.APRILTAG_36H11, marker_size=100, calibration_file=Path(__file__).parent / "calibrations.xml") as camera:
    print("Starting preview... Press 'q' to exit.")  # noqa: T001
    camera.show(annotate=True)