Source code for

# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import os
import sys
import tempfile
from import Callable
from typing import TYPE_CHECKING, Any

import numpy as np
from import Dataset, IterableDataset

from monai.utils.enums import ColorOrder
from monai.utils.module import optional_import

__all__ = ["VideoDataset", "VideoFileDataset", "CameraDataset"]

    import cv2

    has_cv2 = True
    cv2, has_cv2 = None, None

def import_cv():
    """Import cv2. Put it inside a function to avoid webcam lights blinking on ``import monai``."""
    global cv2
    global has_cv2
    cv2, has_cv2 = optional_import("cv2")

class SuppressStderr:
    """Suppress stderr. Useful as OpenCV (and dependencies) can produce a lot of output."""

    def __enter__(self):
        self.errnull_file = open(os.devnull, "w")
        self.old_stderr_fileno_undup = sys.stderr.fileno()
        self.old_stderr_fileno = os.dup(sys.stderr.fileno())
        self.old_stderr = sys.stderr
        os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup)
        sys.stderr = self.errnull_file
        return self

    def __exit__(self, *_):
        sys.stderr = self.old_stderr
        os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup)

[docs] class VideoDataset: # import inside class to avoid webcam blinking on ``import monai``. import_cv() def __init__( self, video_source: str | int, transform: Callable | None = None, max_num_frames: int | None = None, color_order: str = ColorOrder.RGB, multiprocessing: bool = False, channel_dim: int = 0, ) -> None: """ Base video dataset. Args: video_source: filename of video. transform: transform to be applied to each frame. max_num_frames: Max number of frames to iterate across. If `None` is passed, then the dataset will iterate until the end of the file. color_order: Color order to return frame. Default is RGB. multiprocessing: If `True`, open the video source on the fly. This makes things process-safe, which is useful when combined with a DataLoader with `num_workers>0`. However, when using with `num_workers==0`, it makes sense to use `multiprocessing=False`, as the source will then only be opened once, at construction, which will be faster in those circumstances. channel_dim: OpenCV reads with the channel as the last dimension. Use this flag to move it elsewhere. By default this is zero, so the channel dimension is moved to the front. Raises: RuntimeError: OpenCV not installed. NotImplementedError: Unknown color order. """ if not has_cv2: raise RuntimeError("OpenCV not installed.") if color_order not in ColorOrder: raise NotImplementedError self.color_order = color_order self.channel_dim = channel_dim self.video_source = video_source self.multiprocessing = multiprocessing if not multiprocessing: self.cap = self.open_video(video_source) self.transform = transform self.max_num_frames = max_num_frames @staticmethod def open_video(video_source: str | int): """ Use OpenCV to open a video source from either file or capture device. Args: video_source: filename or index referring to capture device. Raises: RuntimeError: Source is a file but file not found. RuntimeError: Failed to open source. """ if isinstance(video_source, str) and not os.path.isfile(video_source): raise RuntimeError("Video file does not exist: " + video_source) with SuppressStderr(): cap = cv2.VideoCapture(video_source) if not cap.isOpened(): raise RuntimeError(f"Failed to open video: {video_source}") return cap def _get_cap(self): """Return the cap. If multiprocessing, create a new one. Else return the one from construction time.""" return self.open_video(self.video_source) if self.multiprocessing else self.cap def get_fps(self) -> int: """Get the FPS of the capture device.""" return self._get_cap().get(cv2.CAP_PROP_FPS) # type: ignore def get_frame(self) -> Any: """Get next frame. For a file, this will be the next frame, whereas for a camera source, it will be the next available frame.""" ret, frame = self._get_cap().read() if not ret: raise RuntimeError("Failed to read frame.") # Switch color order if desired if self.color_order == ColorOrder.RGB: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # move channel dim frame = np.moveaxis(frame, -1, self.channel_dim) return self.transform(frame) if self.transform is not None else frame
[docs] class VideoFileDataset(Dataset, VideoDataset): """ Video dataset from file. This class requires that OpenCV be installed. """ def __init__(self, *args, **kwargs) -> None: VideoDataset.__init__(self, *args, **kwargs) num_frames = self.get_num_frames() if self.max_num_frames is None or num_frames < self.max_num_frames: self.max_num_frames = num_frames @staticmethod def get_available_codecs() -> dict[str, str]: """Try different codecs, see which are available. Returns a dictionary with of available codecs with codecs as keys and file extensions as values.""" if not has_cv2: return {} all_codecs = {"mp4v": ".mp4", "X264": ".avi", "H264": ".mp4", "MP42": ".mp4", "MJPG": ".mjpeg", "DIVX": ".avi"} codecs = {} with SuppressStderr(): with tempfile.TemporaryDirectory() as tmp_dir: for codec, ext in all_codecs.items(): writer = cv2.VideoWriter() fname = os.path.join(tmp_dir, f"test{ext}") fourcc = cv2.VideoWriter_fourcc(*codec) # type: ignore[attr-defined] noviderr =, fourcc, 1, (10, 10)) if noviderr: codecs[codec] = ext writer.release() return codecs def get_num_frames(self) -> int: """ Return the number of frames in a video file. Raises: RuntimeError: no frames found. """ num_frames = int(self._get_cap().get(cv2.CAP_PROP_FRAME_COUNT)) if num_frames == 0: raise RuntimeError("0 frames found") return num_frames def __len__(self): return self.max_num_frames def __getitem__(self, index: int) -> Any: """ Fetch single data item from index. """ if self.max_num_frames is not None and index >= self.max_num_frames: raise IndexError self._get_cap().set(cv2.CAP_PROP_POS_FRAMES, index) return self.get_frame()
[docs] class CameraDataset(IterableDataset, VideoDataset): """ Video dataset from a capture device (e.g., webcam). This class requires that OpenCV be installed. Args: video_source: index of capture device. `get_num_devices` can be used to determine possible devices. transform: transform to be applied to each frame. max_num_frames: Max number of frames to iterate across. If `None` is passed, then the dataset will iterate infinitely. Raises: RuntimeError: OpenCV not installed. """ @staticmethod def get_num_devices() -> int: """Get number of possible devices detected by OpenCV that can be used for capture.""" if not has_cv2: return 0 num_devices = 0 while True: cap = cv2.VideoCapture(num_devices) if not[0]: break num_devices += 1 cap.release() return num_devices def __iter__(self): frame_count = 0 while True: frame = self.get_frame() frame_count += 1 yield frame if self.max_num_frames is not None: if frame_count == self.max_num_frames: break