176 lines
7.8 KiB
Python
176 lines
7.8 KiB
Python
"""Upload a video once, then decode it into the existing footage model."""
|
||
|
||
import hashlib
|
||
import json
|
||
import subprocess
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
from fractions import Fraction
|
||
from pathlib import Path
|
||
|
||
from django.db import close_old_connections, transaction
|
||
|
||
from . import blobs
|
||
from .models import Blob, Extraction, Footage, FootageFrame
|
||
|
||
_active = set()
|
||
_lock = threading.Lock()
|
||
TIMEOUT = 3600
|
||
|
||
|
||
def _command(args):
|
||
result = subprocess.run(args, capture_output=True, text=True, timeout=TIMEOUT)
|
||
if result.returncode:
|
||
raise ValueError((result.stderr or result.stdout or "media tool failed")[-1200:])
|
||
return result.stdout
|
||
|
||
|
||
def _decode_frames(job, source_path, frames_dir, facts, root):
|
||
"""Decode one frame per source frame and publish ffmpeg's live frame count."""
|
||
progress_path = root / "frames.progress"
|
||
log_path = root / "frames.log"
|
||
total = facts.get("reported_frames") or round(facts["duration"] * facts["fps"])
|
||
args = ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
|
||
"-stats_period", "0.25", "-progress", str(progress_path),
|
||
"-i", str(source_path), "-fps_mode", "passthrough",
|
||
str(frames_dir / "%04d.png")]
|
||
with open(log_path, "wb") as log:
|
||
proc = subprocess.Popen(args, stdout=log, stderr=subprocess.STDOUT)
|
||
deadline = time.monotonic() + TIMEOUT
|
||
try:
|
||
while proc.poll() is None:
|
||
if time.monotonic() >= deadline:
|
||
raise TimeoutError("video frame extraction timed out")
|
||
if progress_path.exists():
|
||
lines = progress_path.read_text(errors="replace").splitlines()
|
||
count = next((int(line[6:].strip()) for line in reversed(lines)
|
||
if line.startswith("frame=") and
|
||
line[6:].strip().isdigit()), 0)
|
||
if count and total:
|
||
progress = min(59, int(60 * count / total))
|
||
if progress > job.progress:
|
||
job.progress = progress
|
||
job.save(update_fields=["progress", "updated"])
|
||
time.sleep(0.2)
|
||
finally:
|
||
if proc.poll() is None:
|
||
proc.kill()
|
||
proc.wait()
|
||
if proc.returncode:
|
||
raise ValueError(log_path.read_text(errors="replace")[-1200:] or
|
||
"video frame extraction failed")
|
||
|
||
|
||
def probe(path):
|
||
data = json.loads(_command(["ffprobe", "-v", "error", "-show_streams",
|
||
"-show_format", "-of", "json", str(path)]))
|
||
video = next((s for s in data.get("streams", []) if s.get("codec_type") == "video"), None)
|
||
if not video:
|
||
raise ValueError("the uploaded file has no video stream")
|
||
nominal = Fraction(video.get("r_frame_rate") or "0")
|
||
average = Fraction(video.get("avg_frame_rate") or "0")
|
||
if nominal <= 0 or average <= 0:
|
||
raise ValueError("the video's frame rate is unknown")
|
||
vfr = abs(float(nominal / average) - 1) > 0.001
|
||
if vfr:
|
||
raise ValueError("variable-frame-rate video needs timestamp-aware playback")
|
||
frames = video.get("nb_frames")
|
||
duration = float(data.get("format", {}).get("duration") or 0)
|
||
if ((frames and frames.isdigit() and int(frames) > 900)
|
||
or (duration > 0 and duration * float(average) > 901)):
|
||
raise ValueError("video is longer than the 900-frame footage limit")
|
||
return {"fps": float(average), "nominal_fps": float(nominal),
|
||
"width": int(video["width"]), "height": int(video["height"]),
|
||
"duration": duration,
|
||
"reported_frames": int(frames) if frames and frames.isdigit() else None,
|
||
"has_audio": any(s.get("codec_type") == "audio" for s in data.get("streams", [])),
|
||
"vfr": False}
|
||
|
||
|
||
def extraction_key(source, settings):
|
||
text = json.dumps({"scheme": 1, "source": source.blob_id, "settings": settings},
|
||
sort_keys=True, separators=(",", ":"))
|
||
return "sha256:" + hashlib.sha256(text.encode()).hexdigest()
|
||
|
||
|
||
def _register(job, frames, audio_path, facts):
|
||
width, height = blobs.png_size(frames[0])
|
||
frame_blobs = []
|
||
for index, path in enumerate(frames):
|
||
if blobs.png_size(path) != (width, height):
|
||
raise ValueError(f"decoded frame {index + 1} has different dimensions")
|
||
digest, size = blobs.adopt(path)
|
||
frame_blobs.append((index, digest, size))
|
||
audio_digest, audio_size = blobs.adopt(audio_path)
|
||
h = hashlib.sha256()
|
||
h.update(f"arthur-footage-1/{facts['fps']}/{len(frames)}/{width}x{height}\n".encode())
|
||
for _, digest, _ in frame_blobs:
|
||
h.update(digest.encode())
|
||
h.update(audio_digest.encode())
|
||
with transaction.atomic():
|
||
audio_blob, _ = Blob.objects.get_or_create(
|
||
digest=audio_digest, defaults={"size": audio_size, "media_type": "audio/wav"})
|
||
footage, created = Footage.objects.get_or_create(
|
||
digest=h.hexdigest(),
|
||
defaults={"label": job.source.filename[:200], "source": job.source.filename[:200],
|
||
"fps": facts["fps"],
|
||
"frames": len(frames), "width": width,
|
||
"height": height, "audio": audio_blob})
|
||
if created:
|
||
rows = []
|
||
for index, digest, size in frame_blobs:
|
||
blob, _ = Blob.objects.get_or_create(
|
||
digest=digest, defaults={"size": size, "media_type": "image/png"})
|
||
rows.append(FootageFrame(footage=footage, index=index, blob=blob))
|
||
FootageFrame.objects.bulk_create(rows)
|
||
return footage
|
||
|
||
|
||
def run(key):
|
||
close_old_connections()
|
||
try:
|
||
job = Extraction.objects.select_related("source", "source__blob").get(key=key)
|
||
job.state, job.progress, job.error = "running", 0, ""
|
||
job.save(update_fields=["state", "progress", "error", "updated"])
|
||
facts = job.source.probe
|
||
source_path = blobs.path_for(job.source.blob_id)
|
||
with tempfile.TemporaryDirectory(prefix="arthur-extract-") as directory:
|
||
root = Path(directory)
|
||
frames_dir = root / "frames"
|
||
frames_dir.mkdir()
|
||
_decode_frames(job, source_path, frames_dir, facts, root)
|
||
frames = sorted(frames_dir.glob("*.png"))
|
||
expected = facts.get("reported_frames")
|
||
if not frames or len(frames) > 900 or (expected and len(frames) != expected):
|
||
raise ValueError(f"decoded {len(frames)} frames; expected {expected or '1–900'}")
|
||
job.progress = 60
|
||
job.save(update_fields=["progress", "updated"])
|
||
audio_path = root / "audio.wav"
|
||
if facts["has_audio"]:
|
||
_command(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
|
||
"-i", str(source_path), "-vn", "-ac", "1", "-ar", "44100",
|
||
str(audio_path)])
|
||
else:
|
||
_command(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
|
||
"-f", "lavfi", "-i", "anullsrc=r=44100:cl=mono",
|
||
"-t", str(len(frames) / facts["fps"]), "-c:a", "pcm_s16le",
|
||
str(audio_path)])
|
||
footage = _register(job, frames, audio_path, facts)
|
||
job.footage, job.state, job.progress = footage, "done", 100
|
||
job.save(update_fields=["footage", "state", "progress", "updated"])
|
||
except Exception as exc:
|
||
Extraction.objects.filter(key=key).update(state="failed", error=str(exc)[:2000])
|
||
finally:
|
||
with _lock:
|
||
_active.discard(key)
|
||
close_old_connections()
|
||
|
||
|
||
def enqueue(key):
|
||
with _lock:
|
||
if key in _active:
|
||
return
|
||
_active.add(key)
|
||
threading.Thread(target=run, args=(key,), daemon=True,
|
||
name=f"arthur-extract-{key[7:15]}").start()
|