diff --git a/sleap_io/io/video.py b/sleap_io/io/video.py index 25ccc0d0..893c25b0 100644 --- a/sleap_io/io/video.py +++ b/sleap_io/io/video.py @@ -49,11 +49,13 @@ class VideoBackend: If False, will close the reader after each call. If True (the default), it will keep the reader open and cache it for subsequent calls which may enhance the performance of reading multiple frames. + _frame_rate: Frame rate of the video. """ filename: str | Path | list[str] | list[Path] grayscale: Optional[bool] = None keep_open: bool = True + _frame_rate: Optional[float] = None _cached_shape: Optional[Tuple[int, int, int, int]] = None _open_reader: Optional[object] = None @@ -77,6 +79,7 @@ def from_filename( frames. If False, will close the reader after each call. If True (the default), it will keep the reader open and cache it for subsequent calls which may enhance the performance of reading multiple frames. + _frame_rate: Frame rate of the video. Returns: VideoBackend subclass instance. @@ -189,6 +192,18 @@ def frames(self) -> int: """Number of frames in the video.""" return self.shape[0] + @property + def frame_rate(self) -> Optional[float]: + """Frames per second of the video.""" + video_extensions = ["mp4", "avi", "mov", "mj2", "mkv"] + if not any(self.filename.endswith(ext) for ext in video_extensions): + return None + + if "cv2" in sys.modules: + return cv2.VideoCapture(self.filename).get(cv2.CAP_PROP_FPS) + else: + return iio.immeta(self.filename)["fps"] + def __len__(self) -> int: """Return number of frames in the video.""" return self.shape[0] @@ -314,6 +329,7 @@ class MediaVideo(VideoBackend): If False, will close the reader after each call. If True (the default), it will keep the reader open and cache it for subsequent calls which may enhance the performance of reading multiple frames. + _frame_rate: Frame rate of the video. plugin: Video plugin to use. One of "opencv", "FFMPEG", or "pyav". If `None`, will use the first available plugin in the order listed above. """ @@ -469,6 +485,7 @@ class HDF5Video(VideoBackend): If False, will close the reader after each call. If True (the default), it will keep the reader open and cache it for subsequent calls which may enhance the performance of reading multiple frames. + _frame_rate: Frame rate of the video. dataset: Name of dataset to read from. If `None`, will try to find a rank-4 dataset by iterating through datasets in the file. If specifying an embedded dataset, this can be the group containing a "video" dataset or the dataset @@ -710,6 +727,7 @@ class ImageVideo(VideoBackend): """ EXTS = ("png", "jpg", "jpeg", "tif", "tiff", "bmp") + _frame_rate = None @staticmethod def find_images(folder: str) -> list[str]: diff --git a/sleap_io/model/video.py b/sleap_io/model/video.py index edd49489..22b92189 100644 --- a/sleap_io/model/video.py +++ b/sleap_io/model/video.py @@ -62,6 +62,7 @@ def from_filename( dataset: Optional[str] = None, grayscale: Optional[bool] = None, keep_open: bool = True, + _frame_rate: Optional[float] = None, source_video: Optional[Video] = None, **kwargs, ) -> VideoBackend: @@ -79,6 +80,7 @@ def from_filename( frames. If False, will close the reader after each call. If True (the default), it will keep the reader open and cache it for subsequent calls which may enhance the performance of reading multiple frames. + _frame_rate: The frame rate of the video. source_video: The source video object if this is a proxy video. This is present when the video contains an embedded subset of frames from another video. @@ -93,6 +95,7 @@ def from_filename( dataset=dataset, grayscale=grayscale, keep_open=keep_open, + _frame_rate=_frame_rate, **kwargs, ), source_video=source_video, @@ -107,6 +110,15 @@ def shape(self) -> Tuple[int, int, int, int] | None: """ return self._get_shape() + @property + def frame_rate(self) -> float | None: + """Return the frames per second of the video. + + If the video backend is not set or it cannot determine the frames per second of + the video, this will return None. + """ + return self.backend.frame_rate + def _get_shape(self) -> Tuple[int, int, int, int] | None: """Return the shape of the video as (num_frames, height, width, channels). diff --git a/tests/io/test_video_backends.py b/tests/io/test_video_backends.py index 06f22050..f4ecea90 100644 --- a/tests/io/test_video_backends.py +++ b/tests/io/test_video_backends.py @@ -14,6 +14,7 @@ def test_video_backend_from_filename(centered_pair_low_quality_path, slp_minimal assert type(backend) == MediaVideo assert backend.filename == centered_pair_low_quality_path assert backend.shape == (1100, 384, 384, 1) + assert backend.frame_rate == 15.0 backend = VideoBackend.from_filename(slp_minimal_pkg) assert type(backend) == HDF5Video