merge: W1-sidecar local media inference sidecar
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0138YqDj3ttfayRwuMKKrqVj
This commit is contained in:
Claude
commit
bd5f637501
7
infra/media-worker/.dockerignore
Normal file
7
infra/media-worker/.dockerignore
Normal file
@ -0,0 +1,7 @@
|
||||
.pytest_cache
|
||||
.ruff_cache
|
||||
.mypy_cache
|
||||
.venv
|
||||
.venv-ml
|
||||
**/__pycache__
|
||||
**/*.pyc
|
||||
5
infra/media-worker/.gitignore
vendored
Normal file
5
infra/media-worker/.gitignore
vendored
Normal file
@ -0,0 +1,5 @@
|
||||
.pytest_cache/
|
||||
.venv/
|
||||
.venv
|
||||
.venv-ml
|
||||
__pycache__/
|
||||
68
infra/media-worker/Dockerfile
Normal file
68
infra/media-worker/Dockerfile
Normal file
@ -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=<ca.pem>.
|
||||
|
||||
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
|
||||
53
infra/media-worker/LICENSES.md
Normal file
53
infra/media-worker/LICENSES.md
Normal file
@ -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 |
|
||||
74
infra/media-worker/README.md
Normal file
74
infra/media-worker/README.md
Normal file
@ -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/<operation>` 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/<backend>/`, 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.
|
||||
3
infra/media-worker/app/__init__.py
Normal file
3
infra/media-worker/app/__init__.py
Normal file
@ -0,0 +1,3 @@
|
||||
"""Local Media Studio inference sidecar."""
|
||||
|
||||
VERSION = "0.1.0"
|
||||
35
infra/media-worker/app/__main__.py
Normal file
35
infra/media-worker/app/__main__.py
Normal file
@ -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))
|
||||
22
infra/media-worker/app/auth.py
Normal file
22
infra/media-worker/app/auth.py
Normal file
@ -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")
|
||||
80
infra/media-worker/app/config.py
Normal file
80
infra/media-worker/app/config.py
Normal file
@ -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,
|
||||
)
|
||||
41
infra/media-worker/app/encoders.py
Normal file
41
infra/media-worker/app/encoders.py
Normal file
@ -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))
|
||||
62
infra/media-worker/app/errors.py
Normal file
62
infra/media-worker/app/errors.py
Normal file
@ -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}},
|
||||
)
|
||||
153
infra/media-worker/app/main.py
Normal file
153
infra/media-worker/app/main.py
Normal file
@ -0,0 +1,153 @@
|
||||
"""HTTP API v1 of the Media Studio sidecar.
|
||||
|
||||
GET /healthz is open; everything else needs the bearer token. POST
|
||||
/v1/<operation> 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
|
||||
435
infra/media-worker/app/media_io.py
Normal file
435
infra/media-worker/app/media_io.py
Normal file
@ -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)
|
||||
506
infra/media-worker/app/models.py
Normal file
506
infra/media-worker/app/models.py
Normal file
@ -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/<backend id>/ 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)
|
||||
1
infra/media-worker/app/ops/__init__.py
Normal file
1
infra/media-worker/app/ops/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
"""One module per operation family. Each exposes `run(ctx, body, cancel)`."""
|
||||
163
infra/media-worker/app/ops/audio.py
Normal file
163
infra/media-worker/app/ops/audio.py
Normal file
@ -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="<f4").astype(np.float32)
|
||||
if proc.wait() != 0:
|
||||
raise MediaError("internal", "audio decode failed")
|
||||
finally:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
|
||||
|
||||
def run(ctx: Ctx, body: PathRequest, cancel: threading.Event) -> 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,
|
||||
}
|
||||
62
infra/media-worker/app/ops/common.py
Normal file
62
infra/media-worker/app/ops/common.py
Normal file
@ -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
|
||||
61
infra/media-worker/app/ops/devices.py
Normal file
61
infra/media-worker/app/ops/devices.py
Normal file
@ -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
|
||||
147
infra/media-worker/app/ops/embeddings.py
Normal file
147
infra/media-worker/app/ops/embeddings.py
Normal file
@ -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)],
|
||||
}
|
||||
265
infra/media-worker/app/ops/faces.py
Normal file
265
infra/media-worker/app/ops/faces.py
Normal file
@ -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}
|
||||
163
infra/media-worker/app/ops/inpaint.py
Normal file
163
infra/media-worker/app/ops/inpaint.py
Normal file
@ -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()}
|
||||
54
infra/media-worker/app/ops/masks.py
Normal file
54
infra/media-worker/app/ops/masks.py
Normal file
@ -0,0 +1,54 @@
|
||||
"""Per-frame mask directories (`<tMs>.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 <tMs>.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
|
||||
271
infra/media-worker/app/ops/redact.py
Normal file
271
infra/media-worker/app/ops/redact.py
Normal file
@ -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}
|
||||
58
infra/media-worker/app/ops/saliency.py
Normal file
58
infra/media-worker/app/ops/saliency.py
Normal file
@ -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}
|
||||
103
infra/media-worker/app/ops/scenes.py
Normal file
103
infra/media-worker/app/ops/scenes.py
Normal file
@ -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}
|
||||
244
infra/media-worker/app/ops/segment.py
Normal file
244
infra/media-worker/app/ops/segment.py
Normal file
@ -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 `<tMs>.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}
|
||||
20
infra/media-worker/app/ops/synthesize.py
Normal file
20
infra/media-worker/app/ops/synthesize.py
Normal file
@ -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")
|
||||
57
infra/media-worker/app/ops/transcribe.py
Normal file
57
infra/media-worker/app/ops/transcribe.py
Normal file
@ -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"]]}
|
||||
75
infra/media-worker/app/paths.py
Normal file
75
infra/media-worker/app/paths.py
Normal file
@ -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
|
||||
93
infra/media-worker/app/runtime.py
Normal file
93
infra/media-worker/app/runtime.py
Normal file
@ -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
|
||||
124
infra/media-worker/app/schemas.py
Normal file
124
infra/media-worker/app/schemas.py
Normal file
@ -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
|
||||
5
infra/media-worker/pytest.ini
Normal file
5
infra/media-worker/pytest.ini
Normal file
@ -0,0 +1,5 @@
|
||||
[pytest]
|
||||
testpaths = tests
|
||||
addopts = -p no:cacheprovider
|
||||
filterwarnings =
|
||||
ignore:Using `httpx` with `starlette.testclient` is deprecated
|
||||
3
infra/media-worker/requirements-dev.txt
Normal file
3
infra/media-worker/requirements-dev.txt
Normal file
@ -0,0 +1,3 @@
|
||||
pytest==9.1.1
|
||||
httpx==0.28.1
|
||||
ruff==0.16.10
|
||||
9
infra/media-worker/requirements-ml-cpu.txt
Normal file
9
infra/media-worker/requirements-ml-cpu.txt
Normal file
@ -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
|
||||
9
infra/media-worker/requirements-ml-cuda.txt
Normal file
9
infra/media-worker/requirements-ml-cuda.txt
Normal file
@ -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
|
||||
11
infra/media-worker/requirements-ml-rocm.txt
Normal file
11
infra/media-worker/requirements-ml-rocm.txt
Normal file
@ -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
|
||||
10
infra/media-worker/requirements.txt
Normal file
10
infra/media-worker/requirements.txt
Normal file
@ -0,0 +1,10 @@
|
||||
# Core: HTTP API, ffmpeg-driven operations, OpenCV, scene detection, audio features.
|
||||
# Accelerator-specific ML wheels live in requirements-ml-<accelerator>.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
|
||||
5
infra/media-worker/ruff.toml
Normal file
5
infra/media-worker/ruff.toml
Normal file
@ -0,0 +1,5 @@
|
||||
line-length = 110
|
||||
target-version = "py311"
|
||||
|
||||
[lint]
|
||||
select = ["E", "F", "W", "B", "I", "UP", "SIM"]
|
||||
0
infra/media-worker/tests/__init__.py
Normal file
0
infra/media-worker/tests/__init__.py
Normal file
224
infra/media-worker/tests/conftest.py
Normal file
224
infra/media-worker/tests/conftest.py
Normal file
@ -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()
|
||||
133
infra/media-worker/tests/fakes.py
Normal file
133
infra/media-worker/tests/fakes.py
Normal file
@ -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),
|
||||
}
|
||||
21
infra/media-worker/tests/fixtures/media-capabilities.json
vendored
Normal file
21
infra/media-worker/tests/fixtures/media-capabilities.json
vendored
Normal file
@ -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"]
|
||||
}
|
||||
88
infra/media-worker/tests/test_analysis.py
Normal file
88
infra/media-worker/tests/test_analysis.py
Normal file
@ -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
|
||||
58
infra/media-worker/tests/test_auth.py
Normal file
58
infra/media-worker/tests/test_auth.py
Normal file
@ -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
|
||||
82
infra/media-worker/tests/test_capabilities.py
Normal file
82
infra/media-worker/tests/test_capabilities.py
Normal file
@ -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
|
||||
170
infra/media-worker/tests/test_ml_ops.py
Normal file
170
infra/media-worker/tests/test_ml_ops.py
Normal file
@ -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
|
||||
75
infra/media-worker/tests/test_models.py
Normal file
75
infra/media-worker/tests/test_models.py
Normal file
@ -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
|
||||
126
infra/media-worker/tests/test_paths.py
Normal file
126
infra/media-worker/tests/test_paths.py
Normal file
@ -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
|
||||
174
infra/media-worker/tests/test_redact.py
Normal file
174
infra/media-worker/tests/test_redact.py
Normal file
@ -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())
|
||||
117
infra/media-worker/tests/test_runtime.py
Normal file
117
infra/media-worker/tests/test_runtime.py
Normal file
@ -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]
|
||||
137
infra/media-worker/tests/test_schema.py
Normal file
137
infra/media-worker/tests/test_schema.py
Normal file
@ -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"
|
||||
Loading…
x
Reference in New Issue
Block a user