evOo-local/infra/media-worker/tests/test_ml_ops.py

"""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