| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from io import BytesIO |
| 5 | |
| 6 | import numpy as np |
| 7 | import numpy.typing as npt |
| 8 | |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.utils.import_utils import PlaceholderModule |
| 11 | |
| 12 | from .base import ( |
| 13 | VideoSourceMetadata, |
| 14 | VideoTargetMetadata, |
| 15 | check_frame_pixel_limit, |
| 16 | ) |
| 17 | |
| 18 | try: |
| 19 | import cv2 |
| 20 | import cv2.videoio_registry as vr |
| 21 | except ImportError: |
| 22 | cv2 = PlaceholderModule("cv2") |
| 23 | vr = PlaceholderModule("cv2").placeholder_attr("videoio_registry") |
| 24 | |
| 25 | logger = init_logger(__name__) |
| 26 | |
| 27 | |
| 28 | def decode_opencv( |
| 29 | loader_cls, |
| 30 | data: bytes, |
| 31 | target: VideoTargetMetadata, |
| 32 | sampling_kwargs: dict, |
| 33 | *, |
| 34 | frame_recovery: bool = False, |
| 35 | ) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: |
| 36 | cap = OpenCVVideoBackendMixin.open_video_capture(data) |
| 37 | check_frame_pixel_limit( |
| 38 | int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), |
| 39 | int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), |
| 40 | ) |
| 41 | source = loader_cls._prepare_source(OpenCVVideoBackendMixin.get_video_metadata(cap)) |
| 42 | frame_idx = loader_cls.compute_frames_index_to_sample( |
| 43 | source=source, target=target, **sampling_kwargs |
| 44 | ) |
| 45 | frames, valid = loader_cls.read_frames( |
| 46 | cap, |
| 47 | frame_idx, |
| 48 | total_frames_num=source.total_frames_num, |
| 49 | frame_recovery=frame_recovery, |
| 50 | ) |
| 51 | return frames, source, frame_idx, valid |
| 52 | |
| 53 | |
| 54 | class OpenCVVideoBackendMixin: |
| 55 | @staticmethod |
| 56 | def get_cv2_video_api(): |
| 57 | api_pref = None |
| 58 | for backend in vr.getStreamBufferedBackends(): |
| 59 | if not vr.hasBackend(backend): |
| 60 | continue |
| 61 | if not vr.isBackendBuiltIn(backend): |
| 62 | _, abi, api = vr.getStreamBufferedBackendPluginVersion(backend) |
| 63 | if abi < 1 or (abi == 1 and api < 2): |
| 64 | continue |
| 65 | api_pref = backend |
| 66 | break |
| 67 | return api_pref |
| 68 | |
| 69 | @classmethod |
| 70 | def open_video_capture(cls, data: bytes) -> "cv2.VideoCapture": |
| 71 | backend = cls.get_cv2_video_api() |
| 72 | cap = cv2.VideoCapture(BytesIO(data), backend, []) |
| 73 | if not cap.isOpened(): |
| 74 | raise ValueError("Could not open video stream") |
| 75 | return cap |
| 76 | |
| 77 | @staticmethod |
| 78 | def get_video_metadata(cap: "cv2.VideoCapture") -> VideoSourceMetadata: |
| 79 | total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| 80 | original_fps = cap.get(cv2.CAP_PROP_FPS) |
| 81 | # CAP_PROP_FRAME_COUNT counts every physical sample in the container, |
| 82 | # overstating the presentable frames when an mp4 edit list hides the |
| 83 | # decode lead-in (e.g. a lossless `ffmpeg -ss ... -c copy` trim). |
| 84 | # Cross-check with an end-of-stream seek: the FFMPEG backend reports |
| 85 | # the resulting position against the edit-list-aware timeline, which |
| 86 | # exposes the true frame count without decoding anything. The seek |
| 87 | # may land mid-stream on corrupt files, so only trust it when the |
| 88 | # stream really ends there. A one-frame tolerance keeps mere header |
| 89 | # miscounts from being "corrected". |
| 90 | if total_frames_num > 0 and cap.set(cv2.CAP_PROP_POS_AVI_RATIO, 1.0): |
| 91 | visible_frames = int(cap.get(cv2.CAP_PROP_POS_FRAMES)) |
| 92 | at_stream_end = not cap.grab() |
| 93 | cap.set(cv2.CAP_PROP_POS_FRAMES, 0) |
| 94 | if at_stream_end and 0 < visible_frames < total_frames_num - 1: |
| 95 | logger.warning( |
| 96 | "Video header claims %d frames but only %d are " |
| 97 | "presentable; sampling over the true frame count.", |
| 98 | total_frames_num, |
| 99 | visible_frames, |
| 100 | ) |
| 101 | total_frames_num = visible_frames |
| 102 | duration = total_frames_num / original_fps if original_fps > 0 else 0 |
| 103 | return VideoSourceMetadata( |
| 104 | total_frames_num=total_frames_num, |
| 105 | original_fps=original_fps, |
| 106 | duration=duration, |
| 107 | ) |
| 108 | |
| 109 | @classmethod |
| 110 | def _can_use_for_recovery( |
| 111 | cls, |
| 112 | idx: int, |
| 113 | failed_frames: list[int], |
| 114 | next_target_map: dict[int, int], |
| 115 | total_frames: int, |
| 116 | ) -> bool: |
| 117 | """Check if current frame can recover the oldest failed frame.""" |
| 118 | if not failed_frames: |
| 119 | return False |
| 120 | oldest_failed = failed_frames[0] |
| 121 | limit = next_target_map.get(oldest_failed, total_frames) |
| 122 | return idx < limit |
| 123 | |
| 124 | @classmethod |
| 125 | def _read_frames_with_recovery( |
| 126 | cls, |
| 127 | cap: "cv2.VideoCapture", |
| 128 | frame_indices: list[int], |
| 129 | total_frames: int, |
| 130 | ) -> tuple[npt.NDArray, list[int], dict[int, int]]: |
| 131 | """Read frames with dynamic window forward-scan recovery. |
| 132 | |
| 133 | When a target frame fails to load, the next successfully grabbed |
| 134 | frame (before the next target frame) will be used to recover it. |
| 135 | |
| 136 | Args: |
| 137 | cap: OpenCV VideoCapture object |
| 138 | frame_indices: Sorted list of target frame indices to load |
| 139 | total_frames: Total number of frames in the video |
| 140 | |
| 141 | Returns: |
| 142 | Tuple of (frames_array, valid_frame_indices, recovered_map) |
| 143 | - frames_array: Array of loaded frames |
| 144 | - valid_frame_indices: List of frame indices that were loaded |
| 145 | - recovered_map: Dict mapping recovered_idx -> source_idx |
| 146 | |
| 147 | """ |
| 148 | width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| 149 | height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
| 150 | |
| 151 | assert width > 0 and height > 0, ( |
| 152 | f"Invalid video frame size: width={width}, height={height}" |
| 153 | ) |
| 154 | |
| 155 | frame_idx_set = set(frame_indices) |
| 156 | max_frame_idx = frame_indices[-1] if frame_indices else 0 |
| 157 | |
| 158 | # Build map: target_idx -> next_target_idx (for recovery window) |
| 159 | next_target_map: dict[int, int] = {} |
| 160 | for k in range(len(frame_indices) - 1): |
| 161 | next_target_map[frame_indices[k]] = frame_indices[k + 1] |
| 162 | next_target_map[frame_indices[-1]] = total_frames |
| 163 | |
| 164 | frames_list: list[npt.NDArray] = [] |
| 165 | valid_frame_indices: list[int] = [] |
| 166 | failed_frames_idx: list[int] = [] |
| 167 | recovered_map: dict[int, int] = {} |
| 168 | |
| 169 | i = 0 |
| 170 | for idx in range(max_frame_idx + 1): |
| 171 | is_target_frame = idx in frame_idx_set |
| 172 | |
| 173 | # Attempt to grab the current frame |
| 174 | ok = cap.grab() |
| 175 | |
| 176 | if not ok: |
| 177 | if is_target_frame: |
| 178 | logger.debug( |
| 179 | "Failed to grab frame %d during video loading.", |
| 180 | idx, |
| 181 | ) |
| 182 | failed_frames_idx.append(idx) |
| 183 | continue |
| 184 | |
| 185 | # Check if we should retrieve: target frame OR can recover a failed one |
| 186 | can_recover = cls._can_use_for_recovery( |
| 187 | idx, failed_frames_idx, next_target_map, total_frames |
| 188 | ) |
| 189 | |
| 190 | if is_target_frame or can_recover: |
| 191 | ret, frame = cap.retrieve() |
| 192 | |
| 193 | if ret and frame is not None and frame.size > 0: |
| 194 | rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| 195 | frames_list.append(rgb_frame) |
| 196 | valid_frame_indices.append(idx) |
| 197 | i += 1 |
| 198 | |
| 199 | if can_recover: |
| 200 | recovered_idx = failed_frames_idx.pop(0) |
| 201 | recovered_map[recovered_idx] = idx |
| 202 | logger.info( |
| 203 | "Recovered frame %d using frame %d (delay: %d)", |
| 204 | recovered_idx, |
| 205 | idx, |
| 206 | idx - recovered_idx, |
| 207 | ) |
| 208 | elif is_target_frame: |
| 209 | logger.debug( |
| 210 | "Failed to retrieve frame %d during video loading.", |
| 211 | idx, |
| 212 | ) |
| 213 | failed_frames_idx.append(idx) |
| 214 | |
| 215 | # Log any remaining failed frames |
| 216 | for failed_idx in failed_frames_idx: |
| 217 | logger.debug( |
| 218 | "Frame %d could not be recovered (end of video).", |
| 219 | failed_idx, |
| 220 | ) |
| 221 | |
| 222 | # Stack frames |
| 223 | if frames_list: |
| 224 | frames = np.stack(frames_list) |
| 225 | else: |
| 226 | frames = np.empty((0, height, width, 3), dtype=np.uint8) |
| 227 | |
| 228 | return frames, valid_frame_indices, recovered_map |
| 229 | |
| 230 | @classmethod |
| 231 | def _read_frames_no_recovery( |
| 232 | cls, |
| 233 | cap, |
| 234 | frame_indices: set[int], |
| 235 | max_frame_idx: int, |
| 236 | ) -> tuple[npt.NDArray, list[int]]: |
| 237 | num_expected_frames = len(frame_indices) |
| 238 | width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| 239 | height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
| 240 | frames = np.empty((num_expected_frames, height, width, 3), dtype=np.uint8) |
| 241 | |
| 242 | i = 0 |
| 243 | valid_frame_indices = [] |
| 244 | for idx in range(max_frame_idx + 1): |
| 245 | ok = cap.grab() |
| 246 | if not ok: |
| 247 | # Frame is broken/unreadable, skip it |
| 248 | if idx in frame_indices: |
| 249 | logger.debug( |
| 250 | "Failed to grab frame %d during video loading. " |
| 251 | "This frame will be skipped.", |
| 252 | idx, |
| 253 | ) |
| 254 | continue |
| 255 | if idx in frame_indices: |
| 256 | ret, frame = cap.retrieve() |
| 257 | if ret: |
| 258 | frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| 259 | valid_frame_indices.append(idx) |
| 260 | i += 1 |
| 261 | else: |
| 262 | # retrieve() failed even though grab() succeeded |
| 263 | logger.debug( |
| 264 | "Failed to retrieve frame %d during video loading. " |
| 265 | "This frame will be skipped.", |
| 266 | idx, |
| 267 | ) |
| 268 | |
| 269 | valid_num_frames = len(valid_frame_indices) |
| 270 | if valid_num_frames < num_expected_frames: |
| 271 | logger.warning( |
| 272 | "Video loading completed with %d broken/unreadable frames. " |
| 273 | "Expected %d frames but only loaded %d frames.", |
| 274 | num_expected_frames - valid_num_frames, |
| 275 | num_expected_frames, |
| 276 | valid_num_frames, |
| 277 | ) |
| 278 | |
| 279 | return frames[:valid_num_frames], valid_frame_indices |
| 280 | |
| 281 | @classmethod |
| 282 | def read_frames( |
| 283 | cls, |
| 284 | cap: "cv2.VideoCapture", |
| 285 | frame_idx: list[int], |
| 286 | total_frames_num: int, |
| 287 | *, |
| 288 | frame_recovery: bool = False, |
| 289 | ) -> tuple[npt.NDArray, list[int]]: |
| 290 | if frame_recovery: |
| 291 | num_frames_to_sample = len(frame_idx) |
| 292 | frames, valid_frame_indices, recovered_map = cls._read_frames_with_recovery( |
| 293 | cap, frame_idx, total_frames_num |
| 294 | ) |
| 295 | |
| 296 | if recovered_map: |
| 297 | logger.info( |
| 298 | "Frame recovery: %d frames recovered using forward scan.", |
| 299 | len(recovered_map), |
| 300 | ) |
| 301 | else: |
| 302 | frame_idx_set = set(frame_idx) |
| 303 | num_frames_to_sample = len(frame_idx_set) |
| 304 | frames, valid_frame_indices = cls._read_frames_no_recovery( |
| 305 | cap, frame_idx_set, max(frame_idx) |
| 306 | ) |
| 307 | valid_num_frames = len(valid_frame_indices) |
| 308 | if valid_num_frames < num_frames_to_sample: |
| 309 | logger.warning( |
| 310 | "Video loading completed with %d broken/unreadable frames. " |
| 311 | "Expected to sample %d frames but only loaded %d frames.", |
| 312 | num_frames_to_sample - valid_num_frames, |
| 313 | num_frames_to_sample, |
| 314 | valid_num_frames, |
| 315 | ) |
| 316 | return frames, valid_frame_indices |