File size: 5,568 Bytes
1405c30 | 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 | """
Video I/O utilities: load frames + audio from a video file, save frames back
to video with audio, and trim audio to match video duration.
"""
import os
import subprocess
import tempfile
from pathlib import Path
import numpy as np
from PIL import Image
# ---------------------------------------------------------------------------
# Loading
# ---------------------------------------------------------------------------
def load_video_frames(
path: str,
fps: float = 24.0,
max_frames: int | None = None,
) -> tuple[np.ndarray, float]:
"""
Decode video frames to a uint8 numpy array [N, H, W, 3].
Returns (frames, actual_fps).
Uses decord when available; falls back to opencv.
"""
try:
import decord
decord.bridge.set_bridge("native")
vr = decord.VideoReader(path, ctx=decord.cpu(0))
actual_fps = float(vr.get_avg_fps())
total = len(vr)
if max_frames is not None:
total = min(total, max_frames)
indices = list(range(total))
frames = vr.get_batch(indices).asnumpy() # [N, H, W, 3]
return frames, actual_fps
except ImportError:
pass
import cv2
cap = cv2.VideoCapture(path)
actual_fps = cap.get(cv2.CAP_PROP_FPS) or fps
frames = []
while True:
ret, frame = cap.read()
if not ret:
break
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
if max_frames is not None and len(frames) >= max_frames:
break
cap.release()
return np.stack(frames, axis=0), actual_fps
def extract_audio(video_path: str, output_path: str) -> bool:
"""Extract audio track from video to a WAV file. Returns False if no audio."""
result = subprocess.run(
[
"ffprobe", "-v", "quiet", "-select_streams", "a",
"-show_entries", "stream=codec_type",
"-of", "csv=p=0", video_path,
],
capture_output=True, text=True,
)
if "audio" not in result.stdout:
return False
subprocess.run(
[
"ffmpeg", "-y", "-i", video_path,
"-vn", "-acodec", "pcm_s16le",
"-ar", "44100", "-ac", "2", output_path,
],
capture_output=True, check=True,
)
return True
# ---------------------------------------------------------------------------
# Saving
# ---------------------------------------------------------------------------
def save_video(
frames: np.ndarray,
fps: float,
output_path: str,
audio_path: str | None = None,
audio_duration: float | None = None,
crf: int = 19,
) -> str:
"""
Encode frames [N, H, W, 3] uint8 to an mp4 file.
Optionally mux audio_path (trimmed to audio_duration seconds if provided).
Returns the path to the written file.
"""
N, H, W, _ = frames.shape
tmp_video = output_path + ".noaudio.mp4"
# Write raw video with ffmpeg via stdin pipe
cmd = [
"ffmpeg", "-y",
"-f", "rawvideo",
"-vcodec", "rawvideo",
"-s", f"{W}x{H}",
"-pix_fmt", "rgb24",
"-r", str(fps),
"-i", "pipe:0",
"-vcodec", "libx264",
"-pix_fmt", "yuv420p",
"-crf", str(crf),
"-preset", "fast",
tmp_video,
]
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
for frame in frames:
proc.stdin.write(frame.tobytes())
proc.stdin.close()
proc.wait()
if audio_path and os.path.exists(audio_path):
duration_flag = ["-t", str(audio_duration)] if audio_duration else []
subprocess.run(
[
"ffmpeg", "-y",
"-i", tmp_video,
"-i", audio_path,
*duration_flag,
"-c:v", "copy",
"-c:a", "aac", "-b:a", "192k",
"-shortest",
output_path,
],
capture_output=True, check=True,
)
os.remove(tmp_video)
else:
os.rename(tmp_video, output_path)
return output_path
# ---------------------------------------------------------------------------
# Resolution helpers
# ---------------------------------------------------------------------------
def align_to(value: int, multiple: int = 32) -> int:
"""Round value up to the nearest multiple."""
return ((value + multiple - 1) // multiple) * multiple
def compute_target_size(
orig_w: int,
orig_h: int,
base_resolution: int = 768,
multiple: int = 32,
) -> tuple[int, int]:
"""
Scale the longer edge to base_resolution, preserving aspect ratio,
then align both dimensions to `multiple`.
"""
scale = base_resolution / max(orig_w, orig_h)
new_w = align_to(int(orig_w * scale), multiple)
new_h = align_to(int(orig_h * scale), multiple)
return new_w, new_h
def resize_frames(frames: np.ndarray, target_w: int, target_h: int) -> np.ndarray:
"""Resize [N, H, W, 3] frames to target_w x target_h."""
if frames.shape[2] == target_w and frames.shape[1] == target_h:
return frames
out = np.empty((len(frames), target_h, target_w, 3), dtype=np.uint8)
for i, f in enumerate(frames):
out[i] = np.array(Image.fromarray(f).resize((target_w, target_h), Image.LANCZOS))
return out
def frames_for_duration(fps: float, duration: float) -> int:
"""Return frame count aligned to LTX-2.3 requirements: ((n * fps) // 8) * 8 + 1."""
raw = int(duration * fps)
return ((raw // 8) * 8) + 1
|