Add video upload, extraction progress, and reusable analysis sources

clips/extraction.py

@@ -0,0 +1,176 @@
+"""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()