diff --git a/infra/media-worker/.dockerignore b/infra/media-worker/.dockerignore new file mode 100644 index 0000000..949e720 --- /dev/null +++ b/infra/media-worker/.dockerignore @@ -0,0 +1,7 @@ +.pytest_cache +.ruff_cache +.mypy_cache +.venv +.venv-ml +**/__pycache__ +**/*.pyc diff --git a/infra/media-worker/.gitignore b/infra/media-worker/.gitignore new file mode 100644 index 0000000..01304c1 --- /dev/null +++ b/infra/media-worker/.gitignore @@ -0,0 +1,5 @@ +.pytest_cache/ +.venv/ +.venv +.venv-ml +__pycache__/ diff --git a/infra/media-worker/Dockerfile b/infra/media-worker/Dockerfile new file mode 100644 index 0000000..baa6bee --- /dev/null +++ b/infra/media-worker/Dockerfile @@ -0,0 +1,68 @@ +# syntax=docker/dockerfile:1.7 +# Media Studio inference sidecar. +# docker build --build-arg ACCELERATOR=cpu|cuda|rocm -t rakazo/media-worker infra/media-worker +# docker build --target test -t rakazo/media-worker:test infra/media-worker # offline pytest +# Model weights are not baked in; they download into /models on first use. +# Behind a TLS-intercepting proxy, pass its CA with --secret id=pip_ca,src=. + +ARG ACCELERATOR=cpu + +FROM python:3.11-slim-bookworm@sha256:2333bd330d12de02514770b3585cad313644316047cdee24a7acfdece6de6efb AS runtime +ARG ACCELERATOR + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + TQDM_DISABLE=1 \ + MEDIA_ACCELERATOR=${ACCELERATOR} \ + MEDIA_MODELS_DIR=/models \ + MEDIA_WORKER_PORT=8790 + +RUN case "${ACCELERATOR}" in cpu|cuda|rocm) ;; *) echo "ACCELERATOR must be cpu, cuda or rocm" >&2; exit 1 ;; esac \ + && apt-get update \ + && apt-get install -y --no-install-recommends ffmpeg libsndfile1 \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /srv/media-worker +COPY requirements.txt requirements-ml-cpu.txt requirements-ml-cuda.txt requirements-ml-rocm.txt ./ + +# faster-whisper depends on the CPU onnxruntime wheel; the CUDA image swaps in +# onnxruntime-gpu. SAM2 is an sdist built against the installed torch, without +# its optional CUDA extension. +RUN --mount=type=secret,id=pip_ca,required=false \ + if [ -s /run/secrets/pip_ca ]; then export PIP_CERT=/run/secrets/pip_ca; fi \ + && pip install -r requirements.txt -r "requirements-ml-${ACCELERATOR}.txt" \ + && if [ "${ACCELERATOR}" = "cuda" ]; then \ + pip uninstall -y onnxruntime && pip install --no-deps --force-reinstall onnxruntime-gpu==1.30.0; \ + fi \ + && SAM2_BUILD_CUDA=0 pip install --no-build-isolation --no-deps sam2==1.1.0 \ + && python -c "import cv2, scenedetect, librosa, onnxruntime, faster_whisper, torch, sam2, tokenizers" + +RUN useradd --uid 10001 --create-home --shell /usr/sbin/nologin media \ + && mkdir -p /models \ + && chown media:media /models +VOLUME ["/models"] + +COPY app ./app + +USER media +EXPOSE 8790 +HEALTHCHECK --interval=30s --timeout=5s --start-period=20s --retries=3 \ + CMD ["python", "-c", "import os, urllib.request; urllib.request.urlopen('http://127.0.0.1:' + os.environ.get('MEDIA_WORKER_PORT', '8790') + '/healthz', timeout=4)"] +CMD ["python", "-m", "app"] + + +# Offline test image: same runtime plus pytest; fake backends, no weights. +FROM runtime AS test +USER root +COPY requirements-dev.txt pytest.ini ./ +RUN --mount=type=secret,id=pip_ca,required=false \ + if [ -s /run/secrets/pip_ca ]; then export PIP_CERT=/run/secrets/pip_ca; fi \ + && pip install -r requirements-dev.txt +COPY tests ./tests +USER media +CMD ["python", "-m", "pytest", "-q"] + + +FROM runtime diff --git a/infra/media-worker/LICENSES.md b/infra/media-worker/LICENSES.md new file mode 100644 index 0000000..1a349a8 --- /dev/null +++ b/infra/media-worker/LICENSES.md @@ -0,0 +1,53 @@ +# Licenses: models and libraries of the media sidecar + +Checked October 2026 against the upstream LICENSE files and model cards linked +below. "Commercial use" means the published code or weights license permits it. +Training-data terms were not reviewed legally; where provenance is unclear the +backend is optional, never a default. Weights are downloaded at runtime from the +pinned revision (see `app/models.py`) and are never part of this repository or +the image. + +## Models + +| Name | Purpose | License | Commercial use | Default / optional | Source | +|---|---|---|---|---|---| +| OpenCV YuNet 2023mar | Face detection, 5-point landmarks (`detectFaces`, `trackFaces`, `faceLandmarks`) | MIT | yes | default | https://github.com/opencv/opencv_zoo/tree/main/models/face_detection_yunet | +| Whisper small (CTranslate2 conversion by Systran) | Transcription with word timestamps (`transcribe`) | MIT | yes | default | https://huggingface.co/Systran/faster-whisper-small, https://github.com/openai/whisper | +| SigLIP base patch16-224 (ONNX export by Xenova) | Frame embeddings (`embedFrames`); zero-shot content scores (`classifyContent`, experimental) | Apache-2.0 | yes | default | https://huggingface.co/google/siglip-base-patch16-224, https://huggingface.co/Xenova/siglip-base-patch16-224 | +| SAM 2.1 Hiera tiny | Click/box segmentation with video propagation (`segmentPrompted`, experimental) | Apache-2.0 | yes | default | https://github.com/facebookresearch/sam2, https://huggingface.co/facebook/sam2.1-hiera-tiny | +| Big-LaMa (ONNX export by Carve) | Per-frame inpainting with optical-flow smoothing (`inpaintVideo`, experimental) | Apache-2.0 | yes | default | https://github.com/advimman/lama, https://huggingface.co/Carve/LaMa-ONNX | +| OpenCV SFace 2021dec | Identity embeddings for local person clustering (`trackFaces.embedding`) | Apache-2.0 (weights); training-data provenance unverified | unclear | optional (`MEDIA_OPTIONAL_BACKENDS=faceIdentity`) | https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface | +| MediaPipe Face Landmarker | Dense face landmarks | Apache-2.0 | yes | not used: the `mediapipe` wheel pulls the GUI OpenCV build; YuNet's 5 points cover the partial modes | https://github.com/google-ai-edge/mediapipe | +| OpenCLIP (LAION weights) | Frame embeddings | MIT | yes | not used: SigLIP runs on ONNX Runtime without torch | https://github.com/mlfoundations/open_clip | +| NudeNet v3 | Nudity detection | AGPL-3.0 | copyleft, not shipped | not used | https://github.com/notAI-tech/NudeNet | +| ProPainter | Video inpainting | S-Lab License 1.0 (non-commercial) | no | not used; never default | https://github.com/sczhou/ProPainter | +| InsightFace model zoo (SCRFD, ArcFace, inswapper) | Face detection, recognition, swap | code MIT; models non-commercial research only | no | not used; never default | https://github.com/deepinsight/insightface | +| Synthetic identity generator + swap | `synthesizeIdentity` | none found with a commercial license | — | unavailable: always 501 `unsupported` | — | + +## Methods without weights + +| Name | Purpose | License | Commercial use | Default / optional | Source | +|---|---|---|---|---|---| +| PySceneDetect (AdaptiveDetector) | Shot detection (`detectScenes`) | BSD-3-Clause | yes | default | https://github.com/Breakthrough/PySceneDetect | +| IoU tracker (ByteTrack-style, in `app/ops/faces.py`) | Face tracking | this repository | yes | default | — | +| Spectral residual saliency (Hou & Zhang 2007, in `app/ops/saliency.py`) | `saliency` | this repository | yes | default | https://doi.org/10.1109/CVPR.2007.383267 | +| FFmpeg `ebur128`, `silencedetect` | Loudness, true peak, silences (`audioFeatures`) | LGPL-2.1+ / GPL-2+ (Debian build with libx264) | yes, as a separate process | default | https://ffmpeg.org/legal.html | +| librosa | Beat and tempo tracking (`audioFeatures`) | ISC | yes | default | https://github.com/librosa/librosa | + +## Libraries + +| Name | License | Commercial use | +|---|---|---| +| FastAPI | MIT | yes | +| Uvicorn | BSD-3-Clause | yes | +| Pydantic | MIT | yes | +| NumPy | BSD-3-Clause | yes | +| OpenCV (opencv-python-headless) | Apache-2.0 (wheel bundles FFmpeg under LGPL) | yes | +| soundfile | BSD-3-Clause | yes | +| PyTorch, torchvision | BSD-3-Clause (bundles permissive third-party code: Apache-2.0, MIT, BSL-1.0) | yes | +| ONNX Runtime | MIT | yes | +| faster-whisper | MIT | yes | +| CTranslate2 | MIT | yes | +| Hugging Face tokenizers | Apache-2.0 | yes | +| SAM 2 package (`sam2`) | Apache-2.0 | yes | +| Hydra, iopath | MIT | yes | diff --git a/infra/media-worker/README.md b/infra/media-worker/README.md new file mode 100644 index 0000000..420cb80 --- /dev/null +++ b/infra/media-worker/README.md @@ -0,0 +1,74 @@ +# media-worker + +Local inference sidecar for Media Studio (Python 3.11, FastAPI). It implements +every `MediaInferenceProvider` operation (`packages/adapter-kit/src/media.ts`) +behind HTTP API v1; the TypeScript client is `packages/adapters/src/media-sidecar.ts`. + +## API + +- `GET /healthz` → `{"ok": true, "version": "..."}` (no token) +- `GET /v1/capabilities` → `MediaCapabilities` +- `POST /v1/` with the TypeScript request object (minus `signal`) +- Errors: `{"error": {"code", "message", "operation"}}` with + `bad_request`/`invalid_path` 400, `unauthorized` 401, `too_large` 413, + `unsupported` 501, `busy` 503, `internal` 500. + +Every route but `/healthz` needs `Authorization: Bearer $MEDIA_WORKER_TOKEN`. +Paths must be absolute and resolve (symlinks included) inside `MEDIA_ROOT`. +A request stops when its client disconnects (between frames or chunks) or +after `MEDIA_MAX_SECONDS`. + +## Configuration + +| Variable | Default | | +|---|---|---| +| `MEDIA_WORKER_TOKEN` | — | required, ≥ 32 chars (`MEDIA_WORKER_INSECURE_DEV=1` lifts this for tests) | +| `MEDIA_ROOT` | — | required; mount it at the same absolute path as the API's `DATA_DIR/media` | +| `MEDIA_MODELS_DIR` | `/models` | weight cache (volume) | +| `MEDIA_WORKER_HOST` / `MEDIA_WORKER_PORT` | `127.0.0.1` / `8790` | | +| `MEDIA_GPU_CONCURRENCY` | `1` | model operations running at once | +| `MEDIA_CPU_CONCURRENCY` | `2` | scenes, redact, audio, saliency running at once | +| `MEDIA_QUEUE_SIZE` | `8` | waiting requests per gate before `503 busy` | +| `MEDIA_MAX_INPUT_BYTES` | 20 GiB | larger inputs → 413 | +| `MEDIA_MAX_SECONDS` | 14400 | longer inputs → 413; longer processing is cancelled | +| `MEDIA_MODEL_DOWNLOAD` | `1` | `0` = offline; missing weights → 501 | +| `MEDIA_OPTIONAL_BACKENDS` | empty | opt-in for backends that are not clean for commercial use, e.g. `faceIdentity` | +| `MEDIA_ACCELERATOR` | build arg | `cpu`, `cuda` or `rocm`; falls back to cpu when no GPU is visible | + +## Models + +Backends, licenses, URLs, SHA256 and sizes are pinned in `app/models.py`; the +license review is in `LICENSES.md`. Weights download on first use into +`MEDIA_MODELS_DIR//`, are checked against the pinned SHA256 and size, +and are deleted on mismatch (the operation then reports 501). Warm the volume +ahead of time with `python -m app prefetch`. + +`synthesizeIdentity` always answers 501: there is no commercially licensed +generator and swap pipeline, and the request schema rejects any extra field, +so a target-face image cannot be passed. + +## Redaction geometry + +Boxes interpolate linearly between keyframes; a single keyframe is a static +region for the whole clip. `paddingPct` grows the box by that share of its +size, `featherPct` fades the effect outward over that share of the shorter +side (the region itself stays fully covered). Partial modes use the first two +landmarks as eyes and the third as the nose. Output is FFV1 for `.mkv`, +otherwise x264 CRF 12; audio is copied and metadata dropped. Any frame error is +a 500 and nothing is left at `outPath`. + +## Build and test + +```sh +docker build --build-arg ACCELERATOR=cpu -t rakazo/media-worker infra/media-worker +docker build --target test -t rakazo/media-worker:test infra/media-worker +docker run --rm rakazo/media-worker:test # offline pytest, fake model backends +``` + +Without Docker: `pip install -r requirements.txt -r requirements-dev.txt` and +`python -m pytest` in this directory. Operations whose ML packages are missing +report `available: false`. + +CUDA uses the torch cu126 wheels and `onnxruntime-gpu`. ROCm uses the torch +rocm7.1 wheels for SAM 2; ONNX Runtime and CTranslate2 have no ROCm wheels on +PyPI, so SigLIP, LaMa and Whisper run on the CPU there. diff --git a/infra/media-worker/app/__init__.py b/infra/media-worker/app/__init__.py new file mode 100644 index 0000000..dc450d5 --- /dev/null +++ b/infra/media-worker/app/__init__.py @@ -0,0 +1,3 @@ +"""Local Media Studio inference sidecar.""" + +VERSION = "0.1.0" diff --git a/infra/media-worker/app/__main__.py b/infra/media-worker/app/__main__.py new file mode 100644 index 0000000..e2d507c --- /dev/null +++ b/infra/media-worker/app/__main__.py @@ -0,0 +1,35 @@ +"""Entry point: `python -m app` serves the API, `python -m app prefetch` warms the model volume.""" + +from __future__ import annotations + +import logging +import sys + +import uvicorn + +from .config import ConfigError, Settings +from .main import create_app +from .models import Registry + + +def main(argv: list[str]) -> int: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s") + try: + settings = Settings.from_env() + except ConfigError as error: + print(f"media-worker: {error}", file=sys.stderr) + return 2 + if argv[1:] == ["prefetch"]: + Registry(settings).prefetch() + return 0 + if argv[1:]: + print("usage: python -m app [prefetch]", file=sys.stderr) + return 2 + uvicorn.run( + create_app(settings), host=settings.host, port=settings.port, log_level="info", access_log=False + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/infra/media-worker/app/auth.py b/infra/media-worker/app/auth.py new file mode 100644 index 0000000..01b3258 --- /dev/null +++ b/infra/media-worker/app/auth.py @@ -0,0 +1,22 @@ +"""Bearer-token check for every route except /healthz.""" + +from __future__ import annotations + +import hmac + +from fastapi import Request + +from .errors import MediaError + + +def require_token(request: Request) -> None: + expected: str = request.app.state.settings.token + header = request.headers.get("authorization", "") + scheme, _, presented = header.partition(" ") + if ( + scheme.lower() != "bearer" + or not presented + or not expected + or not hmac.compare_digest(presented.strip().encode(), expected.encode()) + ): + raise MediaError("unauthorized", "missing or invalid bearer token") diff --git a/infra/media-worker/app/config.py b/infra/media-worker/app/config.py new file mode 100644 index 0000000..9cc15aa --- /dev/null +++ b/infra/media-worker/app/config.py @@ -0,0 +1,80 @@ +"""Sidecar settings, read once from the environment.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field + +MIN_TOKEN_LENGTH = 32 +GIB = 1024**3 + + +class ConfigError(Exception): + pass + + +def _int(env: dict[str, str], name: str, default: int, minimum: int = 0) -> int: + raw = env.get(name, "").strip() + if not raw: + return default + try: + value = int(raw) + except ValueError as error: + raise ConfigError(f"{name} must be an integer") from error + if value < minimum: + raise ConfigError(f"{name} must be >= {minimum}") + return value + + +@dataclass(frozen=True) +class Settings: + token: str + media_root: str + models_dir: str + host: str = "127.0.0.1" + port: int = 8790 + accelerator: str = "cpu" + gpu_concurrency: int = 1 + cpu_concurrency: int = 2 + queue_size: int = 8 + max_input_bytes: int = 20 * GIB + max_seconds: int = 4 * 3600 + allow_model_download: bool = True + # Backends that are not licensed for commercial use stay off unless named here. + optional_backends: frozenset[str] = field(default_factory=frozenset) + insecure_dev: bool = False + + @staticmethod + def from_env(env: dict[str, str] | None = None) -> Settings: + env = dict(os.environ if env is None else env) + insecure = env.get("MEDIA_WORKER_INSECURE_DEV", "") == "1" + token = env.get("MEDIA_WORKER_TOKEN", "") + if len(token) < MIN_TOKEN_LENGTH and not insecure: + raise ConfigError(f"MEDIA_WORKER_TOKEN must be at least {MIN_TOKEN_LENGTH} characters") + media_root = env.get("MEDIA_ROOT", "").strip() + if not media_root or not os.path.isabs(media_root): + raise ConfigError("MEDIA_ROOT must be an absolute path") + if not os.path.isdir(media_root): + raise ConfigError("MEDIA_ROOT does not exist") + accelerator = env.get("MEDIA_ACCELERATOR", "cpu").strip() or "cpu" + if accelerator not in ("cpu", "cuda", "rocm"): + raise ConfigError("MEDIA_ACCELERATOR must be cpu, cuda or rocm") + optional = frozenset( + item.strip() for item in env.get("MEDIA_OPTIONAL_BACKENDS", "").split(",") if item.strip() + ) + return Settings( + token=token, + media_root=media_root, + models_dir=env.get("MEDIA_MODELS_DIR", "/models").strip() or "/models", + host=env.get("MEDIA_WORKER_HOST", "127.0.0.1").strip() or "127.0.0.1", + port=_int(env, "MEDIA_WORKER_PORT", 8790, 1), + accelerator=accelerator, + gpu_concurrency=_int(env, "MEDIA_GPU_CONCURRENCY", 1, 1), + cpu_concurrency=_int(env, "MEDIA_CPU_CONCURRENCY", 2, 1), + queue_size=_int(env, "MEDIA_QUEUE_SIZE", 8, 0), + max_input_bytes=_int(env, "MEDIA_MAX_INPUT_BYTES", 20 * GIB, 1), + max_seconds=_int(env, "MEDIA_MAX_SECONDS", 4 * 3600, 1), + allow_model_download=env.get("MEDIA_MODEL_DOWNLOAD", "1") != "0", + optional_backends=optional, + insecure_dev=insecure, + ) diff --git a/infra/media-worker/app/encoders.py b/infra/media-worker/app/encoders.py new file mode 100644 index 0000000..bc042bd --- /dev/null +++ b/infra/media-worker/app/encoders.py @@ -0,0 +1,41 @@ +"""Hardware video encoders that actually encode a frame on this host.""" + +from __future__ import annotations + +import functools +import os +import re +import subprocess + +from .media_io import FFMPEG + +CANDIDATE = re.compile(r"^\s*V\S*\s+((?:h264|hevc|av1)_(?:nvenc|vaapi|qsv|amf|videotoolbox|v4l2m2m))\s") +RENDER_NODE = "/dev/dri/renderD128" + + +def _works(name: str) -> bool: + argv = [FFMPEG, "-v", "error", "-nostdin"] + if name.endswith("_vaapi"): + if not os.path.exists(RENDER_NODE): + return False + argv += ["-vaapi_device", RENDER_NODE] + argv += ["-f", "lavfi", "-i", "color=c=black:s=256x256:d=0.1", "-frames:v", "1"] + if name.endswith("_vaapi"): + argv += ["-vf", "format=nv12,hwupload"] + argv += ["-c:v", name, "-f", "null", "-"] + try: + return subprocess.run(argv, capture_output=True, timeout=15, shell=False).returncode == 0 + except (OSError, subprocess.TimeoutExpired): + return False + + +@functools.cache +def hardware_encoders() -> tuple[str, ...]: + try: + listing = subprocess.run( + [FFMPEG, "-hide_banner", "-encoders"], capture_output=True, timeout=15, shell=False, text=True + ).stdout + except (OSError, subprocess.TimeoutExpired): + return () + names = [m.group(1) for line in listing.splitlines() if (m := CANDIDATE.match(line))] + return tuple(name for name in names if _works(name)) diff --git a/infra/media-worker/app/errors.py b/infra/media-worker/app/errors.py new file mode 100644 index 0000000..ede69d0 --- /dev/null +++ b/infra/media-worker/app/errors.py @@ -0,0 +1,62 @@ +"""Error envelope shared by every route: {"error": {"code", "message", "operation"}}.""" + +from __future__ import annotations + +from fastapi.responses import JSONResponse + +STATUS_BY_CODE = { + "bad_request": 400, + "invalid_path": 400, + "unauthorized": 401, + "too_large": 413, + "unsupported": 501, + "busy": 503, + "internal": 500, +} + + +class MediaError(Exception): + def __init__(self, code: str, message: str, operation: str | None = None) -> None: + super().__init__(message) + self.code = code + self.message = message + self.operation = operation + + @property + def status(self) -> int: + return STATUS_BY_CODE[self.code] + + +class Unsupported(MediaError): + def __init__(self, reason: str, operation: str | None = None) -> None: + super().__init__("unsupported", reason, operation) + + +class InvalidPath(MediaError): + def __init__(self, message: str = "path is outside MEDIA_ROOT") -> None: + super().__init__("invalid_path", message) + + +class BadRequest(MediaError): + def __init__(self, message: str) -> None: + super().__init__("bad_request", message) + + +class Busy(MediaError): + def __init__(self) -> None: + super().__init__("busy", "inference queue is full") + + +class Cancelled(Exception): + """The client went away or the request ran past MEDIA_MAX_SECONDS.""" + + def __init__(self, reason: str = "client disconnected") -> None: + super().__init__(reason) + self.reason = reason + + +def error_response(code: str, message: str, operation: str | None, status: int | None = None) -> JSONResponse: + return JSONResponse( + status_code=status or STATUS_BY_CODE[code], + content={"error": {"code": code, "message": message, "operation": operation}}, + ) diff --git a/infra/media-worker/app/main.py b/infra/media-worker/app/main.py new file mode 100644 index 0000000..5341fc5 --- /dev/null +++ b/infra/media-worker/app/main.py @@ -0,0 +1,153 @@ +"""HTTP API v1 of the Media Studio sidecar. + +GET /healthz is open; everything else needs the bearer token. POST +/v1/ takes the TypeScript request object (minus `signal`) and +answers with the TypeScript result object. +""" + +from __future__ import annotations + +import logging +import threading +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import anyio +from fastapi import Depends, FastAPI, Request +from fastapi.exceptions import RequestValidationError +from fastapi.responses import JSONResponse, Response +from pydantic import BaseModel, ValidationError +from starlette.exceptions import HTTPException as StarletteHTTPException + +from . import VERSION, encoders, schemas +from .auth import require_token +from .config import Settings +from .errors import Cancelled, MediaError, error_response +from .models import OPERATIONS, Registry +from .ops import audio, embeddings, faces, inpaint, redact, saliency, scenes, segment, synthesize, transcribe +from .ops.common import Ctx, jsonable +from .ops.devices import effective_accelerator +from .runtime import Gate, run_job + +log = logging.getLogger("media-worker") + +MAX_BODY_BYTES = 64 * 1024 * 1024 +# Status for a request whose client went away; nobody reads it. +CLIENT_CLOSED = 499 + + +@dataclass(frozen=True) +class Route: + schema: type[BaseModel] + run: Callable[[Ctx, Any, threading.Event], dict[str, Any]] + gpu: bool + + +ROUTES: dict[str, Route] = { + "detectScenes": Route(schemas.PathRequest, scenes.run, False), + "transcribe": Route(schemas.TranscribeRequest, transcribe.run, True), + "detectFaces": Route(schemas.DetectFacesRequest, faces.detect_run, True), + "trackFaces": Route(schemas.TrackFacesRequest, faces.track_run, True), + "faceLandmarks": Route(schemas.TimesRequest, faces.landmarks_run, True), + "embedFrames": Route(schemas.TimesRequest, embeddings.embed_run, True), + "classifyContent": Route(schemas.ClassifyContentRequest, embeddings.classify_run, True), + "segmentPrompted": Route(schemas.SegmentPromptedRequest, segment.run, True), + "inpaintVideo": Route(schemas.InpaintVideoRequest, inpaint.run, True), + "redact": Route(schemas.RedactRequest, redact.run, False), + "synthesizeIdentity": Route(schemas.SynthesizeIdentityRequest, synthesize.run, True), + "audioFeatures": Route(schemas.PathRequest, audio.run, False), + "saliency": Route(schemas.TimesRequest, saliency.run, False), +} +assert set(ROUTES) == set(OPERATIONS) + + +def _operation_of(request: Request) -> str | None: + name = request.url.path.rsplit("/", 1)[-1] + return name if name in ROUTES else None + + +def _validation_message(error: ValidationError) -> str: + first = error.errors()[0] if error.errors() else {} + where = ".".join(str(part) for part in first.get("loc", ())) + return ( + f"{where}: {first.get('msg', 'invalid request')}" + if where + else str(first.get("msg", "invalid request")) + ) + + +def create_app(settings: Settings | None = None, registry: Registry | None = None) -> FastAPI: + settings = settings or Settings.from_env() + registry = registry or Registry(settings) + app = FastAPI(title="media-worker", version=VERSION, docs_url=None, redoc_url=None, openapi_url=None) + app.state.settings = settings + app.state.ctx = Ctx(settings, registry) + app.state.gpu_gate = Gate(settings.gpu_concurrency, settings.queue_size) + app.state.cpu_gate = Gate(settings.cpu_concurrency, settings.queue_size) + + @app.exception_handler(MediaError) + async def media_error(request: Request, error: MediaError) -> JSONResponse: + return error_response(error.code, error.message, error.operation or _operation_of(request)) + + @app.exception_handler(RequestValidationError) + async def validation_error(request: Request, error: RequestValidationError) -> JSONResponse: + return error_response("bad_request", "invalid request", _operation_of(request)) + + @app.exception_handler(StarletteHTTPException) + async def http_error(request: Request, error: StarletteHTTPException) -> JSONResponse: + code = "bad_request" if error.status_code < 500 else "internal" + return error_response(code, str(error.detail), _operation_of(request), status=error.status_code) + + @app.exception_handler(Exception) + async def unexpected(request: Request, error: Exception) -> JSONResponse: + log.exception("unhandled error") + return error_response("internal", "internal error", _operation_of(request)) + + @app.get("/healthz") + async def healthz() -> dict[str, Any]: + return {"ok": True, "version": VERSION} + + @app.get("/v1/capabilities", dependencies=[Depends(require_token)]) + async def capabilities() -> dict[str, Any]: + hardware = await anyio.to_thread.run_sync(encoders.hardware_encoders) + return { + "configured": True, + "operations": registry.operation_capabilities(), + "accelerator": effective_accelerator(settings), + "hardwareEncoders": list(hardware), + } + + @app.post("/v1/{operation}", dependencies=[Depends(require_token)]) + async def operate(operation: str, request: Request) -> Response: + route = ROUTES.get(operation) + if route is None: + return error_response("bad_request", "unknown operation", None, status=404) + raw = bytearray() + async for chunk in request.stream(): + raw.extend(chunk) + if len(raw) > MAX_BODY_BYTES: + raise MediaError("too_large", "request body too large", operation) + try: + body = route.schema.model_validate_json(bytes(raw)) + except ValidationError as error: + raise MediaError("bad_request", _validation_message(error), operation) from error + ctx: Ctx = request.app.state.ctx + gate: Gate = request.app.state.gpu_gate if route.gpu else request.app.state.cpu_gate + try: + result = await run_job( + request, gate, lambda cancel: route.run(ctx, body, cancel), settings.max_seconds + ) + except MediaError as error: + error.operation = operation + raise + except Cancelled as error: + if error.reason == "client disconnected": + return Response(status_code=CLIENT_CLOSED) + raise MediaError("internal", error.reason, operation) from error + except Exception as error: + log.exception("%s failed", operation) + raise MediaError("internal", "internal error", operation) from error + return JSONResponse(jsonable(result)) + + return app diff --git a/infra/media-worker/app/media_io.py b/infra/media-worker/app/media_io.py new file mode 100644 index 0000000..7d3c9a7 --- /dev/null +++ b/infra/media-worker/app/media_io.py @@ -0,0 +1,435 @@ +"""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) diff --git a/infra/media-worker/app/models.py b/infra/media-worker/app/models.py new file mode 100644 index 0000000..3ad8e0c --- /dev/null +++ b/infra/media-worker/app/models.py @@ -0,0 +1,506 @@ +"""Model registry: which backend serves each operation, its license, and its files. + +Weights are never part of the image or the repo. They download lazily on first +use into MEDIA_MODELS_DIR// and are verified against the pinned +SHA256 and size before anything loads them. A mismatch deletes the file and +the operation reports `unsupported` with the reason. + +Default backends are commercially usable. Anything else must be listed in +MEDIA_OPTIONAL_BACKENDS by the operator before it runs. +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import logging +import os +import threading +import urllib.request +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +from .errors import Unsupported + +log = logging.getLogger("media-worker.models") + +OPERATIONS = ( + "detectScenes", + "transcribe", + "detectFaces", + "trackFaces", + "faceLandmarks", + "embedFrames", + "classifyContent", + "segmentPrompted", + "inpaintVideo", + "redact", + "synthesizeIdentity", + "audioFeatures", + "saliency", +) + +HF = "https://huggingface.co" +YUNET_REV = "3cc26e7f1014a5ee5d74a42acee58bafc9d0a310" +SFACE_REV = "3d7082438a6e4551e840c9b2bb60b71e8da4b524" +WHISPER_REV = "536b0662742c02347bc0e980a01041f333bce120" +SIGLIP_REV = "4649052661e53c7000355844105f8a1792088239" +LAMA_REV = "c3c0c9e468934d62e79c329e35d82dd09ff8c444" +SAM2_REV = "de431c4043854a71d8101e17995dfe596bf101a5" + + +@dataclass(frozen=True) +class ModelFile: + name: str + url: str + sha256: str + sizeBytes: int + + +@dataclass(frozen=True) +class BackendSpec: + """One way to serve an operation. `id` doubles as the model cache directory.""" + + id: str + model: str | None + license: str | None + commercialUse: bool | None + experimental: bool + source: str | None + files: tuple[ModelFile, ...] = () + # Python modules that must be importable. + requires: tuple[str, ...] = () + # Set when the backend can never run (e.g. no licensed model exists). + unavailable_reason: str | None = None + + +@dataclass +class Backend: + spec: BackendSpec + # Builds the runtime object from {file name: verified local path} and the settings. + load: Callable[[dict[str, str], Any], Any] | None = None + + +def _yunet_file() -> ModelFile: + return ModelFile( + "face_detection_yunet_2023mar.onnx", + f"{HF}/opencv/face_detection_yunet/resolve/{YUNET_REV}/face_detection_yunet_2023mar.onnx", + "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4", + 232589, + ) + + +def _siglip(name: str, sha256: str, size: int) -> ModelFile: + path = name if name == "tokenizer.json" else f"onnx/{name}" + return ModelFile(name, f"{HF}/Xenova/siglip-base-patch16-224/resolve/{SIGLIP_REV}/{path}", sha256, size) + + +def _whisper(name: str, sha256: str, size: int) -> ModelFile: + return ModelFile(name, f"{HF}/Systran/faster-whisper-small/resolve/{WHISPER_REV}/{name}", sha256, size) + + +def _lazy(module: str, attr: str) -> Callable[[dict[str, str], Any], Any]: + def load(files: dict[str, str], settings: Any) -> Any: + mod = importlib.import_module(f"app.ops.{module}") + return getattr(mod, attr)(files, settings) + + return load + + +YUNET = BackendSpec( + id="yunet-2023mar", + model="opencv-yunet-2023mar", + license="MIT", + commercialUse=True, + experimental=False, + source="https://github.com/opencv/opencv_zoo/tree/main/models/face_detection_yunet", + files=(_yunet_file(),), + requires=("cv2",), +) + +SIGLIP_SOURCE = "https://huggingface.co/google/siglip-base-patch16-224" +SIGLIP_VISION = _siglip( + "vision_model_quantized.onnx", + "ef14a954f3d57e1806666432bd9785004c1dc27100aa260eee0cb0f10a5de058", + 99499129, +) +SIGLIP_FULL = _siglip( + "model_quantized.onnx", + "9171eb00c38b9ec82f924877356d008b79e3285dbac7cd10965827bee30c9a99", + 210977441, +) +SIGLIP_TOKENIZER = _siglip( + "tokenizer.json", + "4a17c975210be5ab4c36b47d8dae4eefb866dbfb1e676e394aad85dc30a3ae08", + 2398744, +) + + +def default_backends() -> dict[str, Backend]: + return { + "detectScenes": Backend( + BackendSpec( + id="pyscenedetect-adaptive", + model="pyscenedetect-adaptive", + license="BSD-3-Clause", + commercialUse=True, + experimental=False, + source="https://github.com/Breakthrough/PySceneDetect", + requires=("scenedetect", "cv2"), + ), + _lazy("scenes", "load_scene_detector"), + ), + "transcribe": Backend( + BackendSpec( + id="faster-whisper-small", + model="faster-whisper-small", + license="MIT", + commercialUse=True, + experimental=False, + source="https://huggingface.co/Systran/faster-whisper-small", + files=( + _whisper( + "config.json", + "b55496ac7940a7ae47d2c01eab40edfd8701feec1229d9cce3b40014383fb828", + 2370, + ), + _whisper( + "model.bin", + "3e305921506d8872816023e4c273e75d2419fb89b24da97b4fe7bce14170d671", + 483546902, + ), + _whisper( + "tokenizer.json", + "fb7b63191e9bb045082c79fd742a3106a12c99513ab30df4a0d47fa6cb6fd0ab", + 2203239, + ), + _whisper( + "vocabulary.txt", + "34ce3fe1c5041027b3f8d42912270993f986dbc4bb34cf27f951e34a1e453913", + 459861, + ), + ), + requires=("faster_whisper",), + ), + _lazy("transcribe", "load_transcriber"), + ), + "detectFaces": Backend(YUNET, _lazy("faces", "load_face_detector")), + "trackFaces": Backend( + BackendSpec( + id=YUNET.id, + model="opencv-yunet-2023mar+iou-tracker", + license="MIT", + commercialUse=True, + experimental=False, + source=YUNET.source, + files=YUNET.files, + requires=YUNET.requires, + ), + _lazy("faces", "load_face_detector"), + ), + "faceLandmarks": Backend( + BackendSpec( + id=YUNET.id, + model="opencv-yunet-2023mar-5pt", + license="MIT", + commercialUse=True, + experimental=False, + source=YUNET.source, + files=YUNET.files, + requires=YUNET.requires, + ), + _lazy("faces", "load_face_detector"), + ), + "embedFrames": Backend( + BackendSpec( + id="siglip-base-patch16-224", + model="siglip-base-patch16-224", + license="Apache-2.0", + commercialUse=True, + experimental=False, + source=SIGLIP_SOURCE, + files=(SIGLIP_VISION,), + requires=("onnxruntime",), + ), + _lazy("embeddings", "load_frame_embedder"), + ), + "classifyContent": Backend( + BackendSpec( + id="siglip-base-patch16-224", + model="siglip-base-patch16-224-zero-shot", + license="Apache-2.0", + commercialUse=True, + # Zero-shot prompts, not a trained moderation head. + experimental=True, + source=SIGLIP_SOURCE, + files=(SIGLIP_FULL, SIGLIP_TOKENIZER), + requires=("onnxruntime", "tokenizers"), + ), + _lazy("embeddings", "load_content_classifier"), + ), + "segmentPrompted": Backend( + BackendSpec( + id="sam2.1-hiera-tiny", + model="sam2.1-hiera-tiny", + license="Apache-2.0", + commercialUse=True, + experimental=True, + source="https://github.com/facebookresearch/sam2", + files=( + ModelFile( + "sam2.1_hiera_tiny.pt", + f"{HF}/facebook/sam2.1-hiera-tiny/resolve/{SAM2_REV}/sam2.1_hiera_tiny.pt", + "7402e0d864fa82708a20fbd15bc84245c2f26dff0eb43a4b5b93452deb34be69", + 156008466, + ), + ), + requires=("torch", "sam2"), + ), + _lazy("segment", "load_segmenter"), + ), + "inpaintVideo": Backend( + BackendSpec( + id="lama-big-onnx", + model="big-lama+dis-flow", + license="Apache-2.0", + commercialUse=True, + experimental=True, + source="https://huggingface.co/Carve/LaMa-ONNX", + files=( + ModelFile( + "lama_fp32.onnx", + f"{HF}/Carve/LaMa-ONNX/resolve/{LAMA_REV}/lama_fp32.onnx", + "1faef5301d78db7dda502fe59966957ec4b79dd64e16f03ed96913c7a4eb68d6", + 208044816, + ), + ), + requires=("onnxruntime", "cv2"), + ), + _lazy("inpaint", "load_inpainter"), + ), + "redact": Backend( + BackendSpec( + id="opencv-redact", + model="opencv-redact", + license="Apache-2.0", + commercialUse=True, + experimental=False, + source="https://github.com/opencv/opencv", + requires=("cv2",), + ), + ), + "synthesizeIdentity": Backend( + BackendSpec( + id="synthetic-identity", + model=None, + license=None, + commercialUse=None, + experimental=True, + source=None, + unavailable_reason="no commercially licensed model", + ), + ), + "audioFeatures": Backend( + BackendSpec( + id="ffmpeg-ebur128-librosa", + model="ffmpeg-ebur128+librosa-beat", + license="ISC", + commercialUse=True, + experimental=False, + source="https://github.com/librosa/librosa", + requires=("librosa",), + ), + ), + "saliency": Backend( + BackendSpec( + id="spectral-residual", + model="spectral-residual", + license="Apache-2.0", + commercialUse=True, + experimental=False, + source="https://doi.org/10.1109/CVPR.2007.383267", + requires=("cv2",), + ), + ), + } + + +# Opt-in backends that are not clean enough to be defaults. +SFACE = Backend( + BackendSpec( + id="sface-2021dec", + model="opencv-sface-2021dec", + license="Apache-2.0 (weights); training data provenance unverified", + commercialUse=None, + experimental=True, + source="https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface", + files=( + ModelFile( + "face_recognition_sface_2021dec.onnx", + f"{HF}/opencv/face_recognition_sface/resolve/{SFACE_REV}/face_recognition_sface_2021dec.onnx", + "0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79", + 38696353, + ), + ), + requires=("cv2",), + ), + _lazy("faces", "load_identity_embedder"), +) +OPTIONAL_BACKENDS = {"faceIdentity": SFACE} + + +def _sha256(path: str) -> str: + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _module_available(name: str) -> bool: + try: + return importlib.util.find_spec(name) is not None + except (ImportError, ValueError): + return False + + +@dataclass +class Registry: + settings: Any + backends: dict[str, Backend] = field(default_factory=default_backends) + optional: dict[str, Backend] = field(default_factory=lambda: dict(OPTIONAL_BACKENDS)) + _runtimes: dict[str, Any] = field(default_factory=dict) + _verified: set[str] = field(default_factory=set) + _failed: dict[str, str] = field(default_factory=dict) + _lock: threading.Lock = field(default_factory=threading.Lock) + _file_locks: dict[str, threading.Lock] = field(default_factory=dict) + + # -- availability ------------------------------------------------------- + + def _license_block(self, spec: BackendSpec, key: str) -> str | None: + if spec.commercialUse is True: + return None + if key in self.settings.optional_backends or spec.id in self.settings.optional_backends: + return None + return "backend is not licensed for commercial use and not enabled in MEDIA_OPTIONAL_BACKENDS" + + def unavailable_reason(self, backend: Backend, key: str) -> str | None: + spec = backend.spec + if spec.unavailable_reason: + return spec.unavailable_reason + blocked = self._license_block(spec, key) + if blocked: + return blocked + missing = [name for name in spec.requires if not _module_available(name)] + if missing: + return f"python module not installed: {', '.join(missing)}" + if spec.id in self._failed: + return self._failed[spec.id] + if not self.settings.allow_model_download: + absent = [f.name for f in spec.files if not os.path.isfile(self._file_path(spec, f))] + if absent: + return "model files not installed and MEDIA_MODEL_DOWNLOAD=0" + return None + + def backend(self, operation: str) -> Backend: + return self.backends[operation] + + def optional_enabled(self, key: str) -> bool: + backend = self.optional.get(key) + return backend is not None and self.unavailable_reason(backend, key) is None + + # -- runtime ------------------------------------------------------------- + + def runtime(self, operation: str) -> Any: + """Return the loaded runtime for an operation, downloading weights on first use.""" + return self._runtime(self.backends[operation], operation) + + def optional_runtime(self, key: str) -> Any: + return self._runtime(self.optional[key], key) + + def _runtime(self, backend: Backend, key: str) -> Any: + reason = self.unavailable_reason(backend, key) + if reason: + raise Unsupported(reason) + cache_key = f"{backend.spec.id}:{key}" + with self._lock: + if cache_key in self._runtimes: + return self._runtimes[cache_key] + files = {f.name: self.ensure_file(backend.spec, f) for f in backend.spec.files} + runtime = backend.load(files, self.settings) if backend.load else None + with self._lock: + self._runtimes.setdefault(cache_key, runtime) + return self._runtimes[cache_key] + + # -- files --------------------------------------------------------------- + + def _file_path(self, spec: BackendSpec, model_file: ModelFile) -> str: + return os.path.join(self.settings.models_dir, spec.id, model_file.name) + + def _file_lock(self, path: str) -> threading.Lock: + with self._lock: + return self._file_locks.setdefault(path, threading.Lock()) + + def ensure_file(self, spec: BackendSpec, model_file: ModelFile) -> str: + path = self._file_path(spec, model_file) + with self._file_lock(path): + if path in self._verified: + return path + if not os.path.isfile(path): + if not self.settings.allow_model_download: + raise Unsupported(f"model file {model_file.name} is not installed") + self._download(model_file, path) + ok = os.path.getsize(path) == model_file.sizeBytes and _sha256(path) == model_file.sha256 + if not ok: + os.unlink(path) + reason = f"model file {model_file.name} failed SHA256 verification" + self._failed[spec.id] = reason + raise Unsupported(reason) + self._verified.add(path) + return path + + def _download(self, model_file: ModelFile, path: str) -> None: + os.makedirs(os.path.dirname(path), exist_ok=True) + partial = f"{path}.part" + log.info("downloading %s", model_file.url) + try: + with ( + urllib.request.urlopen(model_file.url, timeout=60) as response, + open(partial, "wb") as handle, + ): + written = 0 + for chunk in iter(lambda: response.read(1 << 20), b""): + written += len(chunk) + if written > model_file.sizeBytes: + raise ValueError("download larger than pinned size") + handle.write(chunk) + os.replace(partial, path) + except Exception as error: + if os.path.exists(partial): + os.unlink(partial) + raise Unsupported(f"model file {model_file.name} could not be downloaded: {error}") from error + + # -- capabilities ---------------------------------------------------------- + + def operation_capabilities(self) -> dict[str, dict[str, Any]]: + result: dict[str, dict[str, Any]] = {} + for operation in OPERATIONS: + backend = self.backends[operation] + spec = backend.spec + result[operation] = { + "available": self.unavailable_reason(backend, operation) is None, + "experimental": spec.experimental, + "model": spec.model, + "license": spec.license, + "commercialUse": spec.commercialUse, + } + return result + + def prefetch(self) -> None: + """Download and verify every default backend's files (operator warm-up).""" + for operation in OPERATIONS: + backend = self.backends[operation] + if self.unavailable_reason(backend, operation) is None: + for model_file in backend.spec.files: + self.ensure_file(backend.spec, model_file) diff --git a/infra/media-worker/app/ops/__init__.py b/infra/media-worker/app/ops/__init__.py new file mode 100644 index 0000000..64d5a30 --- /dev/null +++ b/infra/media-worker/app/ops/__init__.py @@ -0,0 +1 @@ +"""One module per operation family. Each exposes `run(ctx, body, cancel)`.""" diff --git a/infra/media-worker/app/ops/audio.py b/infra/media-worker/app/ops/audio.py new file mode 100644 index 0000000..b79d917 --- /dev/null +++ b/infra/media-worker/app/ops/audio.py @@ -0,0 +1,163 @@ +"""audioFeatures without ML: ffmpeg ebur128 + silencedetect, librosa beat tracking. + +Audio is decoded once to mono 22.05 kHz float32 and processed in five-minute +chunks, so memory stays bounded and the request can stop between chunks. +""" + +from __future__ import annotations + +import re +import subprocess +import threading +from typing import Any + +import numpy as np + +from ..errors import Cancelled, MediaError +from ..media_io import FFMPEG, parse_ebur128, run_cancellable +from ..schemas import PathRequest +from .common import Ctx, open_input + +SAMPLE_RATE = 22050 +CHUNK_SECONDS = 300 +ENERGY_HOP_MS = 250 +SILENCE_DB = -40 +SILENCE_MIN_S = 0.5 + + +def loudness_and_silence( + path: str, cancel: threading.Event +) -> tuple[float | None, float | None, list[dict[str, int]]]: + result = run_cancellable( + [ + FFMPEG, + "-hide_banner", + "-nostats", + "-nostdin", + "-i", + path, + "-map", + "0:a:0", + "-vn", + "-af", + f"ebur128=peak=true,silencedetect=noise={SILENCE_DB}dB:d={SILENCE_MIN_S}", + "-f", + "null", + "-", + ], + cancel, + ) + stderr = result.stderr.decode("utf-8", "replace") + if result.returncode != 0: + raise MediaError("internal", "loudness measurement failed") + lufs, peak = parse_ebur128(stderr) + silences: list[dict[str, int]] = [] + start: float | None = None + for line in stderr.splitlines(): + if (m := re.search(r"silence_start:\s*(-?[\d.]+)", line)) is not None: + start = max(float(m.group(1)), 0.0) + elif (m := re.search(r"silence_end:\s*([\d.]+)", line)) is not None and start is not None: + silences.append( + {"startMs": int(round(start * 1000)), "endMs": int(round(float(m.group(1)) * 1000))} + ) + start = None + if start is not None: + silences.append({"startMs": int(round(start * 1000)), "endMs": -1}) + return lufs, peak, silences + + +def decode_chunks(path: str, cancel: threading.Event) -> Any: + argv = [ + FFMPEG, + "-v", + "error", + "-nostdin", + "-i", + path, + "-map", + "0:a:0", + "-vn", + "-ac", + "1", + "-ar", + str(SAMPLE_RATE), + "-f", + "f32le", + "pipe:1", + ] + proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, shell=False) + chunk_bytes = SAMPLE_RATE * CHUNK_SECONDS * 4 + assert proc.stdout is not None + try: + while True: + if cancel.is_set(): + raise Cancelled() + buf = proc.stdout.read(chunk_bytes) + if not buf: + break + usable = len(buf) - len(buf) % 4 + yield np.frombuffer(buf[:usable], dtype=" dict[str, Any]: + import librosa # deferred: heavy import (numba), only needed here + + real, info = open_input(ctx, body.path, cancel, need_video=False) + ctx.registry.runtime("audioFeatures") + if not info.has_audio: + return { + "integratedLufs": None, + "truePeakDb": None, + "tempoBpm": None, + "beatsMs": [], + "energy": [], + "silences": [], + } + lufs, peak, silences = loudness_and_silence(real, cancel) + + hop = int(SAMPLE_RATE * ENERGY_HOP_MS / 1000) + beats: list[int] = [] + tempos: list[tuple[float, float]] = [] + rms_values: list[float] = [] + offset_s = 0.0 + for samples in decode_chunks(real, cancel): + length_s = samples.size / SAMPLE_RATE + if samples.size >= SAMPLE_RATE * 2 and float(np.abs(samples).max()) > 1e-4: + tempo, frames = librosa.beat.beat_track(y=samples, sr=SAMPLE_RATE, units="frames") + times = librosa.frames_to_time(frames, sr=SAMPLE_RATE) + beats.extend(int(round((offset_s + t) * 1000)) for t in times) + value = float(np.atleast_1d(tempo)[0]) + if value > 0: + tempos.append((value, length_s)) + for i in range(0, samples.size, hop): + window = samples[i : i + hop] + rms_values.append(float(np.sqrt(np.mean(window.astype(np.float64) ** 2))) if window.size else 0.0) + offset_s += length_s + + duration_ms = int(round(offset_s * 1000)) + for silence in silences: + if silence["endMs"] < 0: + silence["endMs"] = duration_ms + peak_rms = max(rms_values) if rms_values else 0.0 + energy = [ + {"tMs": i * ENERGY_HOP_MS, "value": (v / peak_rms) if peak_rms > 0 else 0.0} + for i, v in enumerate(rms_values) + ] + tempo_bpm = None + if tempos: + weights = np.array([w for _, w in tempos]) + tempo_bpm = float(np.average([t for t, _ in tempos], weights=weights)) + return { + "integratedLufs": lufs, + "truePeakDb": peak, + "tempoBpm": tempo_bpm, + "beatsMs": sorted(set(beats)), + "energy": energy, + "silences": silences, + } diff --git a/infra/media-worker/app/ops/common.py b/infra/media-worker/app/ops/common.py new file mode 100644 index 0000000..84fe793 --- /dev/null +++ b/infra/media-worker/app/ops/common.py @@ -0,0 +1,62 @@ +"""Helpers shared by the operation modules.""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from .. import media_io, paths +from ..config import Settings +from ..errors import BadRequest +from ..media_io import VideoInfo +from ..models import Registry + + +@dataclass +class Ctx: + settings: Settings + registry: Registry + + +def open_input( + ctx: Ctx, path: str, cancel: threading.Event, need_video: bool = True +) -> tuple[str, VideoInfo]: + """Resolve, size-check, probe, and duration-check an input file.""" + real = paths.input_file(ctx.settings.media_root, path, ctx.settings.max_input_bytes) + info = media_io.probe(real, cancel) + media_io.check_duration(info, ctx.settings.max_seconds) + if need_video and not info.has_video: + raise BadRequest("input has no video stream") + return real, info + + +def clamp01(value: float) -> float: + return float(min(max(value, 0.0), 1.0)) + + +def norm_box(x: float, y: float, w: float, h: float, width: int, height: int) -> dict[str, float]: + """Pixel box -> normalized box clipped to the frame.""" + x0, y0 = clamp01(x / width), clamp01(y / height) + x1, y1 = clamp01((x + w) / width), clamp01((y + h) / height) + return {"x": x0, "y": y0, "w": max(x1 - x0, 0.0), "h": max(y1 - y0, 0.0)} + + +def norm_point(x: float, y: float, width: int, height: int) -> dict[str, float]: + return {"x": clamp01(x / width), "y": clamp01(y / height)} + + +def unique_sorted(times: list[int]) -> list[int]: + return sorted(set(int(t) for t in times)) + + +def jsonable(value: Any) -> Any: + if isinstance(value, dict): + return {k: jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [jsonable(v) for v in value] + if isinstance(value, np.generic): + return value.item() + return value diff --git a/infra/media-worker/app/ops/devices.py b/infra/media-worker/app/ops/devices.py new file mode 100644 index 0000000..69054a9 --- /dev/null +++ b/infra/media-worker/app/ops/devices.py @@ -0,0 +1,61 @@ +"""Accelerator detection. ML libraries are imported only when installed.""" + +from __future__ import annotations + +import functools +import importlib.util +from typing import Any + + +def _installed(name: str) -> bool: + try: + return importlib.util.find_spec(name) is not None + except (ImportError, ValueError): + return False + + +@functools.cache +def torch_gpu_available() -> bool: + if not _installed("torch"): + return False + import torch + + return bool(torch.cuda.is_available()) + + +def effective_accelerator(settings: Any) -> str: + """What the image was built for, downgraded to cpu when no GPU is visible. + + ROCm builds of torch expose the GPU through the torch.cuda API. + """ + if settings.accelerator in ("cuda", "rocm") and torch_gpu_available(): + return settings.accelerator + return "cpu" + + +def torch_device(settings: Any) -> str: + return "cuda" if effective_accelerator(settings) != "cpu" else "cpu" + + +def ctranslate2_device(settings: Any) -> str: + # CTranslate2 ships CUDA builds only; ROCm hosts transcribe on the CPU. + if settings.accelerator != "cuda": + return "cpu" + import ctranslate2 + + return "cuda" if ctranslate2.get_cuda_device_count() > 0 else "cpu" + + +def onnx_providers(settings: Any) -> list[str]: + import onnxruntime + + available = set(onnxruntime.get_available_providers()) + providers: list[str] = [] + if settings.accelerator == "cuda" and "CUDAExecutionProvider" in available: + providers.append("CUDAExecutionProvider") + if settings.accelerator == "rocm": + for name in ("MIGraphXExecutionProvider", "ROCMExecutionProvider"): + if name in available: + providers.append(name) + providers.append("CPUExecutionProvider") + return providers diff --git a/infra/media-worker/app/ops/embeddings.py b/infra/media-worker/app/ops/embeddings.py new file mode 100644 index 0000000..60f5124 --- /dev/null +++ b/infra/media-worker/app/ops/embeddings.py @@ -0,0 +1,147 @@ +"""embedFrames and classifyContent with SigLIP (Apache-2.0) on ONNX Runtime. + +classifyContent is zero-shot: each category is a set of text prompts and the +score is the highest sigmoid probability among them. It is reported as +experimental; the gate thresholds decide what the scores mean. +""" + +from __future__ import annotations + +import threading +from typing import Any + +import cv2 +import numpy as np + +from .. import media_io +from ..schemas import ClassifyContentRequest, TimesRequest +from .common import Ctx, clamp01, open_input, unique_sorted +from .devices import onnx_providers + +SIGLIP_SIZE = 224 +SIGLIP_TEXT_LENGTH = 64 +SIGLIP_PAD_ID = 1 +BATCH = 16 + +CATEGORY_PROMPTS: dict[str, tuple[str, ...]] = { + "explicit_nudity": ("a photo of a naked person", "a photo containing explicit nudity"), + "suggestive": ("a photo of a person in lingerie", "a sexually suggestive photo"), + "violence": ("a photo of a violent fight", "a photo showing blood and injuries"), + "minors_risk": ("a photo of a child", "a photo of a young teenager"), + "self_harm": ("a photo of self-harm", "a photo of cuts on a wrist"), + "drugs": ("a photo of illegal drugs", "a photo of drug paraphernalia"), + "weapons": ("a photo of a gun", "a photo of a person holding a knife as a weapon"), +} + + +def preprocess(frames: list[np.ndarray]) -> np.ndarray: + """BGR frames -> SigLIP pixel_values (N, 3, 224, 224), normalized to [-1, 1].""" + batch = [] + for frame in frames: + rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + resized = cv2.resize(rgb, (SIGLIP_SIZE, SIGLIP_SIZE), interpolation=cv2.INTER_CUBIC) + batch.append(resized.astype(np.float32) / 127.5 - 1.0) + return np.ascontiguousarray(np.stack(batch).transpose(0, 3, 1, 2)) + + +def _session(path: str, settings: Any) -> Any: + import onnxruntime + + return onnxruntime.InferenceSession(path, providers=onnx_providers(settings)) + + +class SiglipEmbedder: + model = "siglip-base-patch16-224" + + def __init__(self, path: str, settings: Any) -> None: + self.session = _session(path, settings) + names = [o.name for o in self.session.get_outputs()] + self.output = "pooler_output" if "pooler_output" in names else names[-1] + self.dim = 768 + + def embed(self, frames: list[np.ndarray]) -> np.ndarray: + (vectors,) = self.session.run([self.output], {"pixel_values": preprocess(frames)}) + norms = np.linalg.norm(vectors, axis=1, keepdims=True) + return vectors / np.maximum(norms, 1e-12) + + +def load_frame_embedder(files: dict[str, str], settings: Any) -> SiglipEmbedder: + return SiglipEmbedder(files["vision_model_quantized.onnx"], settings) + + +class SiglipZeroShot: + model = "siglip-base-patch16-224-zero-shot" + + def __init__(self, model_path: str, tokenizer_path: str, settings: Any) -> None: + from tokenizers import Tokenizer + + tokenizer = Tokenizer.from_file(tokenizer_path) + tokenizer.enable_truncation(SIGLIP_TEXT_LENGTH) + tokenizer.enable_padding(length=SIGLIP_TEXT_LENGTH, pad_id=SIGLIP_PAD_ID) + self.labels: list[str] = [] + prompts: list[str] = [] + for category, texts in CATEGORY_PROMPTS.items(): + for text in texts: + self.labels.append(category) + prompts.append(text.lower()) + self.input_ids = np.array([e.ids for e in tokenizer.encode_batch(prompts)], dtype=np.int64) + self.session = _session(model_path, settings) + + def classify(self, frames: list[np.ndarray]) -> list[dict[str, float]]: + (logits,) = self.session.run( + ["logits_per_image"], {"input_ids": self.input_ids, "pixel_values": preprocess(frames)} + ) + probs = 1.0 / (1.0 + np.exp(-logits)) + result = [] + for row in probs: + scores: dict[str, float] = {} + for label, value in zip(self.labels, row, strict=True): + scores[label] = max(scores.get(label, 0.0), clamp01(float(value))) + result.append(scores) + return result + + +def load_content_classifier(files: dict[str, str], settings: Any) -> SiglipZeroShot: + return SiglipZeroShot(files["model_quantized.onnx"], files["tokenizer.json"], settings) + + +def _batched(frames: Any, size: int) -> Any: + times: list[int] = [] + batch: list[np.ndarray] = [] + for t_ms, frame in frames: + times.append(t_ms) + batch.append(frame) + if len(batch) == size: + yield times, batch + times, batch = [], [] + if batch: + yield times, batch + + +def embed_run(ctx: Ctx, body: TimesRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + embedder = ctx.registry.runtime("embedFrames") + vectors: list[list[float]] = [] + # Vectors follow the request order, duplicates included. + for _, batch in _batched(media_io.frames_at(real, info, list(body.timesMs), cancel, 512), BATCH): + vectors.extend([float(v) for v in row] for row in embedder.embed(batch)) + return {"model": embedder.model, "dim": int(embedder.dim), "vectors": vectors} + + +def classify_run(ctx: Ctx, body: ClassifyContentRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + classifier = ctx.registry.runtime("classifyContent") + results: dict[int, dict[str, float]] = {} + + def score(frames: Any) -> None: + for times, batch in _batched(frames, BATCH): + for t_ms, scores in zip(times, classifier.classify(batch), strict=True): + results[t_ms] = scores + + score(media_io.stream_frames(real, info, cancel, body.sampleFps, 512)) + extra = [t for t in unique_sorted(body.extraTimesMs or []) if t not in results] + score(media_io.frames_at(real, info, extra, cancel, 512)) + return { + "model": classifier.model, + "frames": [{"tMs": t, "scores": results[t]} for t in sorted(results)], + } diff --git a/infra/media-worker/app/ops/faces.py b/infra/media-worker/app/ops/faces.py new file mode 100644 index 0000000..167fae2 --- /dev/null +++ b/infra/media-worker/app/ops/faces.py @@ -0,0 +1,265 @@ +"""detectFaces, trackFaces, faceLandmarks. + +Detection is OpenCV YuNet (MIT). Tracking is a ByteTrack-style IoU tracker +implemented here: high-confidence detections associate first, low-confidence +ones only extend existing tracks, and a track ends after 500 ms without a match. +Identity embeddings (for local person clustering) come from an opt-in backend +and are null otherwise. +""" + +from __future__ import annotations + +import math +import threading +from dataclasses import dataclass, field +from typing import Any + +import cv2 +import numpy as np + +from .. import media_io +from ..schemas import DetectFacesRequest, TimesRequest, TrackFacesRequest +from .common import Ctx, clamp01, norm_box, norm_point, open_input, unique_sorted + +DETECT_MAX_SIDE = 1280 +DETECT_THRESHOLD = 0.6 +TRACK_LOW_THRESHOLD = 0.3 +TRACK_MATCH_IOU = 0.3 +TRACK_MAX_GAP_MS = 500 + + +@dataclass +class Detection: + """Pixel-space face: box (x, y, w, h), score, five landmarks (x, y).""" + + box: tuple[float, float, float, float] + score: float + landmarks: list[tuple[float, float]] + + +class YuNetDetector: + def __init__(self, model_path: str) -> None: + self.model_path = model_path + self._local = threading.local() + + def _detector(self, width: int, height: int) -> Any: + detector = getattr(self._local, "detector", None) + if detector is None: + detector = cv2.FaceDetectorYN.create( + self.model_path, "", (width, height), TRACK_LOW_THRESHOLD, 0.3, 5000 + ) + self._local.detector = detector + detector.setInputSize((width, height)) + return detector + + def detect(self, frame: np.ndarray) -> list[Detection]: + height, width = frame.shape[:2] + _, faces = self._detector(width, height).detect(frame) + result: list[Detection] = [] + for row in faces if faces is not None else []: + points = [(float(row[4 + 2 * i]), float(row[5 + 2 * i])) for i in range(5)] + result.append( + Detection( + (float(row[0]), float(row[1]), float(row[2]), float(row[3])), float(row[14]), points + ) + ) + return result + + +def load_face_detector(files: dict[str, str], _settings: Any) -> YuNetDetector: + return YuNetDetector(files["face_detection_yunet_2023mar.onnx"]) + + +class SFaceEmbedder: + def __init__(self, model_path: str) -> None: + self.model_path = model_path + self._local = threading.local() + + def embed(self, frame: np.ndarray, detection: Detection) -> np.ndarray: + recognizer = getattr(self._local, "recognizer", None) + if recognizer is None: + recognizer = cv2.FaceRecognizerSF.create(self.model_path, "") + self._local.recognizer = recognizer + row = np.array( + [*detection.box, *[c for p in detection.landmarks for c in p], detection.score], + dtype=np.float32, + ) + aligned = recognizer.alignCrop(frame, row) + return recognizer.feature(aligned).reshape(-1).astype(np.float32) + + +def load_identity_embedder(files: dict[str, str], _settings: Any) -> SFaceEmbedder: + return SFaceEmbedder(files["face_recognition_sface_2021dec.onnx"]) + + +def _observation(det: Detection, width: int, height: int) -> dict[str, Any]: + return { + "box": norm_box(*det.box, width, height), + "confidence": clamp01(det.score), + "landmarks": [norm_point(x, y, width, height) for x, y in det.landmarks], + } + + +def detect_run(ctx: Ctx, body: DetectFacesRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + detector = ctx.registry.runtime("detectFaces") + frames = ( + media_io.frames_at(real, info, unique_sorted(body.timesMs), cancel, DETECT_MAX_SIDE) + if body.timesMs is not None + else media_io.stream_frames(real, info, cancel, body.sampleFps, DETECT_MAX_SIDE) + ) + out = [] + for t_ms, frame in frames: + height, width = frame.shape[:2] + faces = [ + _observation(d, width, height) for d in detector.detect(frame) if d.score >= DETECT_THRESHOLD + ] + out.append({"tMs": t_ms, "faces": faces}) + return {"frames": out} + + +def landmarks_run(ctx: Ctx, body: TimesRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + detector = ctx.registry.runtime("faceLandmarks") + out = [] + for t_ms, frame in media_io.frames_at(real, info, unique_sorted(body.timesMs), cancel, DETECT_MAX_SIDE): + height, width = frame.shape[:2] + faces = [ + { + "box": norm_box(*d.box, width, height), + "landmarks": [norm_point(x, y, width, height) for x, y in d.landmarks], + } + for d in detector.detect(frame) + if d.score >= DETECT_THRESHOLD + ] + out.append({"tMs": t_ms, "faces": faces}) + return {"frames": out} + + +def iou(a: tuple[float, float, float, float], b: tuple[float, float, float, float]) -> float: + ax1, ay1, ax2, ay2 = a[0], a[1], a[0] + a[2], a[1] + a[3] + bx1, by1, bx2, by2 = b[0], b[1], b[0] + b[2], b[1] + b[3] + iw = max(0.0, min(ax2, bx2) - max(ax1, bx1)) + ih = max(0.0, min(ay2, by2) - max(ay1, by1)) + inter = iw * ih + union = a[2] * a[3] + b[2] * b[3] - inter + return inter / union if union > 0 else 0.0 + + +@dataclass +class Track: + track_id: str + box: tuple[float, float, float, float] + velocity: tuple[float, float] = (0.0, 0.0) + last_ms: int = 0 + keyframes: list[dict[str, Any]] = field(default_factory=list) + # Best (score, frame crop source, detection) for identity embeddings. + best: list[tuple[float, np.ndarray, Detection]] = field(default_factory=list) + + def predicted(self, t_ms: int) -> tuple[float, float, float, float]: + dt = t_ms - self.last_ms + return ( + self.box[0] + self.velocity[0] * dt, + self.box[1] + self.velocity[1] * dt, + self.box[2], + self.box[3], + ) + + +class IouTracker: + """Two-stage greedy IoU association (ByteTrack's idea without the Kalman filter).""" + + def __init__(self, max_gap_ms: int = TRACK_MAX_GAP_MS) -> None: + self.max_gap_ms = max_gap_ms + self.active: list[Track] = [] + self.finished: list[Track] = [] + self._next = 1 + + def _associate( + self, tracks: list[Track], dets: list[Detection], t_ms: int + ) -> tuple[list[tuple[Track, Detection]], list[Track], list[Detection]]: + pairs = sorted( + ( + (iou(tr.predicted(t_ms), d.box), ti, di) + for ti, tr in enumerate(tracks) + for di, d in enumerate(dets) + ), + reverse=True, + ) + used_t: set[int] = set() + used_d: set[int] = set() + matches: list[tuple[Track, Detection]] = [] + for score, ti, di in pairs: + if score < TRACK_MATCH_IOU: + break + if ti in used_t or di in used_d: + continue + used_t.add(ti) + used_d.add(di) + matches.append((tracks[ti], dets[di])) + rest_t = [t for i, t in enumerate(tracks) if i not in used_t] + rest_d = [d for i, d in enumerate(dets) if i not in used_d] + return matches, rest_t, rest_d + + def update( + self, t_ms: int, dets: list[Detection], frame: np.ndarray | None, keep_crops: bool + ) -> list[tuple[Track, Detection]]: + expired = [t for t in self.active if t_ms - t.last_ms > self.max_gap_ms] + self.finished.extend(expired) + self.active = [t for t in self.active if t not in expired] + high = [d for d in dets if d.score >= DETECT_THRESHOLD] + low = [d for d in dets if TRACK_LOW_THRESHOLD <= d.score < DETECT_THRESHOLD] + matches, rest_tracks, rest_high = self._associate(self.active, high, t_ms) + low_matches, _, _ = self._associate(rest_tracks, low, t_ms) + matches += low_matches + for track, det in matches: + dt = max(t_ms - track.last_ms, 1) + track.velocity = ((det.box[0] - track.box[0]) / dt, (det.box[1] - track.box[1]) / dt) + track.box = det.box + track.last_ms = t_ms + for det in rest_high: + track = Track(f"t{self._next}", det.box, last_ms=t_ms) + self._next += 1 + self.active.append(track) + matches.append((track, det)) + if keep_crops and frame is not None: + for track, det in matches: + if det.score >= DETECT_THRESHOLD: + track.best.append((det.score, frame, det)) + track.best.sort(key=lambda item: -item[0]) + del track.best[3:] + return matches + + def close(self) -> list[Track]: + tracks = self.finished + self.active + self.active, self.finished = [], [] + return sorted(tracks, key=lambda t: int(t.track_id[1:])) + + +def track_run(ctx: Ctx, body: TrackFacesRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + detector = ctx.registry.runtime("trackFaces") + embedder = ( + ctx.registry.optional_runtime("faceIdentity") + if ctx.registry.optional_enabled("faceIdentity") + else None + ) + tracker = IouTracker() + for t_ms, frame in media_io.stream_frames(real, info, cancel, body.sampleFps, DETECT_MAX_SIDE): + height, width = frame.shape[:2] + matches = tracker.update(t_ms, detector.detect(frame), frame, embedder is not None) + for track, det in matches: + track.keyframes.append( + {"tMs": t_ms, "box": norm_box(*det.box, width, height), "confidence": clamp01(det.score)} + ) + tracks = [] + for track in tracker.close(): + embedding = None + if embedder is not None and track.best: + vectors = np.stack([embedder.embed(frame, det) for _, frame, det in track.best]) + mean = vectors.mean(axis=0) + norm = float(np.linalg.norm(mean)) + if norm > 0 and math.isfinite(norm): + embedding = [float(v) for v in mean / norm] + tracks.append({"trackId": track.track_id, "keyframes": track.keyframes, "embedding": embedding}) + return {"sampleFps": body.sampleFps, "tracks": tracks} diff --git a/infra/media-worker/app/ops/inpaint.py b/infra/media-worker/app/ops/inpaint.py new file mode 100644 index 0000000..314085c --- /dev/null +++ b/infra/media-worker/app/ops/inpaint.py @@ -0,0 +1,163 @@ +"""inpaintVideo: LaMa (Apache-2.0, ONNX) per frame with optical-flow smoothing. + +Per frame: crop around the mask, inpaint at 512x512, paste back through a +feathered mask. When the previous frame was inpainted too, its result is +warped into the current frame with DIS optical flow and blended in, which +damps flicker. Quality is measured on the output: +- temporalStability: 1 - mean absolute change inside the mask between + consecutive frames after motion compensation (scaled so 32 levels -> 0). +- edgeQuality: 1 - excess gradient on the mask boundary over the gradient + just outside it. +- score: 0.6 * temporalStability + 0.4 * edgeQuality. +Callers fall back to cover/reframe when the score is low. +""" + +from __future__ import annotations + +import os +import threading +from typing import Any + +import cv2 +import numpy as np + +from .. import media_io, paths +from ..errors import BadRequest +from ..schemas import InpaintVideoRequest +from .common import Ctx, clamp01, open_input +from .devices import onnx_providers +from .masks import MaskIndex +from .redact import OUTPUT_EXTENSIONS + +LAMA_SIZE = 512 +TEMPORAL_BLEND = 0.35 + + +class LamaInpainter: + def __init__(self, path: str, settings: Any) -> None: + import onnxruntime + + self.session = onnxruntime.InferenceSession(path, providers=onnx_providers(settings)) + + def inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray: + """image: BGR uint8 512x512, mask: uint8 0/255 512x512 -> BGR uint8.""" + rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 + pixel = rgb.transpose(2, 0, 1)[None] + m = (mask > 127).astype(np.float32)[None, None] + (out,) = self.session.run(["output"], {"image": pixel, "mask": m}) + result = np.clip(out[0].transpose(1, 2, 0), 0, 255).astype(np.uint8) + return cv2.cvtColor(result, cv2.COLOR_RGB2BGR) + + +def load_inpainter(files: dict[str, str], settings: Any) -> LamaInpainter: + return LamaInpainter(files["lama_fp32.onnx"], settings) + + +def _crop_box(mask: np.ndarray) -> tuple[int, int, int, int]: + height, width = mask.shape + ys, xs = np.nonzero(mask) + cx, cy = (xs.min() + xs.max()) / 2, (ys.min() + ys.max()) / 2 + side = max(xs.max() - xs.min() + 1, ys.max() - ys.min() + 1) * 2.2 + side = int(min(max(side, 256), max(width, height))) + x0 = int(min(max(cx - side / 2, 0), max(width - side, 0))) + y0 = int(min(max(cy - side / 2, 0), max(height - side, 0))) + return x0, y0, min(x0 + side, width), min(y0 + side, height) + + +class Smoother: + """Temporal blending and quality bookkeeping across frames.""" + + def __init__(self) -> None: + self.flow = cv2.DISOpticalFlow_create(cv2.DISOPTICAL_FLOW_PRESET_FAST) + self.prev_gray: np.ndarray | None = None + self.prev_out: np.ndarray | None = None + self.prev_masked = False + self.stability: list[float] = [] + self.edges: list[float] = [] + + def _warp_prev(self, gray: np.ndarray) -> np.ndarray | None: + if self.prev_gray is None or self.prev_out is None: + return None + flow = self.flow.calc(gray, self.prev_gray, None) + h, w = gray.shape + grid_x, grid_y = np.meshgrid(np.arange(w, dtype=np.float32), np.arange(h, dtype=np.float32)) + return cv2.remap( + self.prev_out, + grid_x + flow[..., 0], + grid_y + flow[..., 1], + cv2.INTER_LINEAR, + borderMode=cv2.BORDER_REPLICATE, + ) + + def step(self, source: np.ndarray, inpainted: np.ndarray | None, mask: np.ndarray | None) -> np.ndarray: + gray = cv2.cvtColor(source, cv2.COLOR_BGR2GRAY) + out = source + if inpainted is not None and mask is not None: + soft = cv2.GaussianBlur(mask.astype(np.float32) / 255.0, (0, 0), 3) + soft = np.maximum(soft, mask.astype(np.float32) / 255.0)[..., None] + warped = self._warp_prev(gray) if self.prev_masked else None + fill = inpainted.astype(np.float32) + if warped is not None: + fill = fill * (1 - TEMPORAL_BLEND) + warped.astype(np.float32) * TEMPORAL_BLEND + out = np.clip(source.astype(np.float32) * (1 - soft) + fill * soft, 0, 255).astype(np.uint8) + inside = mask > 0 + if warped is not None: + # Raw LaMa output against the motion-compensated previous frame: flicker before smoothing. + diff = np.abs(inpainted.astype(np.float32) - warped.astype(np.float32))[inside].mean() + self.stability.append(clamp01(1.0 - float(diff) / 32.0)) + self.edges.append(_edge_quality(out, mask)) + self.prev_gray = gray + self.prev_out = out + self.prev_masked = mask is not None + return out + + def quality(self) -> dict[str, float]: + temporal = float(np.mean(self.stability)) if self.stability else 1.0 + edge = float(np.mean(self.edges)) if self.edges else 1.0 + return { + "temporalStability": clamp01(temporal), + "edgeQuality": clamp01(edge), + "score": clamp01(0.6 * temporal + 0.4 * edge), + } + + +def _edge_quality(frame: np.ndarray, mask: np.ndarray) -> float: + gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY).astype(np.float32) + grad = cv2.magnitude(cv2.Sobel(gray, cv2.CV_32F, 1, 0), cv2.Sobel(gray, cv2.CV_32F, 0, 1)) + kernel = np.ones((5, 5), np.uint8) + boundary = (cv2.dilate(mask, kernel) > 0) & (cv2.erode(mask, kernel) == 0) + outer = (cv2.dilate(mask, np.ones((15, 15), np.uint8)) > 0) & ~(cv2.dilate(mask, kernel) > 0) + if not boundary.any() or not outer.any(): + return 1.0 + excess = max(0.0, float(grad[boundary].mean()) - float(grad[outer].mean())) + return clamp01(1.0 - excess / (float(grad[outer].mean()) + 40.0)) + + +def run(ctx: Ctx, body: InpaintVideoRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + masks = MaskIndex(paths.input_dir(ctx.settings.media_root, body.maskDir), info.frame_ms()) + out_real = paths.output_file(ctx.settings.media_root, body.outPath) + if os.path.splitext(out_real)[1].lower() not in OUTPUT_EXTENSIONS: + raise BadRequest("outPath must end in .mkv, .mp4 or .mov") + if out_real == real: + raise BadRequest("outPath must differ from path") + inpainter = ctx.registry.runtime("inpaintVideo") + smoother = Smoother() + + def per_frame(_index: int, t_ms: int, frame: np.ndarray) -> np.ndarray: + height, width = frame.shape[:2] + mask = masks.lookup(t_ms, width, height) + if mask is None or not mask.any(): + return smoother.step(frame, None, None) + # Grow slightly so the seam lands on untouched pixels. + mask = cv2.dilate(mask, np.ones((7, 7), np.uint8)) + x0, y0, x1, y1 = _crop_box(mask) + crop = cv2.resize(frame[y0:y1, x0:x1], (LAMA_SIZE, LAMA_SIZE), interpolation=cv2.INTER_AREA) + crop_mask = cv2.resize(mask[y0:y1, x0:x1], (LAMA_SIZE, LAMA_SIZE), interpolation=cv2.INTER_NEAREST) + filled = inpainter.inpaint(crop, crop_mask) + full = frame.copy() + full[y0:y1, x0:x1] = cv2.resize(filled, (x1 - x0, y1 - y0), interpolation=cv2.INTER_CUBIC) + return smoother.step(frame, full, mask) + + media_io.transform_video(real, out_real, info, cancel, per_frame) + return {"outPath": body.outPath, "quality": smoother.quality()} diff --git a/infra/media-worker/app/ops/masks.py b/infra/media-worker/app/ops/masks.py new file mode 100644 index 0000000..f4824cf --- /dev/null +++ b/infra/media-worker/app/ops/masks.py @@ -0,0 +1,54 @@ +"""Per-frame mask directories (`.png`), as written by segmentPrompted.""" + +from __future__ import annotations + +import bisect +import os +import re + +import cv2 +import numpy as np + +from ..errors import BadRequest, MediaError + +MASK_NAME = re.compile(r"^(\d+)\.png$") + + +class MaskIndex: + def __init__(self, directory: str, frame_ms: float) -> None: + entries = [] + for name in os.listdir(directory): + match = MASK_NAME.match(name) + if match: + entries.append((int(match.group(1)), os.path.join(directory, name))) + if not entries: + raise BadRequest("mask directory contains no .png masks") + entries.sort() + self.times = [t for t, _ in entries] + self.files = [p for _, p in entries] + gaps = np.diff(self.times) if len(self.times) > 1 else np.array([frame_ms]) + # A mask covers frames up to 1.5 sampling intervals away (at least one frame). + self.tolerance = max(float(np.median(gaps)) * 1.5, frame_ms) + self._cache: tuple[int, np.ndarray] | None = None + + def lookup(self, t_ms: int, width: int, height: int) -> np.ndarray | None: + """Binary mask (uint8 0/255) for the frame at t_ms, resized to the frame, or None.""" + i = bisect.bisect_left(self.times, t_ms) + best = None + for j in (i - 1, i): + if not 0 <= j < len(self.times) or abs(self.times[j] - t_ms) > self.tolerance: + continue + if best is None or abs(self.times[j] - t_ms) < abs(self.times[best] - t_ms): + best = j + if best is None: + return None + if self._cache is not None and self._cache[0] == best: + return self._cache[1] + mask = cv2.imread(self.files[best], cv2.IMREAD_GRAYSCALE) + if mask is None: + raise MediaError("internal", f"unreadable mask {os.path.basename(self.files[best])}") + if mask.shape != (height, width): + mask = cv2.resize(mask, (width, height), interpolation=cv2.INTER_NEAREST) + mask = np.where(mask > 127, 255, 0).astype(np.uint8) + self._cache = (best, mask) + return mask diff --git a/infra/media-worker/app/ops/redact.py b/infra/media-worker/app/ops/redact.py new file mode 100644 index 0000000..053243e --- /dev/null +++ b/infra/media-worker/app/ops/redact.py @@ -0,0 +1,271 @@ +"""redact: frame-accurate face and region redaction (OpenCV + ffmpeg pipes). + +Geometry per instruction and frame: +- Boxes interpolate linearly between keyframes. A single keyframe is a static + region for the whole video; otherwise the region applies from the first to + the last keyframe (plus half a frame either side). +- `paddingPct` grows the box by that percentage of its width/height (half on + each side). Mask regions are dilated by the same share of their extent. +- `featherPct` fades the effect outward over that percentage of the region's + shorter side. The region itself is always fully covered. + +Fails closed: an unreadable frame, a mask that cannot be read, or an encoder +error is a 500 and no file is left at outPath. +""" + +from __future__ import annotations + +import bisect +import os +import threading +from dataclasses import dataclass +from typing import Any + +import cv2 +import numpy as np + +from .. import media_io, paths +from ..errors import BadRequest +from ..schemas import RedactionInstruction, RedactRequest +from .common import Ctx, open_input +from .masks import MaskIndex + +OUTPUT_EXTENSIONS = (".mkv", ".mp4", ".mov") +DEFAULT_FILL = (0, 0, 0) + + +def hex_to_bgr(color: str | None) -> tuple[int, int, int] | None: + if not color: + return None + value = color.lstrip("#") + r, g, b = int(value[0:2], 16), int(value[2:4], 16), int(value[4:6], 16) + return (b, g, r) + + +@dataclass +class Plan: + mode: str + padding: float + feather: float + color: tuple[int, int, int] | None + times: list[int] + boxes: list[tuple[float, float, float, float]] + landmark_times: list[int] + landmarks: list[list[tuple[float, float]]] + masks: MaskIndex | None + + +def build_plan(ctx: Ctx, instruction: RedactionInstruction, frame_ms: float) -> Plan: + keyframes = sorted(instruction.keyframes, key=lambda k: k.tMs) + landmarks = sorted(instruction.landmarks or [], key=lambda frame: frame.tMs) + masks = None + if instruction.maskDir: + masks = MaskIndex(paths.input_dir(ctx.settings.media_root, instruction.maskDir), frame_ms) + return Plan( + mode=instruction.mode, + padding=instruction.paddingPct / 100.0, + feather=instruction.featherPct / 100.0, + color=hex_to_bgr(instruction.color), + times=[k.tMs for k in keyframes], + boxes=[(k.box.x, k.box.y, k.box.w, k.box.h) for k in keyframes], + landmark_times=[frame.tMs for frame in landmarks], + landmarks=[[(p.x, p.y) for p in frame.points] for frame in landmarks], + masks=masks, + ) + + +def interpolate(times: list[int], values: list[Any], t_ms: float, frame_ms: float, hold_single: bool) -> Any: + """Linear interpolation of tuples/lists over keyframe times; None outside the span.""" + if not times: + return None + if len(times) == 1: + if hold_single or abs(t_ms - times[0]) <= frame_ms / 2: + return values[0] + return None + if t_ms < times[0] - frame_ms / 2 or t_ms > times[-1] + frame_ms / 2: + return None + if t_ms <= times[0]: + return values[0] + if t_ms >= times[-1]: + return values[-1] + i = bisect.bisect_right(times, t_ms) + t0, t1 = times[i - 1], times[i] + a, b = values[i - 1], values[i] + if t1 == t0: + return b + f = (t_ms - t0) / (t1 - t0) + if isinstance(a, list): + if len(a) != len(b): + return a if f < 0.5 else b + return [tuple(pa[k] + (pb[k] - pa[k]) * f for k in range(2)) for pa, pb in zip(a, b, strict=True)] + return tuple(a[k] + (b[k] - a[k]) * f for k in range(len(a))) + + +def _pixel_box( + box: tuple[float, float, float, float], padding: float, width: int, height: int +) -> tuple[int, int, int, int]: + x, y, w, h = box + x -= w * padding / 2 + y -= h * padding / 2 + w *= 1 + padding + h *= 1 + padding + x0 = int(np.floor(max(0.0, x) * width)) + y0 = int(np.floor(max(0.0, y) * height)) + x1 = int(np.ceil(min(1.0, x + w) * width)) + y1 = int(np.ceil(min(1.0, y + h) * height)) + return x0, y0, max(x1, x0), max(y1, y0) + + +def blur_effect(roi: np.ndarray, size: int) -> np.ndarray: + """Strong, size-scaled blur: downsample, upsample, Gaussian.""" + h, w = roi.shape[:2] + factor = max(4, size // 6) + small = cv2.resize(roi, (max(1, w // factor), max(1, h // factor)), interpolation=cv2.INTER_AREA) + up = cv2.resize(small, (w, h), interpolation=cv2.INTER_LINEAR) + k = max(3, (size // 3) | 1) + return cv2.GaussianBlur(up, (k, k), 0) + + +def pixelate_effect(roi: np.ndarray, size: int) -> np.ndarray: + h, w = roi.shape[:2] + block = max(4, size // 10) + small = cv2.resize(roi, (max(1, -(-w // block)), max(1, -(-h // block))), interpolation=cv2.INTER_AREA) + return cv2.resize(small, (w, h), interpolation=cv2.INTER_NEAREST) + + +def _ring_color(frame: np.ndarray, region: np.ndarray, margin: int) -> tuple[int, int, int]: + """Median colour just outside the region: the 'skin tone patch' cover fill.""" + kernel = np.ones((2 * margin + 1, 2 * margin + 1), np.uint8) + ring = (cv2.dilate(region, kernel) > 0) & (region == 0) + if not ring.any(): + return DEFAULT_FILL + b, g, r = np.median(frame[ring], axis=0) + return (int(b), int(g), int(r)) + + +def _region_for( + plan: Plan, t_ms: int, frame_ms: float, width: int, height: int +) -> tuple[np.ndarray, int] | None: + """Binary region mask (uint8 0/255) and its characteristic size in px.""" + region = np.zeros((height, width), np.uint8) + size = 0 + hold = plan.masks is None + box = interpolate(plan.times, plan.boxes, t_ms, frame_ms, hold_single=hold) + if box is not None: + x0, y0, x1, y1 = _pixel_box(box, plan.padding, width, height) + raw_w, raw_h = box[2] * width, box[3] * height + if plan.mode in ("partial_eyes", "partial_upper", "partial_lower"): + _partial_region(plan, region, box, (x0, y0, x1, y1), t_ms, frame_ms, width, height) + else: + region[y0:y1, x0:x1] = 255 + size = max(size, int(min(raw_w, raw_h))) + if plan.masks is not None: + mask = plan.masks.lookup(t_ms, width, height) + if mask is not None: + ys, xs = np.nonzero(mask) + if xs.size: + extent = int(min(xs.max() - xs.min() + 1, ys.max() - ys.min() + 1)) + grow = int(round(extent * plan.padding / 2)) + if grow > 0: + mask = cv2.dilate(mask, np.ones((2 * grow + 1, 2 * grow + 1), np.uint8)) + region |= mask + size = max(size, extent) + if not region.any(): + return None + return region, max(size, 1) + + +def _partial_region( + plan: Plan, + region: np.ndarray, + box: tuple[float, float, float, float], + padded: tuple[int, int, int, int], + t_ms: int, + frame_ms: float, + width: int, + height: int, +) -> None: + x0, y0, x1, y1 = padded + bx, by, bw, bh = box + points = interpolate(plan.landmark_times, plan.landmarks, t_ms, frame_ms, hold_single=False) + # Five-point landmarks: right eye, left eye, nose, mouth corners. + eyes = points[:2] if points and len(points) >= 2 else None + nose_y = points[2][1] if points and len(points) >= 3 else by + bh * 0.55 + if plan.mode == "partial_eyes": + if eyes: + (ex1, ey1), (ex2, ey2) = eyes + else: + ex1, ey1, ex2, ey2 = bx + bw * 0.3, by + bh * 0.4, bx + bw * 0.7, by + bh * 0.4 + cx, cy = (ex1 + ex2) / 2 * width, (ey1 + ey2) / 2 * height + angle = float(np.degrees(np.arctan2((ey2 - ey1) * height, (ex2 - ex1) * width))) + band_w = (x1 - x0) * 1.0 + band_h = max(bh * height * 0.32 * (1 + plan.padding), 2.0) + rect = cv2.boxPoints(((cx, cy), (band_w, band_h), angle)).astype(np.int32) + cv2.fillConvexPoly(region, rect, 255) + elif plan.mode == "partial_upper": + split = int(round(min(max(nose_y, by), by + bh) * height)) + region[y0 : max(split, y0), x0:x1] = 255 + else: + split = int(round(min(max(nose_y, by), by + bh) * height)) + region[min(split, y1) : y1, x0:x1] = 255 + + +def apply_plan(frame: np.ndarray, plan: Plan, t_ms: int, frame_ms: float) -> None: + height, width = frame.shape[:2] + found = _region_for(plan, t_ms, frame_ms, width, height) + if found is None: + return + region, size = found + radius = int(round(plan.feather * size)) + ys, xs = np.nonzero(region) + margin = radius + 2 + x0, x1 = max(0, int(xs.min()) - margin), min(width, int(xs.max()) + margin + 1) + y0, y1 = max(0, int(ys.min()) - margin), min(height, int(ys.max()) + margin + 1) + roi = frame[y0:y1, x0:x1] + roi_region = region[y0:y1, x0:x1] + mode = plan.mode + if mode in ("blur", "partial_upper", "partial_lower"): + effected = blur_effect(roi, size) + elif mode == "pixelate": + effected = pixelate_effect(roi, size) + else: + if mode == "cover" and plan.color is None: + fill = _ring_color(frame, region, max(4, size // 20)) + else: + fill = plan.color or DEFAULT_FILL + effected = np.empty_like(roi) + effected[:] = fill + alpha = roi_region.astype(np.float32) / 255.0 + if radius > 0: + kernel = np.ones((2 * radius + 1, 2 * radius + 1), np.uint8) + grown = cv2.dilate(roi_region, kernel).astype(np.float32) / 255.0 + k = 2 * radius + 1 + soft = cv2.GaussianBlur(grown, (k, k), 0) + alpha = np.maximum(alpha, soft) + alpha = alpha[..., None] + blended = roi.astype(np.float32) * (1.0 - alpha) + effected.astype(np.float32) * alpha + roi[:] = np.clip(np.rint(blended), 0, 255).astype(np.uint8) + + +def apply_redactions(frame: np.ndarray, t_ms: int, plans: list[Plan], frame_ms: float) -> np.ndarray: + for plan in plans: + apply_plan(frame, plan, t_ms, frame_ms) + return frame + + +def run(ctx: Ctx, body: RedactRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + ctx.registry.runtime("redact") + out_real = paths.output_file(ctx.settings.media_root, body.outPath) + if os.path.splitext(out_real)[1].lower() not in OUTPUT_EXTENSIONS: + raise BadRequest("outPath must end in .mkv, .mp4 or .mov") + if out_real == real: + raise BadRequest("outPath must differ from path") + frame_ms = info.frame_ms() + plans = [build_plan(ctx, r, frame_ms) for r in body.redactions] + + def per_frame(_index: int, t_ms: int, frame: np.ndarray) -> np.ndarray: + return apply_redactions(frame, t_ms, plans, frame_ms) + + count = media_io.transform_video(real, out_real, info, cancel, per_frame) + return {"outPath": body.outPath, "framesProcessed": count} diff --git a/infra/media-worker/app/ops/saliency.py b/infra/media-worker/app/ops/saliency.py new file mode 100644 index 0000000..5e02e8b --- /dev/null +++ b/infra/media-worker/app/ops/saliency.py @@ -0,0 +1,58 @@ +"""saliency: spectral residual (Hou & Zhang, 2007), implemented with numpy/OpenCV. + +The point is the saliency-weighted centroid of the top 5 % most salient +pixels; strength is how much those pixels stand out against the frame mean. +""" + +from __future__ import annotations + +import threading +from typing import Any + +import cv2 +import numpy as np + +from .. import media_io +from ..schemas import TimesRequest +from .common import Ctx, clamp01, open_input, unique_sorted + +WORK_SIZE = 64 + + +def saliency_map(frame: np.ndarray) -> np.ndarray: + gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY).astype(np.float32) + small = cv2.resize(gray, (WORK_SIZE, WORK_SIZE), interpolation=cv2.INTER_AREA) + spectrum = np.fft.fft2(small) + log_amp = np.log(np.abs(spectrum) + 1e-8) + phase = np.angle(spectrum) + residual = log_amp - cv2.blur(log_amp, (3, 3)) + sal = np.abs(np.fft.ifft2(np.exp(residual + 1j * phase))) ** 2 + sal = cv2.GaussianBlur(sal.astype(np.float32), (0, 0), 2.5) + peak = float(sal.max()) + return sal / peak if peak > 0 else sal + + +def salient_point(frame: np.ndarray) -> tuple[float, float, float]: + sal = saliency_map(frame) + threshold = float(np.quantile(sal, 0.95)) + top = np.where(sal >= threshold, sal, 0.0) + total = float(top.sum()) + if total <= 0: + return 0.5, 0.5, 0.0 + ys, xs = np.mgrid[0 : sal.shape[0], 0 : sal.shape[1]] + cx = float((xs * top).sum() / total + 0.5) / sal.shape[1] + cy = float((ys * top).sum() / total + 0.5) / sal.shape[0] + mean = float(sal.mean()) + top_mean = float(sal[sal >= threshold].mean()) + strength = (top_mean - mean) / top_mean if top_mean > 0 else 0.0 + return clamp01(cx), clamp01(cy), clamp01(strength) + + +def run(ctx: Ctx, body: TimesRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + ctx.registry.runtime("saliency") + points = [] + for t_ms, frame in media_io.frames_at(real, info, unique_sorted(body.timesMs), cancel, 512): + cx, cy, strength = salient_point(frame) + points.append({"tMs": t_ms, "cx": cx, "cy": cy, "strength": strength}) + return {"points": points} diff --git a/infra/media-worker/app/ops/scenes.py b/infra/media-worker/app/ops/scenes.py new file mode 100644 index 0000000..a0d2fb6 --- /dev/null +++ b/infra/media-worker/app/ops/scenes.py @@ -0,0 +1,103 @@ +"""detectScenes: PySceneDetect (adaptive) for shots, histogram grouping for scenes.""" + +from __future__ import annotations + +import threading +from typing import Any + +import cv2 +import numpy as np +from scenedetect import AdaptiveDetector, SceneManager, StatsManager, open_video + +from .. import media_io +from ..errors import Cancelled +from ..schemas import PathRequest +from .common import Ctx, clamp01, open_input + +# Neighbouring shots whose middle frames correlate at least this much share a scene. +SCENE_HIST_CORRELATION = 0.6 +ADAPTIVE_THRESHOLD = 3.0 + + +class SceneDetector: + def detect(self, path: str, cancel: threading.Event) -> list[tuple[float, float, float]]: + """Return shots as (start_s, end_s, score in [0,1]).""" + video = open_video(path, backend="opencv") + stats = StatsManager() + manager = SceneManager(stats_manager=stats) + manager.add_detector(AdaptiveDetector(adaptive_threshold=ADAPTIVE_THRESHOLD)) + stop = threading.Event() + + def watch() -> None: + while not stop.wait(0.2): + if cancel.is_set(): + manager.stop() + return + + watcher = threading.Thread(target=watch, daemon=True) + watcher.start() + try: + manager.detect_scenes(video=video) + finally: + stop.set() + if cancel.is_set(): + raise Cancelled() + scenes = manager.get_scene_list(start_in_scene=True) + ratio_key = next( + (k for k in AdaptiveDetector().get_metrics() if k.startswith("adaptive_ratio")), None + ) + shots: list[tuple[float, float, float]] = [] + for index, (start, end) in enumerate(scenes): + score = 1.0 + if index > 0 and ratio_key: + value = ( + stats.get_metrics(start.frame_num, [ratio_key])[0] + if stats.metrics_exist(start.frame_num, [ratio_key]) + else None + ) + score = clamp01((value or ADAPTIVE_THRESHOLD) / (2 * ADAPTIVE_THRESHOLD)) + shots.append((start.seconds, end.seconds, score)) + return shots + + +def load_scene_detector(_files: dict[str, str], _settings: Any) -> SceneDetector: + return SceneDetector() + + +def _histogram(frame: np.ndarray) -> np.ndarray: + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) + hist = cv2.calcHist([hsv], [0, 1], None, [32, 16], [0, 180, 0, 256]) + cv2.normalize(hist, hist) + return hist + + +def run(ctx: Ctx, body: PathRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + detector = ctx.registry.runtime("detectScenes") + raw = detector.detect(real, cancel) + duration_ms = int(round(info.duration_s * 1000)) + if not raw: + raw = [(0.0, info.duration_s, 1.0)] + shots: list[dict[str, Any]] = [] + for start_s, end_s, score in raw: + start_ms = int(round(start_s * 1000)) + end_ms = max(int(round(end_s * 1000)), start_ms) + if duration_ms: + end_ms = min(end_ms, duration_ms) + shots.append({"startMs": start_ms, "endMs": end_ms, "score": float(score)}) + + scenes: list[dict[str, int]] = [] + previous: np.ndarray | None = None + for shot in shots: + middle = (shot["startMs"] + shot["endMs"]) // 2 + hist = _histogram(media_io.frame_at(real, info, middle, cancel, max_side=320)) + same = ( + previous is not None + and cv2.compareHist(previous, hist, cv2.HISTCMP_CORREL) >= SCENE_HIST_CORRELATION + ) + if same and scenes: + scenes[-1]["endMs"] = shot["endMs"] + else: + scenes.append({"startMs": shot["startMs"], "endMs": shot["endMs"]}) + previous = hist + return {"shots": shots, "scenes": scenes} diff --git a/infra/media-worker/app/ops/segment.py b/infra/media-worker/app/ops/segment.py new file mode 100644 index 0000000..87334bd --- /dev/null +++ b/infra/media-worker/app/ops/segment.py @@ -0,0 +1,244 @@ +"""segmentPrompted: SAM2 (Apache-2.0) from a click or box, propagated over the video. + +Frames are sampled at up to 15 fps and fed to the SAM2 video predictor in +chunks so memory stays bounded; each chunk is re-prompted with the bounding +box of the previous chunk's last mask. Masks are written as `.png` +(0 or 255) at the source resolution. +""" + +from __future__ import annotations + +import os +import tempfile +import threading +from collections.abc import Iterator +from typing import Any + +import cv2 +import numpy as np + +from .. import paths +from ..errors import BadRequest, Cancelled, MediaError +from ..media_io import FFMPEG, run_cancellable, scaled_size +from ..schemas import SegmentPromptedRequest +from .common import Ctx, clamp01, norm_box, open_input +from .devices import torch_device +from .masks import MASK_NAME + +SEGMENT_MAX_FPS = 15.0 +SEGMENT_MAX_SIDE = 1024 +CHUNK_FRAMES = 64 + +MaskStream = Iterator[tuple[int, np.ndarray, float]] + + +def _mask_box(mask: np.ndarray) -> tuple[int, int, int, int] | None: + ys, xs = np.nonzero(mask) + if xs.size == 0: + return None + return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1 + + +class Sam2Segmenter: + def __init__(self, checkpoint: str, settings: Any) -> None: + import torch + from sam2.build_sam import build_sam2_video_predictor + + self.torch = torch + self.device = torch_device(settings) + self.predictor = build_sam2_video_predictor( + "configs/sam2.1/sam2.1_hiera_t.yaml", checkpoint, device=self.device + ) + self._lock = threading.Lock() + + def _chunk( + self, + frames: list[str], + prompt_at: int, + box: tuple[int, int, int, int] | None, + point: tuple[int, int] | None, + reverse: bool, + cancel: threading.Event, + ) -> MaskStream: + with tempfile.TemporaryDirectory(prefix="sam2-") as chunk_dir: + for i, source in enumerate(frames): + os.symlink(source, os.path.join(chunk_dir, f"{i}.jpg")) + torch = self.torch + with torch.inference_mode(): + state = self.predictor.init_state(video_path=chunk_dir, offload_video_to_cpu=True) + kwargs: dict[str, Any] = {} + if box is not None: + kwargs["box"] = np.array(box, dtype=np.float32) + if point is not None: + kwargs["points"] = np.array([point], dtype=np.float32) + kwargs["labels"] = np.array([1], dtype=np.int32) + self.predictor.add_new_points_or_box(state, frame_idx=prompt_at, obj_id=1, **kwargs) + for index, _, logits in self.predictor.propagate_in_video( + state, start_frame_idx=prompt_at, reverse=reverse + ): + if cancel.is_set(): + raise Cancelled() + values = logits[0, 0].float().cpu().numpy() + mask = values > 0 + confidence = float(1.0 / (1.0 + np.exp(-values[mask])).mean()) if mask.any() else 0.0 + yield index, mask, confidence + self.predictor.reset_state(state) + + def segment( + self, + frames: list[str], + prompt_index: int, + box: tuple[int, int, int, int] | None, + point: tuple[int, int] | None, + propagate: bool, + cancel: threading.Event, + ) -> MaskStream: + with self._lock: + if not propagate: + for _, mask, confidence in self._chunk([frames[prompt_index]], 0, box, point, False, cancel): + yield prompt_index, mask, confidence + return + yield from self._forward(frames, prompt_index, box, point, cancel) + yield from self._backward(frames, prompt_index, box, point, cancel) + + def _forward( + self, + frames: list[str], + anchor: int, + box: tuple[int, int, int, int] | None, + point: tuple[int, int] | None, + cancel: threading.Event, + ) -> MaskStream: + emit_anchor = True + while True: + chunk = frames[anchor : anchor + CHUNK_FRAMES] + last: np.ndarray | None = None + for index, mask, confidence in self._chunk(chunk, 0, box, point, False, cancel): + if index == 0 and not emit_anchor: + continue + last = mask + yield anchor + index, mask, confidence + end = anchor + len(chunk) - 1 + box = _mask_box(last) if last is not None else None + if end >= len(frames) - 1 or box is None: + return + anchor, point, emit_anchor = end, None, False + + def _backward( + self, + frames: list[str], + anchor: int, + box: tuple[int, int, int, int] | None, + point: tuple[int, int] | None, + cancel: threading.Event, + ) -> MaskStream: + while anchor > 0: + start = max(0, anchor - CHUNK_FRAMES + 1) + last: np.ndarray | None = None + for index, mask, confidence in self._chunk( + frames[start : anchor + 1], anchor - start, box, point, True, cancel + ): + if start + index == anchor: + continue # emitted by the forward pass + last = mask + yield start + index, mask, confidence + box = _mask_box(last) if last is not None else None + if start == 0 or box is None: + return + anchor, point = start, None + + +def load_segmenter(files: dict[str, str], settings: Any) -> Sam2Segmenter: + return Sam2Segmenter(files["sam2.1_hiera_tiny.pt"], settings) + + +def _extract_frames( + path: str, + fps: float, + size: tuple[int, int], + directory: str, + cancel: threading.Event, + at_ms: int | None = None, +) -> list[str]: + argv = [FFMPEG, "-v", "error", "-nostdin"] + if at_ms is not None: + argv += ["-ss", f"{at_ms / 1000:.3f}"] + argv += ["-i", path, "-map", "0:v:0", "-an", "-sn"] + filters = [] if at_ms is not None else [f"fps={fps:.6f}"] + filters.append(f"scale={size[0]}:{size[1]}:flags=area") + argv += ["-vf", ",".join(filters)] + if at_ms is not None: + argv += ["-frames:v", "1"] + argv += ["-q:v", "2", "-start_number", "0", os.path.join(directory, "%06d.jpg")] + result = run_cancellable(argv, cancel) + if result.returncode != 0: + raise MediaError("internal", "frame extraction failed") + names = sorted(n for n in os.listdir(directory) if n.endswith(".jpg")) + if not names: + raise MediaError("internal", "no frames extracted") + return [os.path.join(directory, n) for n in names] + + +def run(ctx: Ctx, body: SegmentPromptedRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel) + out_dir = paths.output_dir(ctx.settings.media_root, body.outDir) + segmenter = ctx.registry.runtime("segmentPrompted") + fps = min(info.fps, SEGMENT_MAX_FPS) + width, height = scaled_size(info.width, info.height, SEGMENT_MAX_SIDE) + prompt = body.prompt + if prompt.tMs > int(info.duration_s * 1000) + 1: + raise BadRequest("prompt.tMs is after the end of the video") + with tempfile.TemporaryDirectory(prefix="segment-") as work: + if body.propagate: + frames = _extract_frames(real, fps, (width, height), work, cancel) + times = [int(round(i * 1000.0 / fps)) for i in range(len(frames))] + prompt_index = int(np.argmin([abs(t - prompt.tMs) for t in times])) + else: + frames = _extract_frames(real, fps, (width, height), work, cancel, at_ms=prompt.tMs) + times = [prompt.tMs] + prompt_index = 0 + box = None + if prompt.box is not None: + b = prompt.box + box = ( + int(round(b.x * width)), + int(round(b.y * height)), + int(round(clamp01(b.x + b.w) * width)), + int(round(clamp01(b.y + b.h) * height)), + ) + point = ( + (int(round(prompt.point.x * width)), int(round(prompt.point.y * height))) + if prompt.point + else None + ) + + # Replace masks from an earlier run in the same directory. + for name in os.listdir(out_dir): + if MASK_NAME.match(name): + os.unlink(os.path.join(out_dir, name)) + keyframes = [] + for index, mask, confidence in segmenter.segment( + frames, prompt_index, box, point, body.propagate, cancel + ): + if cancel.is_set(): + raise Cancelled() + t_ms = times[index] + full = cv2.resize( + mask.astype(np.uint8) * 255, (info.width, info.height), interpolation=cv2.INTER_NEAREST + ) + found = _mask_box(full > 0) + if found is None: + continue + if not cv2.imwrite(os.path.join(out_dir, f"{t_ms}.png"), full): + raise MediaError("internal", "could not write mask") + x0, y0, x1, y1 = found + keyframes.append( + { + "tMs": t_ms, + "box": norm_box(x0, y0, x1 - x0, y1 - y0, info.width, info.height), + "confidence": clamp01(confidence), + } + ) + keyframes.sort(key=lambda k: k["tMs"]) + score = float(np.mean([k["confidence"] for k in keyframes])) if keyframes else 0.0 + return {"maskDir": body.outDir, "keyframes": keyframes, "score": score} diff --git a/infra/media-worker/app/ops/synthesize.py b/infra/media-worker/app/ops/synthesize.py new file mode 100644 index 0000000..81244d5 --- /dev/null +++ b/infra/media-worker/app/ops/synthesize.py @@ -0,0 +1,20 @@ +"""synthesizeIdentity: deliberately unavailable. + +No generator + face-swap pipeline with a commercial licence exists that this +sidecar can ship, so the operation always answers 501. The request schema +forbids unknown fields, so a target-face image is rejected before this runs. +""" + +from __future__ import annotations + +import threading +from typing import Any + +from ..errors import Unsupported +from ..schemas import SynthesizeIdentityRequest +from .common import Ctx + + +def run(ctx: Ctx, _body: SynthesizeIdentityRequest, _cancel: threading.Event) -> dict[str, Any]: + spec = ctx.registry.backend("synthesizeIdentity").spec + raise Unsupported(spec.unavailable_reason or "no commercially licensed model") diff --git a/infra/media-worker/app/ops/transcribe.py b/infra/media-worker/app/ops/transcribe.py new file mode 100644 index 0000000..b1fbc63 --- /dev/null +++ b/infra/media-worker/app/ops/transcribe.py @@ -0,0 +1,57 @@ +"""transcribe: faster-whisper with word timestamps.""" + +from __future__ import annotations + +import os +import threading +from typing import Any + +from ..errors import BadRequest, Cancelled +from ..schemas import TranscribeRequest +from .common import Ctx, clamp01, open_input +from .devices import ctranslate2_device + + +class WhisperTranscriber: + def __init__(self, model_dir: str, settings: Any) -> None: + from faster_whisper import WhisperModel + + device = ctranslate2_device(settings) + self.model = WhisperModel( + model_dir, + device=device, + compute_type="float16" if device == "cuda" else "int8", + ) + + def transcribe( + self, path: str, language: str | None, cancel: threading.Event + ) -> tuple[str | None, list[dict[str, Any]]]: + segments, info = self.model.transcribe(path, language=language, word_timestamps=True, vad_filter=True) + words: list[dict[str, Any]] = [] + for segment in segments: + if cancel.is_set(): + raise Cancelled() + for word in segment.words or []: + start = int(round(word.start * 1000)) + words.append( + { + "startMs": start, + "endMs": max(int(round(word.end * 1000)), start), + "text": word.word.strip(), + "confidence": clamp01(word.probability), + } + ) + return info.language, words + + +def load_transcriber(files: dict[str, str], settings: Any) -> WhisperTranscriber: + return WhisperTranscriber(os.path.dirname(files["model.bin"]), settings) + + +def run(ctx: Ctx, body: TranscribeRequest, cancel: threading.Event) -> dict[str, Any]: + real, info = open_input(ctx, body.path, cancel, need_video=False) + if not info.has_audio: + raise BadRequest("input has no audio stream") + transcriber = ctx.registry.runtime("transcribe") + language, words = transcriber.transcribe(real, body.language, cancel) + return {"language": language, "words": [w for w in words if w["text"]]} diff --git a/infra/media-worker/app/paths.py b/infra/media-worker/app/paths.py new file mode 100644 index 0000000..40ff80e --- /dev/null +++ b/infra/media-worker/app/paths.py @@ -0,0 +1,75 @@ +"""MEDIA_ROOT enforcement. Every path a request names goes through here. + +Inputs and outputs must be absolute and stay inside the real MEDIA_ROOT after +symlinks are resolved. A path that does not exist yet is checked through its +nearest existing ancestor (realpath resolves the existing prefix). +""" + +from __future__ import annotations + +import os +import stat + +from .errors import InvalidPath, MediaError + + +def _inside(root: str, candidate: str) -> bool: + return candidate == root or candidate.startswith(root.rstrip(os.sep) + os.sep) + + +def resolve_inside(root: str, requested: str) -> str: + """Return the real absolute location of `requested`, or raise InvalidPath.""" + if not isinstance(requested, str) or not requested or "\0" in requested: + raise InvalidPath() + if not os.path.isabs(requested): + raise InvalidPath("path must be absolute") + real_root = os.path.realpath(root) + lexical = os.path.normpath(requested) + # Reject lexical traversal even when it would land back inside the root. + if ".." in requested.split("/"): + raise InvalidPath() + real = os.path.realpath(lexical) + if not _inside(real_root, real): + raise InvalidPath() + return real + + +def input_file(root: str, requested: str, max_bytes: int | None = None) -> str: + real = resolve_inside(root, requested) + try: + info = os.stat(real) + except FileNotFoundError as error: + raise InvalidPath("input file does not exist") from error + if not stat.S_ISREG(info.st_mode): + raise InvalidPath("input is not a regular file") + if max_bytes is not None and info.st_size > max_bytes: + raise MediaError("too_large", "input exceeds MEDIA_MAX_INPUT_BYTES") + return real + + +def input_dir(root: str, requested: str) -> str: + real = resolve_inside(root, requested) + if not os.path.isdir(real): + raise InvalidPath("directory does not exist") + return real + + +def output_file(root: str, requested: str) -> str: + """Resolve an output file path and create its parent directories inside the root.""" + real = resolve_inside(root, requested) + if os.path.isdir(real): + raise InvalidPath("output path is a directory") + parent = os.path.dirname(real) + os.makedirs(parent, exist_ok=True) + # The parent may have been swapped for a symlink in between; check again. + resolve_inside(root, parent) + return real + + +def output_dir(root: str, requested: str) -> str: + real = resolve_inside(root, requested) + os.makedirs(real, exist_ok=True) + resolve_inside(root, real) + if not os.path.isdir(real): + raise InvalidPath("output directory is not a directory") + return real diff --git a/infra/media-worker/app/runtime.py b/infra/media-worker/app/runtime.py new file mode 100644 index 0000000..e5c2fd2 --- /dev/null +++ b/infra/media-worker/app/runtime.py @@ -0,0 +1,93 @@ +"""Request execution: concurrency gates, cancellation on disconnect, time limit.""" + +from __future__ import annotations + +import asyncio +import threading +import time +from collections.abc import Callable +from typing import Any, Protocol + +from .errors import Busy, Cancelled + +POLL_SECONDS = 0.25 + + +class Gate: + """`concurrency` jobs run; up to `queue_size` more wait; anything beyond is busy. + + Thread-based so it works no matter which event loop serves the request. + """ + + def __init__(self, concurrency: int, queue_size: int) -> None: + self.concurrency = concurrency + self.capacity = concurrency + queue_size + self._slots = threading.Semaphore(concurrency) + self._lock = threading.Lock() + self.inflight = 0 + + def reserve(self) -> None: + with self._lock: + if self.inflight >= self.capacity: + raise Busy() + self.inflight += 1 + + def release_reservation(self) -> None: + with self._lock: + self.inflight -= 1 + + def run(self, fn: Callable[[threading.Event], Any], cancel: threading.Event) -> Any: + """Wait for a slot (reservation already taken), run fn, free both.""" + try: + while not self._slots.acquire(timeout=POLL_SECONDS): + if cancel.is_set(): + raise Cancelled() + try: + if cancel.is_set(): + raise Cancelled() + return fn(cancel) + finally: + self._slots.release() + finally: + self.release_reservation() + + +class DisconnectProbe(Protocol): + async def is_disconnected(self) -> bool: ... + + +async def run_job( + request: DisconnectProbe, + gate: Gate, + fn: Callable[[threading.Event], Any], + max_seconds: float, +) -> Any: + """Run fn in a worker thread. Sets the cancel event when the client + disconnects or the deadline passes; fn checks it between frames/chunks.""" + gate.reserve() + cancel = threading.Event() + task = asyncio.ensure_future(asyncio.to_thread(gate.run, fn, cancel)) + deadline = time.monotonic() + max_seconds + reason: str | None = None + try: + while True: + done, _ = await asyncio.wait({task}, timeout=POLL_SECONDS) + if done: + break + if reason is None: + if time.monotonic() > deadline: + reason = "exceeded MEDIA_MAX_SECONDS" + cancel.set() + elif await request.is_disconnected(): + reason = "client disconnected" + cancel.set() + except asyncio.CancelledError: + cancel.set() + raise + try: + result = task.result() + except Cancelled as error: + raise Cancelled(reason or error.reason) from error + if reason is not None: + raise Cancelled(reason) + return result diff --git a/infra/media-worker/app/schemas.py b/infra/media-worker/app/schemas.py new file mode 100644 index 0000000..e6b0973 --- /dev/null +++ b/infra/media-worker/app/schemas.py @@ -0,0 +1,124 @@ +"""Request bodies. Field names mirror the TypeScript request objects (camelCase). + +Every model forbids unknown fields; that is what keeps a target-face image +out of synthesizeIdentity. +""" + +from __future__ import annotations + +from typing import Annotated, Literal + +from pydantic import BaseModel, ConfigDict, Field, StrictBool, model_validator + +Ms = Annotated[int, Field(ge=0, le=7 * 24 * 3600 * 1000)] +Unit = Annotated[float, Field(ge=0.0, le=1.0)] +Path = Annotated[str, Field(min_length=1, max_length=4096)] +SampleFps = Annotated[float, Field(gt=0.0, le=60.0)] +TimesMs = Annotated[list[Ms], Field(max_length=100_000)] + + +class Strict(BaseModel): + model_config = ConfigDict(extra="forbid", strict=True) + + +class NormalizedBox(Strict): + x: Unit + y: Unit + w: Unit + h: Unit + + +class NormalizedPoint(Strict): + x: Unit + y: Unit + + +class BoxKeyframe(Strict): + tMs: Ms + box: NormalizedBox + + +class PathRequest(Strict): + path: Path + + +class TranscribeRequest(PathRequest): + language: Annotated[str, Field(min_length=2, max_length=16, pattern=r"^[A-Za-z-]+$")] | None = None + + +class DetectFacesRequest(PathRequest): + sampleFps: SampleFps + timesMs: TimesMs | None = None + + +class TrackFacesRequest(PathRequest): + sampleFps: SampleFps + + +class TimesRequest(PathRequest): + timesMs: TimesMs + + +class ClassifyContentRequest(PathRequest): + sampleFps: SampleFps + extraTimesMs: TimesMs | None = None + + +class MarkPrompt(Strict): + tMs: Ms + box: NormalizedBox | None = None + point: NormalizedPoint | None = None + + @model_validator(mode="after") + def needs_box_or_point(self) -> MarkPrompt: + if self.box is None and self.point is None: + raise ValueError("prompt needs a box or a point") + return self + + +class SegmentPromptedRequest(PathRequest): + prompt: MarkPrompt + propagate: StrictBool + outDir: Path + + +class InpaintVideoRequest(PathRequest): + maskDir: Path + outPath: Path + + +RedactionMode = Literal[ + "blur", "pixelate", "solid", "sticker", "partial_eyes", "partial_upper", "partial_lower", "cover" +] + + +class LandmarkFrame(Strict): + tMs: Ms + points: Annotated[list[NormalizedPoint], Field(max_length=512)] + + +class RedactionInstruction(Strict): + keyframes: Annotated[list[BoxKeyframe], Field(max_length=200_000)] + mode: RedactionMode + paddingPct: Annotated[float, Field(ge=0.0, le=100.0)] + featherPct: Annotated[float, Field(ge=0.0, le=50.0)] + color: Annotated[str, Field(pattern=r"^#[0-9a-fA-F]{6}$")] | None = None + landmarks: Annotated[list[LandmarkFrame], Field(max_length=200_000)] | None = None + maskDir: Path | None = None + + @model_validator(mode="after") + def needs_region(self) -> RedactionInstruction: + if not self.keyframes and not self.maskDir: + raise ValueError("redaction needs keyframes or a maskDir") + return self + + +class RedactRequest(PathRequest): + redactions: Annotated[list[RedactionInstruction], Field(min_length=1, max_length=1000)] + outPath: Path + + +class SynthesizeIdentityRequest(PathRequest): + track: Annotated[list[BoxKeyframe], Field(min_length=1, max_length=200_000)] + seed: Annotated[int, Field(ge=-(2**53 - 1), le=2**53 - 1)] + outPath: Path diff --git a/infra/media-worker/pytest.ini b/infra/media-worker/pytest.ini new file mode 100644 index 0000000..c87f771 --- /dev/null +++ b/infra/media-worker/pytest.ini @@ -0,0 +1,5 @@ +[pytest] +testpaths = tests +addopts = -p no:cacheprovider +filterwarnings = + ignore:Using `httpx` with `starlette.testclient` is deprecated diff --git a/infra/media-worker/requirements-dev.txt b/infra/media-worker/requirements-dev.txt new file mode 100644 index 0000000..7684ca1 --- /dev/null +++ b/infra/media-worker/requirements-dev.txt @@ -0,0 +1,3 @@ +pytest==9.1.1 +httpx==0.28.1 +ruff==0.16.10 diff --git a/infra/media-worker/requirements-ml-cpu.txt b/infra/media-worker/requirements-ml-cpu.txt new file mode 100644 index 0000000..68371cc --- /dev/null +++ b/infra/media-worker/requirements-ml-cpu.txt @@ -0,0 +1,9 @@ +--extra-index-url https://download.pytorch.org/whl/cpu +torch==2.13.0+cpu +torchvision==0.28.0+cpu +onnxruntime==1.30.0 +faster-whisper==1.2.1 +ctranslate2==4.8.2 +tokenizers==0.23.2 +hydra-core==1.3.7 +iopath==0.1.10 diff --git a/infra/media-worker/requirements-ml-cuda.txt b/infra/media-worker/requirements-ml-cuda.txt new file mode 100644 index 0000000..b723cc5 --- /dev/null +++ b/infra/media-worker/requirements-ml-cuda.txt @@ -0,0 +1,9 @@ +--extra-index-url https://download.pytorch.org/whl/cu126 +torch==2.13.0+cu126 +torchvision==0.28.0+cu126 +onnxruntime-gpu==1.30.0 +faster-whisper==1.2.1 +ctranslate2==4.8.2 +tokenizers==0.23.2 +hydra-core==1.3.7 +iopath==0.1.10 diff --git a/infra/media-worker/requirements-ml-rocm.txt b/infra/media-worker/requirements-ml-rocm.txt new file mode 100644 index 0000000..a9e87ce --- /dev/null +++ b/infra/media-worker/requirements-ml-rocm.txt @@ -0,0 +1,11 @@ +# ROCm: torch (SAM2) runs on the GPU. ONNX Runtime and CTranslate2 have no ROCm +# wheels on PyPI, so SigLIP, LaMa and Whisper run on the CPU in this build. +--extra-index-url https://download.pytorch.org/whl/rocm7.1 +torch==2.13.0+rocm7.1 +torchvision==0.28.0+rocm7.1 +onnxruntime==1.30.0 +faster-whisper==1.2.1 +ctranslate2==4.8.2 +tokenizers==0.23.2 +hydra-core==1.3.7 +iopath==0.1.10 diff --git a/infra/media-worker/requirements.txt b/infra/media-worker/requirements.txt new file mode 100644 index 0000000..b45dc11 --- /dev/null +++ b/infra/media-worker/requirements.txt @@ -0,0 +1,10 @@ +# Core: HTTP API, ffmpeg-driven operations, OpenCV, scene detection, audio features. +# Accelerator-specific ML wheels live in requirements-ml-.txt. +fastapi==0.142.2 +uvicorn[standard]==0.54.0 +pydantic==2.13.5 +numpy==2.4.6 +opencv-python-headless==4.13.0.92 +scenedetect-headless==0.7.1 +librosa==0.11.0 +soundfile==0.14.0 diff --git a/infra/media-worker/ruff.toml b/infra/media-worker/ruff.toml new file mode 100644 index 0000000..ee25629 --- /dev/null +++ b/infra/media-worker/ruff.toml @@ -0,0 +1,5 @@ +line-length = 110 +target-version = "py311" + +[lint] +select = ["E", "F", "W", "B", "I", "UP", "SIM"] diff --git a/infra/media-worker/tests/__init__.py b/infra/media-worker/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/infra/media-worker/tests/conftest.py b/infra/media-worker/tests/conftest.py new file mode 100644 index 0000000..601b5f6 --- /dev/null +++ b/infra/media-worker/tests/conftest.py @@ -0,0 +1,224 @@ +"""Shared fixtures. All media is synthetic (ffmpeg testsrc/sine, drawn shapes).""" + +from __future__ import annotations + +import subprocess +from collections.abc import Callable +from pathlib import Path + +import cv2 +import numpy as np +import pytest +import soundfile +from fastapi.testclient import TestClient + +from app.config import Settings +from app.main import create_app +from app.models import Backend, Registry, default_backends + +from .fakes import fake_backends + +TOKEN = "test-token-" + "x" * 32 +AUTH = {"Authorization": f"Bearer {TOKEN}"} + + +def ffmpeg(*args: str) -> None: + subprocess.run(["ffmpeg", "-v", "error", "-y", *args], check=True) + + +def write_video(path: Path, frames: list[np.ndarray], fps: int = 25) -> None: + height, width = frames[0].shape[:2] + proc = subprocess.Popen( + [ + "ffmpeg", + "-v", + "error", + "-y", + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "-s", + f"{width}x{height}", + "-r", + str(fps), + "-i", + "pipe:0", + "-c:v", + "libx264", + "-crf", + "10", + "-pix_fmt", + "yuv420p", + str(path), + ], + stdin=subprocess.PIPE, + ) + assert proc.stdin is not None + for frame in frames: + proc.stdin.write(frame.tobytes()) + proc.stdin.close() + assert proc.wait() == 0 + + +def decode(path: Path) -> np.ndarray: + probe = ( + subprocess.run( + [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=width,height", + "-of", + "csv=p=0", + str(path), + ], + check=True, + capture_output=True, + text=True, + ) + .stdout.strip() + .split(",") + ) + width, height = int(probe[0]), int(probe[1]) + raw = subprocess.run( + [ + "ffmpeg", + "-v", + "error", + "-i", + str(path), + "-map", + "0:v:0", + "-fps_mode", + "passthrough", + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "pipe:1", + ], + check=True, + capture_output=True, + ).stdout + return np.frombuffer(raw, np.uint8).reshape(-1, height, width, 3) + + +@pytest.fixture(scope="session") +def media_root(tmp_path_factory: pytest.TempPathFactory) -> Path: + root = tmp_path_factory.mktemp("media") + ffmpeg( + "-f", + "lavfi", + "-i", + "testsrc=size=320x240:rate=25:duration=2", + "-f", + "lavfi", + "-i", + "sine=frequency=440:duration=2", + "-c:v", + "libx264", + "-crf", + "10", + "-pix_fmt", + "yuv420p", + "-c:a", + "aac", + "-shortest", + str(root / "testsrc.mp4"), + ) + ffmpeg( + "-f", + "lavfi", + "-i", + "testsrc=size=320x240:rate=25:duration=2", + "-f", + "lavfi", + "-i", + "mandelbrot=size=320x240:rate=25", + "-filter_complex", + "[1:v]trim=duration=2,setpts=PTS-STARTPTS[b];[0:v][b]concat=n=2:v=1:a=0", + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + str(root / "cut.mp4"), + ) + # Seeded 4-px block noise with a tone: lots of local detail for blur/pixelate checks. + rng = np.random.default_rng(7) + noise = [ + np.kron(rng.integers(0, 256, (60, 80, 3), dtype=np.uint8), np.ones((4, 4, 1), np.uint8)) + for _ in range(25) + ] + write_video(root / "noise-silent.mp4", noise) + ffmpeg( + "-i", + str(root / "noise-silent.mp4"), + "-f", + "lavfi", + "-i", + "sine=frequency=330:duration=1", + "-c:v", + "copy", + "-c:a", + "aac", + "-shortest", + str(root / "noise.mp4"), + ) + # A white disc moving left to right on black: the synthetic "face"/"mark". + frames = [] + for i in range(40): + frame = np.zeros((240, 320, 3), np.uint8) + cv2.circle(frame, (60 + i * 5, 120), 25, (255, 255, 255), -1) + frames.append(frame) + write_video(root / "disc.mp4", frames) + # A bright square in the right half for saliency. + frame = np.full((240, 320, 3), 40, np.uint8) + cv2.rectangle(frame, (240, 100), (280, 140), (255, 255, 255), -1) + write_video(root / "spot.mp4", [frame] * 10) + # 120 BPM click track (8 s) followed by 1.5 s of silence. + sr = 22050 + clicks = np.zeros(int(sr * 9.5), np.float32) + for beat in range(16): + start = int(beat * 0.5 * sr) + t = np.arange(int(0.03 * sr)) / sr + clicks[start : start + t.size] += 0.8 * np.sin(2 * np.pi * 1000 * t) * np.exp(-t * 80) + soundfile.write(str(root / "clicks.wav"), clicks, sr) + return root + + +def make_settings(media_root: Path, tmp_path: Path, **overrides: object) -> Settings: + values: dict[str, object] = { + "token": TOKEN, + "media_root": str(media_root), + "models_dir": str(tmp_path / "models"), + "allow_model_download": False, + } + values.update(overrides) + return Settings(**values) # type: ignore[arg-type] + + +ClientFactory = Callable[..., TestClient] + + +@pytest.fixture +def make_client(media_root: Path, tmp_path: Path) -> ClientFactory: + def make( + backends: dict[str, Backend] | None = None, fakes: bool = True, **overrides: object + ) -> TestClient: + settings = make_settings(media_root, tmp_path, **overrides) + merged = default_backends() + if fakes: + merged.update(fake_backends()) + merged.update(backends or {}) + return TestClient(create_app(settings, Registry(settings, backends=merged))) + + return make + + +@pytest.fixture +def client(make_client: ClientFactory) -> TestClient: + return make_client() diff --git a/infra/media-worker/tests/fakes.py b/infra/media-worker/tests/fakes.py new file mode 100644 index 0000000..8be7527 --- /dev/null +++ b/infra/media-worker/tests/fakes.py @@ -0,0 +1,133 @@ +"""Deterministic offline stand-ins for the model runtimes. + +They implement the same small runtime interfaces as the real backends, so the +operation code around them (frame IO, tracking, masks, encoding) runs for real. +Synthetic inputs only: bright shapes on dark backgrounds stand in for faces +and marks. +""" + +from __future__ import annotations + +import threading +from typing import Any + +import cv2 +import numpy as np + +from app.models import Backend, BackendSpec +from app.ops.faces import Detection + + +def fake_backend(runtime: Any, model: str = "fake", experimental: bool = False) -> Backend: + return Backend( + BackendSpec( + id=f"fake-{model}", + model=model, + license="MIT", + commercialUse=True, + experimental=experimental, + source=None, + ), + lambda _files, _settings: runtime, + ) + + +def bright_blobs(frame: np.ndarray, threshold: int = 200) -> list[tuple[int, int, int, int]]: + gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + _, binary = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY) + contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + boxes = [cv2.boundingRect(c) for c in contours if cv2.contourArea(c) > 20] + return sorted(boxes) + + +class FakeFaceDetector: + """Every bright blob is a 'face'; landmarks sit at fixed box fractions.""" + + def detect(self, frame: np.ndarray) -> list[Detection]: + result = [] + for x, y, w, h in bright_blobs(frame): + marks = [(0.3, 0.4), (0.7, 0.4), (0.5, 0.6), (0.35, 0.8), (0.65, 0.8)] + result.append(Detection((x, y, w, h), 0.9, [(x + fx * w, y + fy * h) for fx, fy in marks])) + return result + + +class FakeEmbedder: + model = "fake-embed" + dim = 4 + + def embed(self, frames: list[np.ndarray]) -> np.ndarray: + vectors = np.array([[*frame.reshape(-1, 3).mean(axis=0), 1.0] for frame in frames], dtype=np.float32) + return vectors / np.linalg.norm(vectors, axis=1, keepdims=True) + + +class FakeClassifier: + model = "fake-classifier" + + def classify(self, frames: list[np.ndarray]) -> list[dict[str, float]]: + return [{"violence": float(frame[..., 2].mean() / 255.0), "weapons": 0.0} for frame in frames] + + +class FakeTranscriber: + def transcribe( + self, _path: str, language: str | None, _cancel: threading.Event + ) -> tuple[str | None, list[dict[str, Any]]]: + words = [ + {"startMs": 0, "endMs": 400, "text": "synthetic", "confidence": 0.9}, + {"startMs": 450, "endMs": 900, "text": "tone", "confidence": 0.8}, + ] + return language or "en", words + + +class FakeSegmenter: + """Segments the bright blob in each sampled frame.""" + + def segment( + self, + frames: list[str], + prompt_index: int, + _box: Any, + _point: Any, + propagate: bool, + cancel: threading.Event, + ) -> Any: + indices = range(len(frames)) if propagate else [prompt_index] + for index in indices: + if cancel.is_set(): + return + image = cv2.imread(frames[index if propagate else 0]) + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + yield index, gray > 200, 0.95 + + +class FakeInpainter: + """Classical OpenCV inpainting (Telea) instead of LaMa.""" + + def inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray: + return cv2.inpaint(image, (mask > 127).astype(np.uint8), 5, cv2.INPAINT_TELEA) + + +class BlockingEmbedder(FakeEmbedder): + """Holds the GPU slot until released, to exercise queue overflow.""" + + def __init__(self) -> None: + self.entered = threading.Event() + self.release = threading.Event() + + def embed(self, frames: list[np.ndarray]) -> np.ndarray: + self.entered.set() + self.release.wait(10) + return super().embed(frames) + + +def fake_backends() -> dict[str, Backend]: + detector = FakeFaceDetector() + return { + "transcribe": fake_backend(FakeTranscriber(), "fake-whisper"), + "detectFaces": fake_backend(detector, "fake-faces"), + "trackFaces": fake_backend(detector, "fake-faces"), + "faceLandmarks": fake_backend(detector, "fake-faces"), + "embedFrames": fake_backend(FakeEmbedder(), "fake-embed"), + "classifyContent": fake_backend(FakeClassifier(), "fake-classifier", experimental=True), + "segmentPrompted": fake_backend(FakeSegmenter(), "fake-segment", experimental=True), + "inpaintVideo": fake_backend(FakeInpainter(), "fake-inpaint", experimental=True), + } diff --git a/infra/media-worker/tests/fixtures/media-capabilities.json b/infra/media-worker/tests/fixtures/media-capabilities.json new file mode 100644 index 0000000..e3ea669 --- /dev/null +++ b/infra/media-worker/tests/fixtures/media-capabilities.json @@ -0,0 +1,21 @@ +{ + "$comment": "Shape of MediaCapabilitiesSchema in packages/contracts/src/media.ts. Keep in sync with the contract.", + "keys": ["configured", "operations", "accelerator", "hardwareEncoders"], + "operations": [ + "detectScenes", + "transcribe", + "detectFaces", + "trackFaces", + "faceLandmarks", + "embedFrames", + "classifyContent", + "segmentPrompted", + "inpaintVideo", + "redact", + "synthesizeIdentity", + "audioFeatures", + "saliency" + ], + "operationKeys": ["available", "experimental", "model", "license", "commercialUse"], + "accelerators": ["cpu", "cuda", "rocm", "metal", "none"] +} diff --git a/infra/media-worker/tests/test_analysis.py b/infra/media-worker/tests/test_analysis.py new file mode 100644 index 0000000..b02bc1b --- /dev/null +++ b/infra/media-worker/tests/test_analysis.py @@ -0,0 +1,88 @@ +"""Operations without model weights, run for real on synthetic media.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from fastapi.testclient import TestClient + +from app.media_io import parse_ebur128 + +from .conftest import AUTH + + +def test_audio_features_on_click_track(client: TestClient, media_root: Path) -> None: + response = client.post("/v1/audioFeatures", headers=AUTH, json={"path": f"{media_root}/clicks.wav"}) + assert response.status_code == 200, response.text + body = response.json() + assert body["integratedLufs"] is not None and -60 < body["integratedLufs"] < 0 + assert body["truePeakDb"] is not None and body["truePeakDb"] <= 1.0 + tempo = body["tempoBpm"] + assert tempo is not None and any(abs(tempo - bpm) < 6 for bpm in (60, 120, 240)) + beats = body["beatsMs"] + assert len(beats) >= 6 + gaps = np.diff(beats) + assert np.median(gaps) % 500 < 40 or 500 - np.median(gaps) % 500 < 40 + # The trailing 1.5 s are silent. + assert any(7400 <= s["startMs"] <= 7700 and s["endMs"] >= 9400 for s in body["silences"]), body[ + "silences" + ] + energy = body["energy"] + assert energy[0]["tMs"] == 0 and all(0.0 <= e["value"] <= 1.0 for e in energy) + assert max(e["value"] for e in energy) == 1.0 + assert all(isinstance(t, int) for t in beats) + + +def test_audio_features_without_audio_stream(client: TestClient, media_root: Path) -> None: + response = client.post("/v1/audioFeatures", headers=AUTH, json={"path": f"{media_root}/disc.mp4"}) + assert response.status_code == 200 + assert response.json() == { + "integratedLufs": None, + "truePeakDb": None, + "tempoBpm": None, + "beatsMs": [], + "energy": [], + "silences": [], + } + + +def test_parse_ebur128_summary() -> None: + stderr = """ +[Parsed_ebur128_0 @ 0x0] Summary: + + Integrated loudness: + I: -23.4 LUFS + Threshold: -33.6 LUFS + + True peak: + Peak: -1.2 dBFS +""" + assert parse_ebur128(stderr) == (-23.4, -1.2) + assert parse_ebur128("Summary:\n I: -inf LUFS\n Peak: -inf dBFS") == (None, None) + + +def test_scene_detection_finds_hard_cut(client: TestClient, media_root: Path) -> None: + response = client.post("/v1/detectScenes", headers=AUTH, json={"path": f"{media_root}/cut.mp4"}) + assert response.status_code == 200, response.text + body = response.json() + shots = body["shots"] + assert len(shots) == 2, shots + assert abs(shots[1]["startMs"] - 2000) <= 80 + assert shots[0]["startMs"] == 0 and abs(shots[-1]["endMs"] - 4000) <= 80 + assert all(0.0 <= s["score"] <= 1.0 for s in shots) + assert body["scenes"][0]["startMs"] == 0 + assert body["scenes"][-1]["endMs"] == shots[-1]["endMs"] + + +def test_saliency_points_at_bright_spot(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/saliency", headers=AUTH, json={"path": f"{media_root}/spot.mp4", "timesMs": [0, 200]} + ) + assert response.status_code == 200, response.text + points = response.json()["points"] + assert [p["tMs"] for p in points] == [0, 200] + for point in points: + assert point["cx"] > 0.6 + assert 0.3 < point["cy"] < 0.7 + assert 0.0 < point["strength"] <= 1.0 diff --git a/infra/media-worker/tests/test_auth.py b/infra/media-worker/tests/test_auth.py new file mode 100644 index 0000000..8866032 --- /dev/null +++ b/infra/media-worker/tests/test_auth.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +from app.config import ConfigError, Settings + +from .conftest import AUTH, TOKEN + + +def test_healthz_needs_no_token(client: TestClient) -> None: + response = client.get("/healthz") + assert response.status_code == 200 + assert response.json()["ok"] is True + assert isinstance(response.json()["version"], str) + + +@pytest.mark.parametrize( + "headers", + [ + {}, + {"Authorization": "Bearer wrong-token-" + "y" * 32}, + {"Authorization": TOKEN}, + {"Authorization": "Basic " + TOKEN}, + ], +) +def test_routes_reject_missing_or_wrong_token(client: TestClient, headers: dict[str, str]) -> None: + for method, url in [("GET", "/v1/capabilities"), ("POST", "/v1/detectScenes")]: + response = client.request(method, url, headers=headers, json={"path": "/x"}) + assert response.status_code == 401 + assert response.json()["error"]["code"] == "unauthorized" + + +def test_valid_token_is_accepted(client: TestClient) -> None: + assert client.get("/v1/capabilities", headers=AUTH).status_code == 200 + + +def test_refuses_to_start_without_a_long_token(tmp_path: Path) -> None: + env = {"MEDIA_ROOT": str(tmp_path)} + with pytest.raises(ConfigError): + Settings.from_env(env) + with pytest.raises(ConfigError): + Settings.from_env({**env, "MEDIA_WORKER_TOKEN": "short"}) + assert Settings.from_env({**env, "MEDIA_WORKER_TOKEN": "z" * 32}).token == "z" * 32 + assert Settings.from_env({**env, "MEDIA_WORKER_INSECURE_DEV": "1"}).insecure_dev + + +def test_settings_validate_root_and_defaults(tmp_path: Path) -> None: + with pytest.raises(ConfigError): + Settings.from_env({"MEDIA_WORKER_TOKEN": "z" * 32, "MEDIA_ROOT": "relative/path"}) + settings = Settings.from_env({"MEDIA_WORKER_TOKEN": "z" * 32, "MEDIA_ROOT": str(tmp_path)}) + assert settings.host == "127.0.0.1" + assert settings.port == 8790 + assert settings.gpu_concurrency == 1 + assert settings.max_input_bytes == 20 * 1024**3 + assert settings.max_seconds == 4 * 3600 diff --git a/infra/media-worker/tests/test_capabilities.py b/infra/media-worker/tests/test_capabilities.py new file mode 100644 index 0000000..18640f1 --- /dev/null +++ b/infra/media-worker/tests/test_capabilities.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import json +import re +from pathlib import Path + +from fastapi.testclient import TestClient + +from app.models import OPTIONAL_BACKENDS, Registry, default_backends + +from .conftest import AUTH, make_settings + +CONTRACT = json.loads((Path(__file__).parent / "fixtures" / "media-capabilities.json").read_text()) + + +def _assert_shape(body: dict) -> None: + assert sorted(body) == sorted(CONTRACT["keys"]) + assert body["configured"] is True + assert body["accelerator"] in CONTRACT["accelerators"] + assert isinstance(body["hardwareEncoders"], list) + assert all(isinstance(name, str) for name in body["hardwareEncoders"]) + assert sorted(body["operations"]) == sorted(CONTRACT["operations"]) + for entry in body["operations"].values(): + assert sorted(entry) == sorted(CONTRACT["operationKeys"]) + assert isinstance(entry["available"], bool) + assert isinstance(entry["experimental"], bool) + assert entry["model"] is None or isinstance(entry["model"], str) + assert entry["license"] is None or isinstance(entry["license"], str) + assert entry["commercialUse"] in (True, False, None) + + +def test_capabilities_match_the_contract(make_client) -> None: + # Real registry, no fakes, downloads off: whatever is installed here. + body = make_client(fakes=False).get("/v1/capabilities", headers=AUTH).json() + _assert_shape(body) + ops = body["operations"] + # Operations that need no model weights work offline with the core requirements. + for name in ("detectScenes", "redact", "audioFeatures", "saliency"): + assert ops[name]["available"] is True, name + assert ops["synthesizeIdentity"] == { + "available": False, + "experimental": True, + "model": None, + "license": None, + "commercialUse": None, + } + # Weights are not installed and downloads are off. + for name in ("transcribe", "detectFaces", "embedFrames", "inpaintVideo", "segmentPrompted"): + assert ops[name]["available"] is False, name + + +def test_with_fake_backends_everything_but_synthesis_is_available(client: TestClient) -> None: + body = client.get("/v1/capabilities", headers=AUTH).json() + _assert_shape(body) + unavailable = [name for name, entry in body["operations"].items() if not entry["available"]] + assert unavailable == ["synthesizeIdentity"] + + +def test_default_backends_are_commercially_usable_and_pinned() -> None: + for operation, backend in default_backends().items(): + spec = backend.spec + if spec.unavailable_reason: + continue + assert spec.commercialUse is True, operation + assert spec.license, operation + for model_file in spec.files: + assert re.fullmatch(r"[0-9a-f]{64}", model_file.sha256), model_file.name + assert model_file.sizeBytes > 0 + assert model_file.url.startswith("https://") + assert "/resolve/main/" not in model_file.url, "pin a revision, not a branch" + for backend in OPTIONAL_BACKENDS.values(): + assert backend.spec.commercialUse is not True + + +def test_optional_backend_needs_opt_in(media_root: Path, tmp_path: Path) -> None: + registry = Registry(make_settings(media_root, tmp_path)) + assert not registry.optional_enabled("faceIdentity") + reason = registry.unavailable_reason(OPTIONAL_BACKENDS["faceIdentity"], "faceIdentity") + assert reason is not None and "MEDIA_OPTIONAL_BACKENDS" in reason + opted = Registry(make_settings(media_root, tmp_path, optional_backends=frozenset({"faceIdentity"}))) + reason = opted.unavailable_reason(OPTIONAL_BACKENDS["faceIdentity"], "faceIdentity") + assert reason is None or "MEDIA_OPTIONAL_BACKENDS" not in reason diff --git a/infra/media-worker/tests/test_ml_ops.py b/infra/media-worker/tests/test_ml_ops.py new file mode 100644 index 0000000..f21d24a --- /dev/null +++ b/infra/media-worker/tests/test_ml_ops.py @@ -0,0 +1,170 @@ +"""Model-backed operations with fake runtimes; frame IO, tracking, masks and encoding are real.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from fastapi.testclient import TestClient + +from app.ops.faces import Detection, IouTracker + +from .conftest import AUTH, decode + + +def test_transcribe(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/transcribe", headers=AUTH, json={"path": f"{media_root}/testsrc.mp4", "language": "de"} + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["language"] == "de" + assert [w["text"] for w in body["words"]] == ["synthetic", "tone"] + + +def test_transcribe_rejects_silent_video(client: TestClient, media_root: Path) -> None: + response = client.post("/v1/transcribe", headers=AUTH, json={"path": f"{media_root}/disc.mp4"}) + assert response.status_code == 400 + + +def test_detect_faces_sampled_and_at_times(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/detectFaces", headers=AUTH, json={"path": f"{media_root}/disc.mp4", "sampleFps": 5} + ) + assert response.status_code == 200, response.text + frames = response.json()["frames"] + assert [f["tMs"] for f in frames][:3] == [0, 200, 400] + assert all(len(f["faces"]) == 1 for f in frames) + face = frames[0]["faces"][0] + assert abs(face["box"]["x"] - 35 / 320) < 0.02 and abs(face["box"]["w"] - 51 / 320) < 0.02 + assert len(face["landmarks"]) == 5 + response = client.post( + "/v1/detectFaces", + headers=AUTH, + json={"path": f"{media_root}/disc.mp4", "sampleFps": 5, "timesMs": [800, 0]}, + ) + assert [f["tMs"] for f in response.json()["frames"]] == [0, 800] + + +def test_track_faces_follows_the_disc(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/trackFaces", headers=AUTH, json={"path": f"{media_root}/disc.mp4", "sampleFps": 10} + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["sampleFps"] == 10 + assert len(body["tracks"]) == 1 + track = body["tracks"][0] + assert track["trackId"] == "t1" + assert track["embedding"] is None + xs = [k["box"]["x"] for k in track["keyframes"]] + assert len(xs) >= 15 and xs == sorted(xs) and xs[-1] > xs[0] + 0.3 + + +def test_tracker_splits_after_gap_and_uses_low_confidence() -> None: + tracker = IouTracker(max_gap_ms=500) + box = (10.0, 10.0, 50.0, 50.0) + tracker.update(0, [Detection(box, 0.9, [])], None, False) + # Low-confidence detection extends the existing track but never starts one. + matches = tracker.update( + 100, + [Detection((12.0, 10.0, 50.0, 50.0), 0.4, []), Detection((200.0, 200.0, 30.0, 30.0), 0.4, [])], + None, + False, + ) + assert [m[0].track_id for m in matches] == ["t1"] + # 600 ms without a match: the next detection starts a new track. + tracker.update(800, [Detection(box, 0.9, [])], None, False) + assert [t.track_id for t in tracker.close()] == ["t1", "t2"] + + +def test_face_landmarks(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/faceLandmarks", headers=AUTH, json={"path": f"{media_root}/disc.mp4", "timesMs": [0, 400]} + ) + assert response.status_code == 200, response.text + frames = response.json()["frames"] + assert [f["tMs"] for f in frames] == [0, 400] + assert len(frames[1]["faces"][0]["landmarks"]) == 5 + + +def test_embed_frames_keeps_request_order(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/embedFrames", + headers=AUTH, + json={"path": f"{media_root}/testsrc.mp4", "timesMs": [1000, 0, 1000]}, + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["model"] == "fake-embed" and body["dim"] == 4 + assert len(body["vectors"]) == 3 and body["vectors"][0] == body["vectors"][2] + assert all(abs(np.linalg.norm(v) - 1) < 1e-5 for v in body["vectors"]) + + +def test_classify_content_samples_plus_extra_times(client: TestClient, media_root: Path) -> None: + response = client.post( + "/v1/classifyContent", + headers=AUTH, + json={"path": f"{media_root}/testsrc.mp4", "sampleFps": 2, "extraTimesMs": [250, 500]}, + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["model"] == "fake-classifier" + times = [f["tMs"] for f in body["frames"]] + assert times == sorted(times) and {0, 250, 500} <= set(times) + # Samples are the source frames nearest to each 500 ms slot (25 fps source). + for slot in (1000, 1500): + assert min(abs(t - slot) for t in times) <= 20 + assert all(0 <= f["scores"]["violence"] <= 1 for f in body["frames"]) + + +def test_segment_prompted_writes_masks_and_keyframes(client: TestClient, media_root: Path) -> None: + out_dir = media_root / "marks" / "m1" + body = { + "path": f"{media_root}/disc.mp4", + "prompt": {"tMs": 400, "point": {"x": 0.25, "y": 0.5}}, + "propagate": True, + "outDir": str(out_dir), + } + response = client.post("/v1/segmentPrompted", headers=AUTH, json=body) + assert response.status_code == 200, response.text + result = response.json() + assert result["maskDir"] == str(out_dir) + assert len(result["keyframes"]) == 24 # 1.6 s at 15 fps + assert 0.9 < result["score"] <= 1.0 + names = sorted(int(p.stem) for p in out_dir.glob("*.png")) + assert names == [k["tMs"] for k in result["keyframes"]] + single = client.post("/v1/segmentPrompted", headers=AUTH, json=body | {"propagate": False}) + assert single.status_code == 200 + assert [k["tMs"] for k in single.json()["keyframes"]] == [400] + assert [p.name for p in out_dir.glob("*.png")] == ["400.png"] + + +def test_inpaint_video_removes_masked_disc(client: TestClient, media_root: Path) -> None: + out_dir = media_root / "marks" / "inpaint" + seg = client.post( + "/v1/segmentPrompted", + headers=AUTH, + json={ + "path": f"{media_root}/disc.mp4", + "prompt": {"tMs": 0, "point": {"x": 0.2, "y": 0.5}}, + "propagate": True, + "outDir": str(out_dir), + }, + ) + assert seg.status_code == 200 + out = media_root / "inpainted" / "disc.mkv" + response = client.post( + "/v1/inpaintVideo", + headers=AUTH, + json={"path": f"{media_root}/disc.mp4", "maskDir": str(out_dir), "outPath": str(out)}, + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["outPath"] == str(out) + quality = body["quality"] + assert set(quality) == {"temporalStability", "edgeQuality", "score"} + assert all(0.0 <= v <= 1.0 for v in quality.values()) + before, after = decode(media_root / "disc.mp4"), decode(out) + assert after.shape == before.shape + assert after.mean() < before.mean() * 0.3 diff --git a/infra/media-worker/tests/test_models.py b/infra/media-worker/tests/test_models.py new file mode 100644 index 0000000..f6d9617 --- /dev/null +++ b/infra/media-worker/tests/test_models.py @@ -0,0 +1,75 @@ +"""Lazy download and SHA256 pinning, exercised with local file:// URLs.""" + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import pytest + +from app.errors import Unsupported +from app.models import Backend, BackendSpec, ModelFile, Registry + +from .conftest import make_settings + + +def _backend(source: Path, sha256: str, size: int, loads: list[dict[str, str]]) -> Backend: + spec = BackendSpec( + id="test-model", + model="test-model", + license="MIT", + commercialUse=True, + experimental=False, + source=None, + files=(ModelFile("weights.bin", source.as_uri(), sha256, size),), + ) + return Backend(spec, lambda files, _settings: loads.append(files) or "runtime") + + +def test_downloads_once_and_verifies(media_root: Path, tmp_path: Path) -> None: + source = tmp_path / "upstream.bin" + source.write_bytes(b"weights" * 100) + digest = hashlib.sha256(source.read_bytes()).hexdigest() + loads: list[dict[str, str]] = [] + settings = make_settings(media_root, tmp_path, allow_model_download=True) + registry = Registry(settings, backends={"embedFrames": _backend(source, digest, 700, loads)}) + assert registry.runtime("embedFrames") == "runtime" + assert registry.runtime("embedFrames") == "runtime" + assert len(loads) == 1 + cached = Path(loads[0]["weights.bin"]) + assert cached == tmp_path / "models" / "test-model" / "weights.bin" + assert cached.read_bytes() == source.read_bytes() + + +def test_hash_mismatch_deletes_file_and_reports_unsupported(media_root: Path, tmp_path: Path) -> None: + source = tmp_path / "tampered.bin" + source.write_bytes(b"tampered") + settings = make_settings(media_root, tmp_path, allow_model_download=True) + backend = _backend(source, "0" * 64, 8, []) + registry = Registry(settings, backends={"embedFrames": backend}) + with pytest.raises(Unsupported) as error: + registry.runtime("embedFrames") + assert "SHA256" in error.value.message + assert not (tmp_path / "models" / "test-model" / "weights.bin").exists() + assert registry.unavailable_reason(backend, "embedFrames") == error.value.message + + +def test_oversized_download_is_aborted(media_root: Path, tmp_path: Path) -> None: + source = tmp_path / "big.bin" + source.write_bytes(b"x" * 5000) + settings = make_settings(media_root, tmp_path, allow_model_download=True) + registry = Registry(settings, backends={"embedFrames": _backend(source, "0" * 64, 10, [])}) + with pytest.raises(Unsupported): + registry.runtime("embedFrames") + assert not list((tmp_path / "models" / "test-model").iterdir()) + + +def test_offline_without_weights_is_unsupported(media_root: Path, tmp_path: Path) -> None: + source = tmp_path / "w.bin" + source.write_bytes(b"w") + registry = Registry( + make_settings(media_root, tmp_path), backends={"embedFrames": _backend(source, "0" * 64, 1, [])} + ) + with pytest.raises(Unsupported) as error: + registry.runtime("embedFrames") + assert "MEDIA_MODEL_DOWNLOAD=0" in error.value.message diff --git a/infra/media-worker/tests/test_paths.py b/infra/media-worker/tests/test_paths.py new file mode 100644 index 0000000..673bee8 --- /dev/null +++ b/infra/media-worker/tests/test_paths.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import os +import shutil +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +from app.errors import InvalidPath +from app.paths import input_file, output_file, resolve_inside + +from .conftest import AUTH + + +@pytest.fixture +def root(tmp_path: Path) -> Path: + path = tmp_path / "root" + path.mkdir() + (path / "a.txt").write_text("a") + return path + + +@pytest.mark.parametrize( + "requested", + [ + "relative.txt", + "", + "/etc/passwd", + "{root}/../outside.txt", + "{root}/sub/../../outside.txt", + "{root}/a\0.txt", + ], +) +def test_rejects_paths_outside_root(root: Path, requested: str) -> None: + with pytest.raises(InvalidPath): + resolve_inside(str(root), requested.format(root=root)) + + +def test_accepts_paths_inside_root(root: Path) -> None: + assert resolve_inside(str(root), f"{root}/a.txt") == os.path.realpath(root / "a.txt") + assert resolve_inside(str(root), f"{root}/new/dir/file.mkv").endswith("/new/dir/file.mkv") + + +def test_rejects_symlink_escape(root: Path, tmp_path: Path) -> None: + outside = tmp_path / "outside" + outside.mkdir() + (outside / "secret.mp4").write_text("x") + os.symlink(outside, root / "link") + os.symlink(outside / "secret.mp4", root / "file-link.mp4") + with pytest.raises(InvalidPath): + input_file(str(root), f"{root}/link/secret.mp4") + with pytest.raises(InvalidPath): + input_file(str(root), f"{root}/file-link.mp4") + with pytest.raises(InvalidPath): + output_file(str(root), f"{root}/link/new/out.mkv") + assert not (outside / "new").exists() + + +def test_symlink_inside_root_is_allowed(root: Path) -> None: + os.symlink(root / "a.txt", root / "alias.txt") + assert input_file(str(root), f"{root}/alias.txt") == os.path.realpath(root / "a.txt") + + +def test_missing_input_and_directory_input(root: Path) -> None: + with pytest.raises(InvalidPath): + input_file(str(root), f"{root}/missing.mp4") + (root / "dir").mkdir() + with pytest.raises(InvalidPath): + input_file(str(root), f"{root}/dir") + + +def test_output_parent_is_created(root: Path) -> None: + real = output_file(str(root), f"{root}/renders/r1/out.mkv") + assert os.path.isdir(os.path.dirname(real)) + + +def test_http_rejects_traversal_and_symlink_escape( + client: TestClient, media_root: Path, tmp_path: Path +) -> None: + outside = tmp_path / "elsewhere" + outside.mkdir() + shutil.copy(media_root / "testsrc.mp4", outside / "clip.mp4") + link = media_root / "escape-link" + if not link.exists(): + os.symlink(outside, link) + for path in ["/etc/hostname", f"{media_root}/../x.mp4", f"{link}/clip.mp4", "testsrc.mp4"]: + response = client.post("/v1/detectScenes", headers=AUTH, json={"path": path}) + assert response.status_code == 400, path + assert response.json()["error"] == { + "code": "invalid_path", + "message": response.json()["error"]["message"], + "operation": "detectScenes", + } + response = client.post( + "/v1/redact", + headers=AUTH, + json={ + "path": f"{media_root}/testsrc.mp4", + "outPath": f"{link}/out.mkv", + "redactions": [ + { + "keyframes": [{"tMs": 0, "box": {"x": 0, "y": 0, "w": 0.1, "h": 0.1}}], + "mode": "blur", + "paddingPct": 0, + "featherPct": 0, + } + ], + }, + ) + assert response.status_code == 400 + assert response.json()["error"]["code"] == "invalid_path" + assert not (outside / "out.mkv").exists() + + +def test_input_size_limit(make_client, media_root: Path) -> None: + client = make_client(max_input_bytes=10) + response = client.post("/v1/audioFeatures", headers=AUTH, json={"path": f"{media_root}/testsrc.mp4"}) + assert response.status_code == 413 + assert response.json()["error"]["code"] == "too_large" + + +def test_duration_limit(make_client, media_root: Path) -> None: + client = make_client(max_seconds=1) + response = client.post("/v1/audioFeatures", headers=AUTH, json={"path": f"{media_root}/testsrc.mp4"}) + assert response.status_code == 413 diff --git a/infra/media-worker/tests/test_redact.py b/infra/media-worker/tests/test_redact.py new file mode 100644 index 0000000..92091e8 --- /dev/null +++ b/infra/media-worker/tests/test_redact.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import cv2 +import numpy as np +import pytest +from fastapi.testclient import TestClient + +import app.ops.redact as redact_module + +from .conftest import AUTH, decode + +BOX = {"x": 0.25, "y": 0.25, "w": 0.5, "h": 0.5} +# Pixel region of BOX in the 320x240 fixture. +X0, Y0, X1, Y1 = 80, 60, 240, 180 + + +def _body(media_root: Path, out: Path, mode: str, **extra: object) -> dict[str, object]: + return { + "path": f"{media_root}/testsrc.mp4", + "outPath": str(out), + "redactions": [ + { + "keyframes": [{"tMs": 0, "box": BOX}, {"tMs": 2000, "box": BOX}], + "mode": mode, + "paddingPct": 0, + "featherPct": 0, + **extra, + } + ], + } + + +def _outside_mask(shape: tuple[int, ...], margin: int) -> np.ndarray: + mask = np.ones(shape[:2], bool) + mask[Y0 - margin : Y1 + margin, X0 - margin : X1 + margin] = False + return mask + + +def _streams(path: Path) -> list[str]: + out = subprocess.run( + ["ffprobe", "-v", "error", "-show_entries", "stream=codec_type,codec_name", "-of", "json", str(path)], + check=True, + capture_output=True, + text=True, + ).stdout + return [f"{s['codec_type']}:{s['codec_name']}" for s in json.loads(out)["streams"]] + + +def test_solid_redaction_changes_only_the_box(client: TestClient, media_root: Path, tmp_path: Path) -> None: + out = media_root / "redact" / "solid.mkv" + response = client.post("/v1/redact", headers=AUTH, json=_body(media_root, out, "solid", color="#000000")) + assert response.status_code == 200, response.text + assert response.json() == {"outPath": str(out), "framesProcessed": 50} + source, result = decode(media_root / "testsrc.mp4"), decode(out) + assert result.shape == source.shape + inside = result[:, Y0 + 2 : Y1 - 2, X0 + 2 : X1 - 2] + assert inside.max() <= 8 + assert np.abs(source[:, Y0:Y1, X0:X1].astype(int)).mean() > 50 + outside = _outside_mask(source.shape[1:], 2) + diff = np.abs(result.astype(int) - source.astype(int))[:, outside] + assert diff.mean() < 2.0 + assert _streams(out) == ["video:ffv1", "audio:aac"] + assert not [p for p in out.parent.iterdir() if ".partial-" in p.name] + + +@pytest.mark.parametrize("mode", ["blur", "pixelate"]) +def test_blur_and_pixelate_destroy_detail_inside(client: TestClient, media_root: Path, mode: str) -> None: + out = media_root / "redact" / f"{mode}.mp4" + body = _body(media_root, out, mode) | {"path": f"{media_root}/noise.mp4"} + response = client.post("/v1/redact", headers=AUTH, json=body) + assert response.status_code == 200, response.text + source, result = decode(media_root / "noise.mp4"), decode(out) + region = (slice(None), slice(Y0 + 4, Y1 - 4), slice(X0 + 4, X1 - 4)) + + # Local detail (gradient energy) collapses inside the box. + def detail(frames: np.ndarray) -> float: + gray = frames.astype(np.float32).mean(axis=3) + return float(np.abs(np.diff(gray, axis=2)).mean()) + + assert detail(result[region]) < detail(source[region]) * 0.2 + assert ( + result[region].astype(np.float32).var(axis=(1, 2)).mean() + < source[region].astype(np.float32).var(axis=(1, 2)).mean() * 0.2 + ) + outside = _outside_mask(source.shape[1:], 4) + assert np.abs(result.astype(int) - source.astype(int))[:, outside].mean() < 3.0 + assert _streams(out) == ["video:h264", "audio:aac"] + + +@pytest.mark.parametrize("mode", ["sticker", "cover", "partial_eyes", "partial_upper", "partial_lower"]) +def test_other_modes_produce_output(client: TestClient, media_root: Path, mode: str) -> None: + out = media_root / "redact" / f"{mode}.mkv" + body = _body(media_root, out, mode) + body["redactions"][0]["featherPct"] = 10 # type: ignore[index] + body["redactions"][0]["paddingPct"] = 20 # type: ignore[index] + if mode == "partial_eyes": + body["redactions"][0]["landmarks"] = [ # type: ignore[index] + {"tMs": 0, "points": [{"x": 0.4, "y": 0.4}, {"x": 0.6, "y": 0.4}, {"x": 0.5, "y": 0.5}]} + ] + response = client.post("/v1/redact", headers=AUTH, json=body) + assert response.status_code == 200, response.text + source, result = decode(media_root / "testsrc.mp4"), decode(out) + changed = np.abs(result.astype(int) - source.astype(int))[:, Y0:Y1, X0:X1].mean() + assert changed > 10 + corner = np.abs(result.astype(int) - source.astype(int))[:, :20, :20].mean() + assert corner < 2.0 + + +def test_mask_dir_redaction(client: TestClient, media_root: Path) -> None: + masks = media_root / "redact-masks" + masks.mkdir(exist_ok=True) + for t in range(0, 2000, 200): + mask = np.zeros((240, 320), np.uint8) + mask[100:140, 140:180] = 255 + cv2.imwrite(str(masks / f"{t}.png"), mask) + out = media_root / "redact" / "mask.mkv" + body = { + "path": f"{media_root}/testsrc.mp4", + "outPath": str(out), + "redactions": [ + { + "keyframes": [], + "maskDir": str(masks), + "mode": "cover", + "color": "#123456", + "paddingPct": 0, + "featherPct": 0, + } + ], + } + response = client.post("/v1/redact", headers=AUTH, json=body) + assert response.status_code == 200, response.text + result = decode(out) + patch = result[:, 104:136, 144:176].reshape(-1, 3).mean(axis=0) + assert np.allclose(patch, [0x56, 0x34, 0x12], atol=6) + + +def test_failure_mid_video_leaves_no_output( + client: TestClient, media_root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + calls = {"n": 0} + original = redact_module.apply_redactions + + def flaky(frame, t_ms, plans, frame_ms): # type: ignore[no-untyped-def] + calls["n"] += 1 + if calls["n"] == 10: + raise RuntimeError("simulated frame failure") + return original(frame, t_ms, plans, frame_ms) + + monkeypatch.setattr(redact_module, "apply_redactions", flaky) + out = media_root / "redact-fail" / "out.mkv" + response = client.post("/v1/redact", headers=AUTH, json=_body(media_root, out, "blur")) + assert response.status_code == 500 + assert response.json()["error"]["code"] == "internal" + assert not out.exists() + assert list(out.parent.iterdir()) == [] + + +def test_unreadable_input_fails_closed(client: TestClient, media_root: Path) -> None: + broken = media_root / "broken.mp4" + data = (media_root / "testsrc.mp4").read_bytes() + broken.write_bytes(data[: len(data) // 3]) + out = media_root / "redact-broken" / "out.mkv" + response = client.post( + "/v1/redact", headers=AUTH, json=_body(media_root, out, "blur") | {"path": str(broken)} + ) + assert response.status_code in (400, 500) + assert not out.exists() + if out.parent.exists(): + assert not list(out.parent.iterdir()) diff --git a/infra/media-worker/tests/test_runtime.py b/infra/media-worker/tests/test_runtime.py new file mode 100644 index 0000000..75ebb03 --- /dev/null +++ b/infra/media-worker/tests/test_runtime.py @@ -0,0 +1,117 @@ +"""Cancellation on disconnect, the time limit, and GPU queue overflow.""" + +from __future__ import annotations + +import asyncio +import threading +from pathlib import Path + +import numpy as np +import pytest + +from app import media_io +from app.errors import Busy, Cancelled +from app.runtime import Gate, run_job + +from .conftest import AUTH +from .fakes import BlockingEmbedder, fake_backend + + +class FakeRequest: + def __init__(self, disconnect_after: int | None) -> None: + self.calls = 0 + self.disconnect_after = disconnect_after + + async def is_disconnected(self) -> bool: + self.calls += 1 + return self.disconnect_after is not None and self.calls > self.disconnect_after + + +def _spin(seen: dict[str, bool]): # type: ignore[no-untyped-def] + def work(cancel: threading.Event) -> str: + for _ in range(400): + if cancel.wait(0.01): + seen["cancelled"] = True + raise Cancelled() + return "done" + + return work + + +def test_disconnect_cancels_work() -> None: + seen: dict[str, bool] = {} + gate = Gate(1, 0) + with pytest.raises(Cancelled) as error: + asyncio.run(run_job(FakeRequest(disconnect_after=1), gate, _spin(seen), max_seconds=60)) + assert error.value.reason == "client disconnected" + assert seen["cancelled"] + assert gate.inflight == 0 + + +def test_time_limit_cancels_work() -> None: + seen: dict[str, bool] = {} + with pytest.raises(Cancelled) as error: + asyncio.run(run_job(FakeRequest(None), Gate(1, 0), _spin(seen), max_seconds=0.3)) + assert error.value.reason == "exceeded MEDIA_MAX_SECONDS" + + +def test_completed_work_returns_result() -> None: + assert asyncio.run(run_job(FakeRequest(None), Gate(1, 0), lambda _c: "ok", max_seconds=5)) == "ok" + + +def test_gate_overflow_is_busy() -> None: + gate = Gate(1, 1) + gate.reserve() + gate.reserve() + with pytest.raises(Busy): + gate.reserve() + gate.release_reservation() + gate.reserve() + + +def test_cancelled_transform_leaves_no_output(media_root: Path) -> None: + cancel = threading.Event() + path = str(media_root / "testsrc.mp4") + info = media_io.probe(path, threading.Event()) + out = media_root / "cancelled" / "out.mkv" + out.parent.mkdir(exist_ok=True) + + def frame_fn(index: int, _t: int, frame: np.ndarray) -> np.ndarray: + if index == 5: + cancel.set() + return frame + + with pytest.raises(Cancelled): + media_io.transform_video(path, str(out), info, cancel, frame_fn) + assert list(out.parent.iterdir()) == [] + + +def test_http_busy_when_gpu_queue_is_full(make_client, media_root: Path) -> None: + blocking = BlockingEmbedder() + client = make_client( + backends={"embedFrames": fake_backend(blocking, "fake-embed")}, gpu_concurrency=1, queue_size=0 + ) + body = {"path": f"{media_root}/testsrc.mp4", "timesMs": [0]} + results: list[int] = [] + first = threading.Thread( + target=lambda: results.append(client.post("/v1/embedFrames", headers=AUTH, json=body).status_code) + ) + first.start() + assert blocking.entered.wait(20) + second = client.post("/v1/embedFrames", headers=AUTH, json=body) + assert second.status_code == 503 + assert second.json()["error"] == { + "code": "busy", + "message": "inference queue is full", + "operation": "embedFrames", + } + # CPU-bound operations use their own gate and still run. + assert ( + client.post( + "/v1/saliency", headers=AUTH, json={"path": f"{media_root}/spot.mp4", "timesMs": [0]} + ).status_code + == 200 + ) + blocking.release.set() + first.join(20) + assert results == [200] diff --git a/infra/media-worker/tests/test_schema.py b/infra/media-worker/tests/test_schema.py new file mode 100644 index 0000000..e2db207 --- /dev/null +++ b/infra/media-worker/tests/test_schema.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +from .conftest import AUTH + + +def _synth(media_root: Path, **extra: object) -> dict[str, object]: + return { + "path": f"{media_root}/testsrc.mp4", + "track": [{"tMs": 0, "box": {"x": 0.1, "y": 0.1, "w": 0.2, "h": 0.2}}], + "seed": 42, + "outPath": f"{media_root}/synth/out.mkv", + **extra, + } + + +@pytest.mark.parametrize("field", ["targetFace", "targetImage", "faceImagePath", "referencePath", "image"]) +def test_synthesize_identity_rejects_any_target_face_field( + client: TestClient, media_root: Path, field: str +) -> None: + response = client.post("/v1/synthesizeIdentity", headers=AUTH, json=_synth(media_root, **{field: "data"})) + assert response.status_code == 400 + error = response.json()["error"] + assert error["code"] == "bad_request" + assert error["operation"] == "synthesizeIdentity" + + +def test_synthesize_identity_is_unsupported_without_licensed_model( + client: TestClient, media_root: Path +) -> None: + response = client.post("/v1/synthesizeIdentity", headers=AUTH, json=_synth(media_root)) + assert response.status_code == 501 + assert response.json()["error"] == { + "code": "unsupported", + "message": "no commercially licensed model", + "operation": "synthesizeIdentity", + } + assert not (media_root / "synth" / "out.mkv").exists() + + +@pytest.mark.parametrize( + ("operation", "body"), + [ + ("detectScenes", {"path": "/x", "extra": 1}), + ("detectFaces", {"path": "/x"}), + ("detectFaces", {"path": "/x", "sampleFps": 0}), + ("detectFaces", {"path": "/x", "sampleFps": "2"}), + ("embedFrames", {"path": "/x", "timesMs": [1.5]}), + ("embedFrames", {"path": "/x", "timesMs": [-1]}), + ("segmentPrompted", {"path": "/x", "prompt": {"tMs": 0}, "propagate": True, "outDir": "/y"}), + ( + "segmentPrompted", + { + "path": "/x", + "prompt": {"tMs": 0, "box": {"x": 2, "y": 0, "w": 1, "h": 1}}, + "propagate": True, + "outDir": "/y", + }, + ), + ("redact", {"path": "/x", "outPath": "/y", "redactions": []}), + ( + "redact", + { + "path": "/x", + "outPath": "/y", + "redactions": [{"keyframes": [], "mode": "blur", "paddingPct": 0, "featherPct": 0}], + }, + ), + ( + "redact", + { + "path": "/x", + "outPath": "/y", + "redactions": [ + { + "keyframes": [{"tMs": 0, "box": {"x": 0, "y": 0, "w": 1, "h": 1}}], + "mode": "reframe_out", + "paddingPct": 0, + "featherPct": 0, + } + ], + }, + ), + ( + "redact", + { + "path": "/x", + "outPath": "/y", + "redactions": [ + { + "keyframes": [{"tMs": 0, "box": {"x": 0, "y": 0, "w": 1, "h": 1}}], + "mode": "solid", + "paddingPct": 0, + "featherPct": 0, + "color": "red", + } + ], + }, + ), + ( + "redact", + { + "path": "/x", + "outPath": "/y", + "redactions": [ + { + "keyframes": [{"tMs": 0, "box": {"x": 0, "y": 0, "w": 1, "h": 1, "z": 1}}], + "mode": "blur", + "paddingPct": 0, + "featherPct": 0, + } + ], + }, + ), + ], +) +def test_schema_violations_are_bad_requests( + client: TestClient, operation: str, body: dict[str, object] +) -> None: + response = client.post(f"/v1/{operation}", headers=AUTH, json=body) + assert response.status_code == 400, response.text + assert response.json()["error"]["code"] == "bad_request" + assert response.json()["error"]["operation"] == operation + + +def test_malformed_json_and_unknown_operation(client: TestClient) -> None: + response = client.post( + "/v1/detectScenes", headers={**AUTH, "content-type": "application/json"}, content=b"{nope" + ) + assert response.status_code == 400 + response = client.post("/v1/makeDeepfake", headers=AUTH, json={}) + assert response.status_code == 404 + assert response.json()["error"]["code"] == "bad_request"