Files
reTraining/backend/video.py
T
asus 5c7c122105 feat: add counting bench, triage, and dataset modules
This commit includes major additions and updates to the frontend and backend architectures, introducing new dataset management, live counting features, batch processing, and triage logic. Includes new UI pages, components, and API routes.
2026-08-14 16:28:52 +07:00

203 lines
6.8 KiB
Python

"""ffmpeg/ffprobe wrappers: read a video's metadata, cut a range into frames.
ffmpeg is a hard dependency rather than an OpenCV fallback because seeking to a
timestamp in a long recording has to be exact — an off-by-a-few-seconds trim
silently produces frames of the wrong thing (REQ-020…022).
"""
import json
import os
import shutil
import subprocess
import tempfile
from typing import Callable, List, Optional
VIDEO_EXTS = (".mp4", ".mkv", ".mov", ".avi", ".webm", ".m4v")
_probe_cache: dict = {}
class VideoError(Exception):
pass
def available() -> bool:
return shutil.which("ffmpeg") is not None and shutil.which("ffprobe") is not None
def probe(path: str) -> dict:
"""Duration, resolution and fps of a video, cached by (path, mtime, size).
The cache matters: the Library page probes every file in a date folder, and
ffprobe on a cold cache over dozens of long recordings is slow (REQ-012).
"""
try:
stat = os.stat(path)
except OSError as exc:
raise VideoError(str(exc))
key = (os.path.realpath(path), stat.st_mtime, stat.st_size)
if key in _probe_cache:
return _probe_cache[key]
if not available():
raise VideoError("ffprobe is not installed in this environment")
result = subprocess.run(
["ffprobe", "-v", "error", "-select_streams", "v:0",
"-show_entries", "stream=width,height,avg_frame_rate,codec_name",
"-show_entries", "format=duration",
"-of", "json", path],
capture_output=True, text=True,
)
if result.returncode != 0:
raise VideoError(result.stderr.strip().splitlines()[-1] if result.stderr else "ffprobe failed")
payload = json.loads(result.stdout or "{}")
streams = payload.get("streams") or [{}]
info = {
"duration": float(payload.get("format", {}).get("duration") or 0.0),
"width": int(streams[0].get("width") or 0),
"height": int(streams[0].get("height") or 0),
"fps": _parse_fps(streams[0].get("avg_frame_rate")),
"codec_name": str(streams[0].get("codec_name") or ""),
"size": stat.st_size,
}
_probe_cache[key] = info
return info
def _parse_fps(value: Optional[str]) -> float:
# ffprobe reports "30000/1001", and "0/0" for streams it cannot work out.
if not value or "/" not in value:
return 0.0
numerator, denominator = value.split("/", 1)
try:
return round(float(numerator) / float(denominator), 3) if float(denominator) else 0.0
except ValueError:
return 0.0
def frame_count(start_sec: float, end_sec: float, fps: float) -> int:
"""How many frames a trim will produce — shown before extraction (REQ-021)."""
span = max(0.0, end_sec - start_sec)
return max(0, int(span * fps))
def extract_frames(
video_path: str,
out_dir: str,
start_sec: float,
end_sec: float,
fps: float,
on_progress: Optional[Callable[[int], None]] = None,
should_stop: Optional[Callable[[], bool]] = None,
) -> List[str]:
"""Cut [start, end] at `fps` into numbered JPEGs. Returns their filenames.
`-ss` before `-i` seeks by keyframe (fast) and ffmpeg then decodes accurately
from there, which is what makes a two-minute cut out of a two-hour recording
quick instead of a full decode.
"""
if not available():
raise VideoError("ffmpeg is not installed in this environment")
if end_sec <= start_sec:
raise VideoError("The end of the range must be after its start")
if fps <= 0:
raise VideoError("fps must be greater than 0")
os.makedirs(out_dir, exist_ok=True)
command = [
"ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
"-ss", f"{start_sec:.3f}", "-to", f"{end_sec:.3f}", "-i", video_path,
"-vf", f"fps={fps}", "-q:v", "2",
os.path.join(out_dir, "%06d.jpg"),
]
# stderr goes to a file, not a pipe: a damaged recording emits decode errors
# without end, and nothing reads the pipe until ffmpeg exits — so a 64KB
# buffer would fill and block ffmpeg forever, hanging the whole job queue.
with tempfile.TemporaryFile("w+") as stderr_log:
process = subprocess.Popen(command, stdout=subprocess.DEVNULL,
stderr=stderr_log, text=True)
expected = frame_count(start_sec, end_sec, fps)
while process.poll() is None:
if should_stop is not None and should_stop():
process.terminate()
process.wait(timeout=10)
raise VideoError("cancelled")
if on_progress is not None:
on_progress(min(_written(out_dir), expected))
try:
process.wait(timeout=1)
except subprocess.TimeoutExpired:
pass
if process.returncode != 0:
stderr_log.seek(0)
lines = stderr_log.read().strip().splitlines()
raise VideoError(lines[-1] if lines else "ffmpeg failed")
return sorted(name for name in os.listdir(out_dir) if name.endswith(".jpg"))
def _written(out_dir: str) -> int:
try:
return sum(1 for name in os.listdir(out_dir) if name.endswith(".jpg"))
except OSError:
return 0
def convert_to_h264(
input_path: str,
output_path: Optional[str] = None,
crf: int = 18,
preset: str = "medium",
) -> str:
"""Transcode input video to browser-supported H.264 / AAC MP4.
Uses CRF 18 for visually lossless compression (no perceptible visual loss,
significantly reduced file size), yuv420p for universal browser playback,
and +faststart for instant web streaming.
"""
if not available():
raise VideoError("ffmpeg is not installed in this environment")
if not os.path.isfile(input_path):
raise VideoError(f"Input video file not found: {input_path}")
if not output_path:
base, _ = os.path.splitext(input_path)
output_path = f"{base}_h264.mp4"
temp_output_path = output_path.replace(".mp4", "_tmp.mp4") if output_path.endswith(".mp4") else output_path + ".tmp.mp4"
os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
command = [
"ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
"-i", input_path,
"-c:v", "libx264",
"-crf", str(crf),
"-preset", preset,
"-pix_fmt", "yuv420p",
"-movflags", "+faststart",
"-c:a", "aac",
"-b:a", "128k",
"-f", "mp4",
temp_output_path,
]
result = subprocess.run(command, capture_output=True, text=True)
if result.returncode != 0:
if os.path.isfile(temp_output_path):
os.remove(temp_output_path)
err_msg = result.stderr.strip().splitlines()[-1] if result.stderr else "ffmpeg conversion failed"
raise VideoError(err_msg)
os.replace(temp_output_path, output_path)
return output_path