5e2094f2ad
FastAPI service in infra/media-worker implementing every MediaInferenceProvider operation behind HTTP API v1, plus capabilities. Bearer auth, MEDIA_ROOT realpath enforcement, strict request schemas, cancellation on disconnect, GPU/CPU gates with 503 overflow, input size/duration limits, lazy SHA256-pinned model downloads, fail-closed redaction with atomic output, offline pytest suite with fake backends, CPU/CUDA/ROCm Dockerfile and a license review of every model. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0138YqDj3ttfayRwuMKKrqVj
436 lines
14 KiB
Python
436 lines
14 KiB
Python
"""ffmpeg/ffprobe plumbing. Every call is an argv list; nothing goes through a shell."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import json
|
|
import os
|
|
import queue
|
|
import re
|
|
import subprocess
|
|
import threading
|
|
import uuid
|
|
from collections import deque
|
|
from collections.abc import Callable, Iterator
|
|
from dataclasses import dataclass
|
|
from fractions import Fraction
|
|
|
|
import numpy as np
|
|
|
|
from .errors import BadRequest, Cancelled, MediaError
|
|
|
|
FFMPEG = "ffmpeg"
|
|
FFPROBE = "ffprobe"
|
|
SHOWINFO_PTS = re.compile(rb"pts_time:\s*(-?[0-9.]+)")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VideoInfo:
|
|
width: int
|
|
height: int
|
|
fps: float
|
|
fps_text: str
|
|
duration_s: float
|
|
has_audio: bool
|
|
has_video: bool
|
|
start_s: float = 0.0
|
|
# Video stream duration when the container reports it (audio may run longer).
|
|
video_duration_s: float = 0.0
|
|
|
|
def frame_ms(self) -> float:
|
|
return 1000.0 / self.fps
|
|
|
|
|
|
def run_cancellable(
|
|
argv: list[str], cancel: threading.Event, timeout: float | None = None
|
|
) -> subprocess.CompletedProcess[bytes]:
|
|
"""Run a short-lived tool, killing it when the request is cancelled."""
|
|
proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False)
|
|
killer = _kill_on_cancel(proc, cancel)
|
|
try:
|
|
stdout, stderr = proc.communicate(timeout=timeout)
|
|
finally:
|
|
killer.set()
|
|
if cancel.is_set():
|
|
raise Cancelled()
|
|
return subprocess.CompletedProcess(argv, proc.returncode, stdout, stderr)
|
|
|
|
|
|
def _kill_on_cancel(proc: subprocess.Popen[bytes], cancel: threading.Event) -> threading.Event:
|
|
done = threading.Event()
|
|
|
|
def watch() -> None:
|
|
while not done.is_set():
|
|
if cancel.wait(0.2):
|
|
if proc.poll() is None:
|
|
proc.kill()
|
|
return
|
|
|
|
threading.Thread(target=watch, daemon=True).start()
|
|
return done
|
|
|
|
|
|
def _parse_rate(text: str | None) -> float:
|
|
if not text or text in ("0/0", "0"):
|
|
return 0.0
|
|
try:
|
|
return float(Fraction(text))
|
|
except (ValueError, ZeroDivisionError):
|
|
return 0.0
|
|
|
|
|
|
def probe(path: str, cancel: threading.Event) -> VideoInfo:
|
|
result = run_cancellable(
|
|
[
|
|
FFPROBE,
|
|
"-v",
|
|
"error",
|
|
"-print_format",
|
|
"json",
|
|
"-show_streams",
|
|
"-show_format",
|
|
path,
|
|
],
|
|
cancel,
|
|
)
|
|
if result.returncode != 0:
|
|
raise BadRequest("input is not a readable media file")
|
|
data = json.loads(result.stdout or b"{}")
|
|
streams = data.get("streams", [])
|
|
video = next(
|
|
(
|
|
s
|
|
for s in streams
|
|
if s.get("codec_type") == "video" and not s.get("disposition", {}).get("attached_pic")
|
|
),
|
|
None,
|
|
)
|
|
audio = next((s for s in streams if s.get("codec_type") == "audio"), None)
|
|
duration = float(data.get("format", {}).get("duration") or 0.0)
|
|
if video is None:
|
|
return VideoInfo(0, 0, 0.0, "0", duration, audio is not None, False)
|
|
fps_text = video.get("avg_frame_rate") or "0/0"
|
|
fps = _parse_rate(fps_text)
|
|
if fps <= 0 or fps > 240:
|
|
fps_text = video.get("r_frame_rate") or "0/0"
|
|
fps = _parse_rate(fps_text)
|
|
if fps <= 0:
|
|
fps_text, fps = "25", 25.0
|
|
width, height = int(video.get("width") or 0), int(video.get("height") or 0)
|
|
rotation = _rotation(video)
|
|
if rotation in (90, 270):
|
|
width, height = height, width
|
|
if not duration:
|
|
duration = float(video.get("duration") or 0.0)
|
|
start = float(video.get("start_time") or 0.0)
|
|
video_duration = float(video.get("duration") or 0.0) or duration
|
|
return VideoInfo(width, height, fps, fps_text, duration, audio is not None, True, start, video_duration)
|
|
|
|
|
|
def _rotation(stream: dict) -> int:
|
|
for side in stream.get("side_data_list", []) or []:
|
|
if "rotation" in side:
|
|
return int(abs(int(side["rotation"]))) % 360
|
|
tag = (stream.get("tags") or {}).get("rotate")
|
|
return int(abs(int(tag))) % 360 if tag else 0
|
|
|
|
|
|
def check_duration(info: VideoInfo, max_seconds: int) -> None:
|
|
if info.duration_s > max_seconds:
|
|
raise MediaError("too_large", "input duration exceeds MEDIA_MAX_SECONDS")
|
|
|
|
|
|
def scaled_size(width: int, height: int, max_side: int | None) -> tuple[int, int]:
|
|
if not max_side or max(width, height) <= max_side:
|
|
return width - width % 2, height - height % 2
|
|
scale = max_side / max(width, height)
|
|
w = max(2, int(round(width * scale / 2)) * 2)
|
|
h = max(2, int(round(height * scale / 2)) * 2)
|
|
return w, h
|
|
|
|
|
|
class _RawReader:
|
|
"""Read fixed-size BGR frames from an ffmpeg stdout pipe."""
|
|
|
|
def __init__(
|
|
self,
|
|
argv: list[str],
|
|
width: int,
|
|
height: int,
|
|
cancel: threading.Event,
|
|
pts: queue.Queue[float] | None = None,
|
|
) -> None:
|
|
self.size = width * height * 3
|
|
self.shape = (height, width, 3)
|
|
self.proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False)
|
|
self.cancel = cancel
|
|
self.pts = pts
|
|
self._stderr: deque[bytes] = deque(maxlen=20)
|
|
self._drain = threading.Thread(target=self._read_stderr, daemon=True)
|
|
self._drain.start()
|
|
|
|
def _read_stderr(self) -> None:
|
|
assert self.proc.stderr is not None
|
|
for line in self.proc.stderr:
|
|
if self.pts is not None and b"Parsed_showinfo" in line:
|
|
match = SHOWINFO_PTS.search(line)
|
|
if match:
|
|
self.pts.put(float(match.group(1)))
|
|
continue
|
|
self._stderr.append(line)
|
|
|
|
def frames(self) -> Iterator[np.ndarray]:
|
|
assert self.proc.stdout is not None
|
|
try:
|
|
while True:
|
|
if self.cancel.is_set():
|
|
raise Cancelled()
|
|
buf = self.proc.stdout.read(self.size)
|
|
if not buf:
|
|
break
|
|
if len(buf) != self.size:
|
|
raise MediaError("internal", "truncated frame from decoder")
|
|
yield np.frombuffer(buf, np.uint8).reshape(self.shape)
|
|
code = self.proc.wait()
|
|
self._drain.join(timeout=5)
|
|
if code != 0:
|
|
raise MediaError("internal", "decoder failed: " + self.stderr_tail())
|
|
finally:
|
|
self.close()
|
|
|
|
def stderr_tail(self) -> str:
|
|
return b"".join(self._stderr).decode("utf-8", "replace").strip()[-500:]
|
|
|
|
def close(self) -> None:
|
|
if self.proc.poll() is None:
|
|
self.proc.kill()
|
|
self.proc.wait()
|
|
|
|
|
|
def stream_frames(
|
|
path: str,
|
|
info: VideoInfo,
|
|
cancel: threading.Event,
|
|
sample_fps: float | None = None,
|
|
max_side: int | None = None,
|
|
) -> Iterator[tuple[int, np.ndarray]]:
|
|
"""Yield (tMs, frame) for every frame (CFR-normalized) or about `sample_fps` per second.
|
|
|
|
Sampled frames are real source frames; their timestamps come from
|
|
ffmpeg's showinfo, so tMs is exact even for variable frame rates.
|
|
"""
|
|
width, height = scaled_size(info.width, info.height, max_side)
|
|
filters: list[str] = []
|
|
if sample_fps:
|
|
# The source frame nearest to each k / sample_fps slot (phase-locked, no drift).
|
|
half = 0.5 / info.fps
|
|
slot = f"floor((t+{half:.6f})*{sample_fps:.6f})"
|
|
prev = f"floor((prev_selected_t+{half:.6f})*{sample_fps:.6f})"
|
|
filters += [f"select='isnan(prev_selected_t)+gt({slot}\\,{prev})'", "showinfo"]
|
|
if (width, height) != (info.width, info.height):
|
|
filters.append(f"scale={width}:{height}:flags=area")
|
|
level = "info" if sample_fps else "error"
|
|
argv = [FFMPEG, "-v", level, "-nostats", "-nostdin", "-i", path, "-map", "0:v:0", "-an", "-sn"]
|
|
if filters:
|
|
argv += ["-vf", ",".join(filters)]
|
|
if sample_fps:
|
|
argv += ["-fps_mode", "passthrough"]
|
|
else:
|
|
argv += ["-fps_mode", "cfr", "-r", info.fps_text]
|
|
argv += ["-f", "rawvideo", "-pix_fmt", "bgr24", "pipe:1"]
|
|
pts: queue.Queue[float] | None = queue.Queue() if sample_fps else None
|
|
reader = _RawReader(argv, width, height, cancel, pts)
|
|
for index, frame in enumerate(reader.frames()):
|
|
if pts is None:
|
|
yield int(round(index * 1000.0 / info.fps)), frame
|
|
continue
|
|
try:
|
|
t = pts.get(timeout=30)
|
|
except queue.Empty as error:
|
|
raise MediaError("internal", "missing frame timestamp from decoder") from error
|
|
yield max(int(round((t - info.start_s) * 1000)), 0), frame
|
|
|
|
|
|
def frame_at(
|
|
path: str, info: VideoInfo, t_ms: int, cancel: threading.Event, max_side: int | None = None
|
|
) -> np.ndarray:
|
|
"""Decode the frame shown at t_ms (accurate seek)."""
|
|
width, height = scaled_size(info.width, info.height, max_side)
|
|
t = min(max(t_ms / 1000.0, 0.0), max(info.duration_s - 0.5 / info.fps, 0.0))
|
|
argv = [
|
|
FFMPEG,
|
|
"-v",
|
|
"error",
|
|
"-nostdin",
|
|
"-ss",
|
|
f"{t:.3f}",
|
|
"-i",
|
|
path,
|
|
"-map",
|
|
"0:v:0",
|
|
"-frames:v",
|
|
"1",
|
|
]
|
|
if (width, height) != (info.width, info.height):
|
|
argv += ["-vf", f"scale={width}:{height}:flags=area"]
|
|
argv += ["-f", "rawvideo", "-pix_fmt", "bgr24", "pipe:1"]
|
|
frames = list(_RawReader(argv, width, height, cancel).frames())
|
|
if not frames:
|
|
raise MediaError("internal", f"no frame decoded at {t_ms} ms")
|
|
return frames[0]
|
|
|
|
|
|
def frames_at(
|
|
path: str, info: VideoInfo, times_ms: list[int], cancel: threading.Event, max_side: int | None = None
|
|
) -> Iterator[tuple[int, np.ndarray]]:
|
|
for t_ms in times_ms:
|
|
if cancel.is_set():
|
|
raise Cancelled()
|
|
yield t_ms, frame_at(path, info, t_ms, cancel, max_side)
|
|
|
|
|
|
def temp_sibling(final_path: str) -> str:
|
|
directory, name = os.path.split(final_path)
|
|
stem, ext = os.path.splitext(name)
|
|
return os.path.join(directory, f".{stem}.partial-{uuid.uuid4().hex}{ext}")
|
|
|
|
|
|
def encoder_args(out_path: str) -> list[str]:
|
|
"""Near-lossless intermediate: FFV1 in Matroska, otherwise x264 CRF 12."""
|
|
if out_path.lower().endswith(".mkv"):
|
|
return ["-c:v", "ffv1", "-level", "3", "-pix_fmt", "yuv444p", "-g", "1"]
|
|
return ["-c:v", "libx264", "-preset", "medium", "-crf", "12", "-pix_fmt", "yuv420p"]
|
|
|
|
|
|
FrameFn = Callable[[int, int, np.ndarray], np.ndarray]
|
|
|
|
|
|
def transform_video(
|
|
path: str,
|
|
out_path: str,
|
|
info: VideoInfo,
|
|
cancel: threading.Event,
|
|
frame_fn: FrameFn,
|
|
) -> int:
|
|
"""Decode every frame, apply frame_fn(index, tMs, frame), encode, copy audio.
|
|
|
|
Fails closed: any decode, transform, or encode error removes the temporary
|
|
output and raises; out_path only ever appears complete (atomic rename).
|
|
"""
|
|
if not info.has_video:
|
|
raise BadRequest("input has no video stream")
|
|
width, height = scaled_size(info.width, info.height, None)
|
|
tmp = temp_sibling(out_path)
|
|
enc_argv = [
|
|
FFMPEG,
|
|
"-v",
|
|
"error",
|
|
"-nostdin",
|
|
"-y",
|
|
"-f",
|
|
"rawvideo",
|
|
"-pix_fmt",
|
|
"bgr24",
|
|
"-s",
|
|
f"{width}x{height}",
|
|
"-r",
|
|
info.fps_text,
|
|
"-i",
|
|
"pipe:0",
|
|
"-i",
|
|
path,
|
|
"-map",
|
|
"0:v:0",
|
|
"-map",
|
|
"1:a?",
|
|
"-map_metadata",
|
|
"-1",
|
|
"-map_chapters",
|
|
"-1",
|
|
*encoder_args(out_path),
|
|
"-c:a",
|
|
"copy",
|
|
tmp,
|
|
]
|
|
dec_argv = [
|
|
FFMPEG,
|
|
"-v",
|
|
"error",
|
|
"-xerror",
|
|
"-nostdin",
|
|
"-i",
|
|
path,
|
|
"-map",
|
|
"0:v:0",
|
|
"-an",
|
|
"-sn",
|
|
"-fps_mode",
|
|
"cfr",
|
|
"-r",
|
|
info.fps_text,
|
|
]
|
|
if (width, height) != (info.width, info.height):
|
|
dec_argv += ["-vf", f"crop={width}:{height}:0:0"]
|
|
dec_argv += ["-f", "rawvideo", "-pix_fmt", "bgr24", "pipe:1"]
|
|
encoder = subprocess.Popen(
|
|
enc_argv, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, shell=False
|
|
)
|
|
enc_err: list[bytes] = []
|
|
drain = threading.Thread(target=lambda: enc_err.extend(encoder.stderr or []), daemon=True)
|
|
drain.start()
|
|
reader = _RawReader(dec_argv, width, height, cancel)
|
|
count = 0
|
|
ok = False
|
|
try:
|
|
for index, frame in enumerate(reader.frames()):
|
|
t_ms = int(round(index * 1000.0 / info.fps))
|
|
out = frame_fn(index, t_ms, frame.copy())
|
|
if out.shape != frame.shape or out.dtype != np.uint8:
|
|
raise MediaError("internal", "frame transform changed the frame shape")
|
|
assert encoder.stdin is not None
|
|
try:
|
|
encoder.stdin.write(np.ascontiguousarray(out).tobytes())
|
|
except BrokenPipeError as error:
|
|
raise MediaError("internal", "encoder stopped accepting frames") from error
|
|
count += 1
|
|
expected = (info.video_duration_s or info.duration_s) * info.fps
|
|
if count == 0 or (expected > 0 and count < expected * 0.98 - 2):
|
|
raise MediaError("internal", f"decoded {count} of ~{int(expected)} frames")
|
|
assert encoder.stdin is not None
|
|
encoder.stdin.close()
|
|
if encoder.wait() != 0:
|
|
drain.join(timeout=5)
|
|
raise MediaError(
|
|
"internal", "encoder failed: " + b"".join(enc_err[-5:]).decode("utf-8", "replace")[-500:]
|
|
)
|
|
drain.join(timeout=5)
|
|
os.replace(tmp, out_path)
|
|
ok = True
|
|
return count
|
|
finally:
|
|
reader.close()
|
|
if not ok:
|
|
if encoder.poll() is None:
|
|
encoder.kill()
|
|
with contextlib.suppress(subprocess.TimeoutExpired):
|
|
encoder.wait(timeout=10)
|
|
if os.path.exists(tmp):
|
|
os.unlink(tmp)
|
|
|
|
|
|
_NUMBER = r"(-?(?:inf|\d+(?:\.\d+)?))"
|
|
|
|
|
|
def parse_ebur128(stderr: str) -> tuple[float | None, float | None]:
|
|
"""Integrated loudness and true peak from the ebur128 summary block."""
|
|
summary = stderr[stderr.rfind("Summary:") :] if "Summary:" in stderr else ""
|
|
lufs = re.search(r"I:\s+" + _NUMBER + r"\s+LUFS", summary)
|
|
peak = re.search(r"Peak:\s+" + _NUMBER + r"\s+dBFS", summary)
|
|
|
|
def num(match: re.Match[str] | None) -> float | None:
|
|
if not match or "inf" in match.group(1):
|
|
return None
|
|
return float(match.group(1))
|
|
|
|
return num(lufs), num(peak)
|