File size: 8,347 Bytes
319eb16 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 | #!/usr/bin/env python3
"""
Video helper utilities for robust loading of videos with codec handling and
backend fallbacks. Centralized here to be reused across dataset loaders.
"""
import gc
import os
import shutil
import subprocess
import tempfile
import time
from signal import SIGTERM
from typing import Optional
import cv2
import numpy as np
# Configuration constants
DEFAULT_SUBPROCESS_TIMEOUT = 30 # seconds for ffprobe/ffmpeg operations
DEFAULT_FRAME_READ_TIMEOUT = 60 # seconds for reading all frames from a video
MAX_FRAMES_SANITY_CHECK = 10000 # maximum number of frames to prevent infinite loops
def _ffprobe_codec_name(path: str, timeout: int = DEFAULT_SUBPROCESS_TIMEOUT) -> str | None:
"""Return codec_name for the first video stream using ffprobe, or None on failure."""
if shutil.which("ffprobe") is None:
return None
try:
completed = subprocess.run(
[
"ffprobe",
"-v",
"error",
"-select_streams",
"v:0",
"-show_entries",
"stream=codec_name",
"-of",
"default=nk=1:nw=1",
path,
],
capture_output=True,
check=True,
text=True,
timeout=timeout,
)
codec_name = completed.stdout.strip().splitlines()[0] if completed.stdout.strip() else None
return codec_name
except subprocess.TimeoutExpired:
print(f"Warning: ffprobe timed out after {timeout}s for {path}")
return None
except Exception as e:
print(f"Warning: ffprobe failed for {path}: {e}")
return None
def _reencode_to_h264(input_path: str, timeout: int = DEFAULT_SUBPROCESS_TIMEOUT) -> str | None:
"""Re-encode input video to H.264 yuv420p if ffmpeg is available. Returns output path or None."""
if shutil.which("ffmpeg") is None:
return None
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_out:
output_path = tmp_out.name
try:
subprocess.run(
[
"ffmpeg",
"-y",
"-loglevel",
"error",
"-i",
input_path,
"-c:v",
"libx264",
"-pix_fmt",
"yuv420p",
"-movflags",
"+faststart",
output_path,
],
check=True,
timeout=timeout,
)
return output_path
except subprocess.TimeoutExpired:
print(f"Warning: ffmpeg re-encode timed out after {timeout}s for {input_path}")
try:
os.unlink(output_path)
except Exception:
pass
return None
except Exception as e:
print(f"Warning: ffmpeg re-encode failed for {input_path}: {e}")
try:
os.unlink(output_path)
except Exception:
pass
return None
def _open_with_best_backend(path: str) -> cv2.VideoCapture | None:
"""Try multiple OpenCV backends and return an opened capture or None."""
backends: list[int] = [getattr(cv2, "CAP_FFMPEG", cv2.CAP_ANY), cv2.CAP_ANY]
for backend in backends:
try:
cap_try = cv2.VideoCapture(path, backend)
if cap_try.isOpened():
# Validate we can read at least one frame
ret, test_frame = cap_try.read()
if ret and test_frame is not None:
cap_try.set(cv2.CAP_PROP_POS_FRAMES, 0)
return cap_try
cap_try.release()
except Exception:
try:
cap_try.release()
except Exception:
pass
return None
def load_video_frames(
video_input,
timeout: Optional[int] = DEFAULT_FRAME_READ_TIMEOUT,
max_frames: int = MAX_FRAMES_SANITY_CHECK,
) -> np.ndarray:
"""Load video frames (RGB uint8) from a file path (str/Path) or video bytes.
- For byte inputs, detect AV1 and re-encode to H.264 for compatibility
- Uses OpenCV with FFMPEG backend when available
- Returns numpy array of shape (T, H, W, 3) in RGB order
- Protected by timeout and maximum frame count to prevent hangs
Args:
video_input: Path to video file or video bytes
timeout: Maximum time in seconds to spend reading frames (None to disable)
max_frames: Maximum number of frames to read (sanity check against infinite loops)
Raises:
TimeoutError: If frame reading exceeds the timeout
ValueError: If video cannot be opened or no frames extracted
"""
temp_files_to_cleanup: list[str] = []
cap: cv2.VideoCapture | None = None
video_path_str = str(video_input) if isinstance(video_input, (str, os.PathLike)) else "bytes"
if isinstance(video_input, (str, os.PathLike)):
video_path = str(video_input)
decodable_path = video_path
# If file is AV1, re-encode to H.264 so OpenCV can decode (many platforms lack AV1)
codec_name = _ffprobe_codec_name(video_path)
if codec_name == "av1":
h264_path = _reencode_to_h264(video_path, timeout=120)
if h264_path is not None:
temp_files_to_cleanup.append(h264_path)
decodable_path = h264_path
cap = _open_with_best_backend(decodable_path)
else:
# Save bytes to temp file first
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
temp_file.write(video_input)
temp_input_path = temp_file.name
temp_files_to_cleanup.append(temp_input_path)
# If the input is AV1, transcode to H.264 for compatibility
codec_name = _ffprobe_codec_name(temp_input_path)
decodable_path = temp_input_path
if codec_name == "av1":
h264_path = _reencode_to_h264(temp_input_path)
if h264_path is not None:
temp_files_to_cleanup.append(h264_path)
decodable_path = h264_path
cap = _open_with_best_backend(decodable_path)
try:
frames: list[np.ndarray] = []
if cap is None or not cap.isOpened():
raise ValueError(
f"Could not open video file {video_path_str} with available backends. "
"If the source is AV1, install AV1 support or enable ffmpeg re-encode."
)
start_time = time.time()
frame_count = 0
while True:
# Check timeout
if timeout is not None and (time.time() - start_time) > timeout:
raise TimeoutError(
f"Video frame reading timed out after {timeout}s for {video_path_str}. "
f"Read {frame_count} frames before timeout."
)
# Check max frames sanity limit
if frame_count >= max_frames:
print(
f"Warning: Reached maximum frame limit ({max_frames}) for {video_path_str}. "
f"This might indicate a corrupted or unusually long video."
)
frames = []
gc.collect()
break
ret, frame = cap.read()
if not ret:
break
# Convert BGR to RGB for consistency
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frames.append(rgb_frame)
frame_count += 1
cap.release()
# Clean up temp file(s) if we created any
for path in temp_files_to_cleanup:
try:
os.unlink(path)
except Exception:
pass
if len(frames) == 0:
raise ValueError(f"No frames could be extracted from video {video_path_str}")
return np.array(frames)
except Exception as e:
if cap is not None:
cap.release()
# Clean up temp files in case of error
for path in temp_files_to_cleanup:
try:
os.unlink(path)
except Exception:
pass
# Add more context to the error
if isinstance(e, (TimeoutError, ValueError)):
raise
raise RuntimeError(f"Error loading video {video_path_str}: {e}") from e
|