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
(
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,
)
| 64 | import_cv() |
| 65 | |
| 66 | def __init__( |
| 67 | self, |
| 68 | video_source: str | int, |
| 69 | transform: Callable | None = None, |
| 70 | max_num_frames: int | None = None, |
| 71 | color_order: str = ColorOrder.RGB, |
| 72 | multiprocessing: bool = False, |
| 73 | channel_dim: int = 0, |
| 74 | ) -> None: |
| 75 | """ |
| 76 | Base video dataset. |
| 77 | |
| 78 | Args: |
| 79 | video_source: filename of video. |
| 80 | transform: transform to be applied to each frame. |
| 81 | max_num_frames: Max number of frames to iterate across. If `None` is passed, |
| 82 | then the dataset will iterate until the end of the file. |
| 83 | color_order: Color order to return frame. Default is RGB. |
| 84 | multiprocessing: If `True`, open the video source on the fly. This makes |
| 85 | things process-safe, which is useful when combined with a DataLoader |
| 86 | with `num_workers>0`. However, when using with `num_workers==0`, it |
| 87 | makes sense to use `multiprocessing=False`, as the source will then |
| 88 | only be opened once, at construction, which will be faster in those |
| 89 | circumstances. |
| 90 | channel_dim: OpenCV reads with the channel as the last dimension. Use this |
| 91 | flag to move it elsewhere. By default this is zero, so the channel |
| 92 | dimension is moved to the front. |
| 93 | |
| 94 | Raises: |
| 95 | RuntimeError: OpenCV not installed. |
| 96 | NotImplementedError: Unknown color order. |
| 97 | """ |
| 98 | if not has_cv2: |
| 99 | raise RuntimeError("OpenCV not installed.") |
| 100 | if color_order not in ColorOrder: |
| 101 | raise NotImplementedError |
| 102 | |
| 103 | self.color_order = color_order |
| 104 | self.channel_dim = channel_dim |
| 105 | self.video_source = video_source |
| 106 | self.multiprocessing = multiprocessing |
| 107 | if not multiprocessing: |
| 108 | self.cap = self.open_video(video_source) |
| 109 | self.transform = transform |
| 110 | self.max_num_frames = max_num_frames |
| 111 | |
| 112 | @staticmethod |
| 113 | def open_video(video_source: str | int): |