61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
"""Video I/O utilities: read frames (cv2.VideoCapture) + encode mp4 (imageio/ffmpeg).
|
|
|
|
Shared by scripts/demo.py.
|
|
"""
|
|
|
|
from typing import List, Tuple
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def read_video_frames(path: str) -> Tuple[List[np.ndarray], float, Tuple[int, int]]:
|
|
"""Read all video frames (BGR uint8).
|
|
|
|
Returns:
|
|
frames_bgr: list of (H, W, 3) BGR uint8
|
|
fps: frame rate (30 if unreadable)
|
|
(H, W): frame size
|
|
"""
|
|
cap = cv2.VideoCapture(path)
|
|
if not cap.isOpened():
|
|
raise FileNotFoundError(f"Cannot open video: {path}")
|
|
|
|
fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
|
frames: List[np.ndarray] = []
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
frames.append(frame)
|
|
cap.release()
|
|
|
|
if not frames:
|
|
raise RuntimeError(f"Video has no valid frames: {path}")
|
|
|
|
H, W = frames[0].shape[:2]
|
|
return frames, float(fps), (H, W)
|
|
|
|
|
|
def write_video_ffmpeg(
|
|
frames_bgr: List[np.ndarray], path: str, fps: float, quality: int = 8
|
|
) -> None:
|
|
"""Encode a BGR frame sequence as mp4 (imageio + libx264).
|
|
|
|
Args:
|
|
frames_bgr: list of (H, W, 3) BGR uint8
|
|
path: output mp4 path
|
|
fps: frame rate
|
|
quality: imageio quality (0-10, higher is better)
|
|
"""
|
|
import imageio.v2 as imageio
|
|
|
|
writer = imageio.get_writer(path, fps=fps, codec="libx264", quality=quality,
|
|
macro_block_size=1)
|
|
try:
|
|
for frame in frames_bgr:
|
|
# BGR -> RGB
|
|
writer.append_data(frame[:, :, ::-1])
|
|
finally:
|
|
writer.close()
|