Source code for zoloto.cameras.file

from __future__ import annotations

from pathlib import Path
from typing import Generator

from cv2 import VideoCapture, imread
from numpy.typing import NDArray

from zoloto.exceptions import CameraReadError
from zoloto.marker_type import MarkerType

from .base import BaseCamera
from .mixins import IterableCameraMixin, VideoCaptureMixin, ViewableCameraMixin
from .utils import (
    get_video_capture_resolution,
    validate_calibrated_video_capture_resolution,
)


[docs]class ImageFileCamera(BaseCamera):
[docs] def __init__( self, image_path: Path, *, marker_size: int | None = None, marker_type: MarkerType, calibration_file: Path | None = None, ) -> None: self.image_path = image_path super().__init__( marker_size=marker_size, marker_type=marker_type, calibration_file=calibration_file, )
def __repr__(self) -> str: return f"<{self.__class__.__name__}: {self.image_path}>"
[docs] def capture_frame(self) -> NDArray: return imread(str(self.image_path))
[docs]class VideoFileCamera( VideoCaptureMixin, IterableCameraMixin, BaseCamera, ViewableCameraMixin ):
[docs] def __init__( self, video_path: Path, *, marker_size: int | None = None, marker_type: MarkerType, calibration_file: Path | None = None, ) -> None: super().__init__( marker_size=marker_size, marker_type=marker_type, calibration_file=calibration_file, ) self.video_path = video_path self.video_capture = VideoCapture(str(self.video_path)) if self.calibration_params is not None: validate_calibrated_video_capture_resolution( self.video_capture, self.calibration_params, override=False )
def __repr__(self) -> str: return f"<{self.__class__.__name__}: {self.video_path}>"
[docs] def close(self) -> None: super().close() self.video_capture.release()
[docs] def get_resolution(self) -> tuple[int, int]: return get_video_capture_resolution(self.video_capture)
[docs] def __iter__(self) -> Generator[NDArray, None, None]: try: yield from super().__iter__() except CameraReadError as e: if e.frame is not None: raise