Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| _IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".bmp", ".webp"} | |
| class VideoInput: | |
| def __init__( | |
| self, | |
| source: str | Path, | |
| *, | |
| K: np.ndarray | None = None, | |
| D: np.ndarray | None = None, | |
| stride: int = 1, | |
| ) -> None: | |
| self.source = Path(source) | |
| self._mode = "" | |
| self._cap: cv2.VideoCapture | None = None | |
| self._frame_paths: list[Path] = [] | |
| # `_sample_index` counts frames returned by read() (after stride subsampling). | |
| # `_cap_pos` tracks the next raw frame index the video capture will return | |
| # (grab-only increments this without decoding, so we can cheaply skip). | |
| self._sample_index = 0 | |
| self._cap_pos = 0 | |
| self._total_frames = 0 | |
| self._num_frames = 0 | |
| self._image_width = 0 | |
| self._image_height = 0 | |
| self._fps: float | None = None | |
| self._stride = max(1, int(stride)) | |
| # Optional OpenCV pinhole intrinsics + distortion. When both are provided | |
| # (and D has any non-zero element), every frame returned by read() / | |
| # read_frame() is automatically passed through cv2.undistort. This keeps | |
| # downstream inference / camera projection in a clean pinhole space and | |
| # avoids the caller having to wrap VideoInput just for undistortion. | |
| self._K = None if K is None else np.asarray(K, dtype=np.float32) | |
| self._D = None if D is None else np.asarray(D, dtype=np.float32).reshape(-1) | |
| self._undistort_enabled = ( | |
| self._K is not None | |
| and self._D is not None | |
| and bool(np.any(self._D != 0)) | |
| ) | |
| if self.source.is_dir(): | |
| self._mode = "frames" | |
| self._frame_paths = [ | |
| path for path in self.source.iterdir() if path.is_file() and path.suffix.lower() in _IMAGE_SUFFIXES | |
| ] | |
| self._frame_paths.sort() | |
| if len(self._frame_paths) == 0: | |
| raise ValueError(f"No image frames found in directory: {self.source}") | |
| first_frame = cv2.imread(str(self._frame_paths[0]), cv2.IMREAD_COLOR) | |
| if first_frame is None: | |
| raise ValueError(f"Failed to read first frame: {self._frame_paths[0]}") | |
| self._total_frames = len(self._frame_paths) | |
| self._image_height, self._image_width = first_frame.shape[:2] | |
| elif self.source.is_file(): | |
| self._mode = "video" | |
| self._cap = cv2.VideoCapture(str(self.source)) | |
| if not self._cap.isOpened(): | |
| raise ValueError(f"Failed to open video: {self.source}") | |
| self._total_frames = int(self._cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| self._image_width = int(self._cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| self._image_height = int(self._cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = float(self._cap.get(cv2.CAP_PROP_FPS)) | |
| self._fps = fps if fps > 0 else None | |
| else: | |
| raise ValueError(f"VideoInput source does not exist: {self.source}") | |
| # Sampled frame indices (into the raw video) after stride subsampling. | |
| # With stride == 1 this is [0, 1, ..., total_frames - 1] and behavior is | |
| # backwards-compatible. | |
| self._frame_ids = np.arange(0, self._total_frames, self._stride, dtype=np.int64) | |
| self._num_frames = int(self._frame_ids.size) | |
| def num_frames(self) -> int: | |
| """Number of frames yielded by read() (after stride subsampling).""" | |
| return self._num_frames | |
| def total_num_frames(self) -> int: | |
| """Raw frame count of the underlying video / directory, ignoring stride.""" | |
| return self._total_frames | |
| def stride(self) -> int: | |
| return self._stride | |
| def frame_ids(self) -> np.ndarray: | |
| """Raw frame indices corresponding to each read() call, shape (num_frames,).""" | |
| return self._frame_ids | |
| def image_width(self) -> int: | |
| return self._image_width | |
| def image_height(self) -> int: | |
| return self._image_height | |
| def fps(self) -> float | None: | |
| return self._fps | |
| def reset(self) -> None: | |
| self._sample_index = 0 | |
| self._cap_pos = 0 | |
| if self._cap is not None: | |
| self._cap.set(cv2.CAP_PROP_POS_FRAMES, 0) | |
| def _maybe_undistort(self, frame_bgr: np.ndarray) -> np.ndarray: | |
| if self._undistort_enabled: | |
| return cv2.undistort(frame_bgr, self._K, self._D) | |
| return frame_bgr | |
| def read_frame(self, frame_index: int) -> np.ndarray: | |
| """Random-access read by *raw* frame index (unaffected by stride).""" | |
| frame_index = int(frame_index) | |
| if frame_index < 0 or frame_index >= self._total_frames: | |
| raise IndexError(f"frame_index out of range: {frame_index}") | |
| if self._mode == "frames": | |
| frame_path = self._frame_paths[frame_index] | |
| frame_bgr = cv2.imread(str(frame_path), cv2.IMREAD_COLOR) | |
| if frame_bgr is None: | |
| raise ValueError(f"Failed to read frame: {frame_path}") | |
| frame_bgr = self._maybe_undistort(frame_bgr) | |
| return cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) | |
| if self._cap is None: | |
| raise ValueError("VideoInput is not initialized with a readable source.") | |
| saved_pos = self._cap_pos | |
| self._cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index) | |
| ret, frame_bgr = self._cap.read() | |
| # Restore the streaming pointer so an interleaved read() sequence keeps its place. | |
| self._cap.set(cv2.CAP_PROP_POS_FRAMES, saved_pos) | |
| if not ret: | |
| raise ValueError(f"Failed to read frame at index: {frame_index}") | |
| frame_bgr = self._maybe_undistort(frame_bgr) | |
| return cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) | |
| def read(self) -> np.ndarray | None: | |
| """Return the next sampled frame (advances by `stride` under the hood).""" | |
| if self._sample_index >= self._num_frames: | |
| return None | |
| target = int(self._frame_ids[self._sample_index]) | |
| self._sample_index += 1 | |
| if self._mode == "frames": | |
| frame_path = self._frame_paths[target] | |
| frame_bgr = cv2.imread(str(frame_path), cv2.IMREAD_COLOR) | |
| if frame_bgr is None: | |
| raise ValueError(f"Failed to read frame: {frame_path}") | |
| frame_bgr = self._maybe_undistort(frame_bgr) | |
| return cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) | |
| if self._cap is None: | |
| raise ValueError("VideoInput is not initialized with a readable source.") | |
| # Cheap forward skip via grab() (does not decode), then decode the target frame. | |
| # For stride == 1 this reduces to the plain read() path (single grab+retrieve). | |
| while self._cap_pos < target: | |
| if not self._cap.grab(): | |
| return None | |
| self._cap_pos += 1 | |
| ret, frame_bgr = self._cap.read() | |
| if not ret: | |
| return None | |
| self._cap_pos += 1 | |
| frame_bgr = self._maybe_undistort(frame_bgr) | |
| return cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) | |
| def close(self) -> None: | |
| if self._cap is not None: | |
| self._cap.release() | |
| self._cap = None | |