Files
imagefind/scripts/live-search-acceptance.py

234 lines
8.1 KiB
Python

#!/usr/bin/env python3
"""Verify the six recognition filters against run-owned live videos."""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
from typing import Any
import httpx
def call(client: httpx.Client, method: str, path: str, **kwargs: Any) -> Any:
response = client.request(method, path, **kwargs)
if response.is_error:
raise RuntimeError(f"{method} {path} -> {response.status_code}: {response.text[:1000]}")
return response.json()
def find_owned(videos: list[dict[str, Any]], source_id: str, filename: str) -> dict[str, Any]:
match = next(
(
item
for item in videos
if item.get("source_id") == source_id
and str(item.get("source_key") or "").rsplit("/", 1)[-1] == filename
),
None,
)
if match is None:
raise AssertionError(f"run-owned search fixture is missing: {filename}")
return match
def search(
client: httpx.Client,
*,
text: str,
recognition_type: str,
expected_video_id: str | None,
expected_source: str,
forbidden_video_id: str | None = None,
) -> dict[str, Any]:
result = call(
client,
"POST",
"/api/v1/search",
json={"text": text, "recognition_types": [recognition_type], "limit": 100},
)
items = result.get("items", [])
ids = [item["video_id"] for item in items]
if expected_video_id is not None:
if expected_video_id not in ids:
raise AssertionError(f"{recognition_type} search omitted expected video for {text!r}")
matched = next(item for item in items if item["video_id"] == expected_video_id)
if expected_source not in matched.get("match_sources", []):
raise AssertionError(
f"{recognition_type} search returned the target without {expected_source!r}: "
f"{matched.get('match_sources')}"
)
if forbidden_video_id is not None and forbidden_video_id in ids:
raise AssertionError(f"{recognition_type} search returned a forbidden video for {text!r}")
return {
"query": text,
"result_count": len(items),
"matched_expected": expected_video_id in ids if expected_video_id else not items,
"expected_rank": ids.index(expected_video_id) + 1 if expected_video_id in ids else None,
"match_sources": sorted(
{
source
for item in items
if expected_video_id is None or item["video_id"] == expected_video_id
for source in item.get("match_sources", [])
}
),
}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--base-url", required=True)
parser.add_argument("--token-file", type=Path, required=True)
parser.add_argument("--run-dir", type=Path, required=True)
parser.add_argument("--timeout", type=int, default=600)
args = parser.parse_args()
state = json.loads((args.run_dir / "state.json").read_text())
client = httpx.Client(
base_url=args.base_url.rstrip("/"),
headers={"Authorization": f"Bearer {args.token_file.read_text().strip()}"},
timeout=httpx.Timeout(60, connect=10),
)
run_id = state["run_id"]
source_id = state["source_id"]
expected_names = {
"positive": f"{run_id}-e2e-positive.mp4",
"negative": f"{run_id}-e2e-negative.mp4",
"subtitle": f"{run_id}-subtitle-e2e-subtitle.mp4",
"audio": f"{run_id}-e2e-speech.mp4",
}
deadline = time.monotonic() + args.timeout
while True:
videos = call(client, "GET", "/api/v1/videos?limit=500")
fixtures = {name: find_owned(videos, source_id, filename) for name, filename in expected_names.items()}
terminal = all(
item.get("index_state", {}).get("basic") in {"ready", "failed"}
and item.get("index_state", {}).get("visual") in {"ready", "failed", "skipped"}
and item.get("index_state", {}).get("ocr") in {"ready", "failed", "skipped"}
for item in fixtures.values()
)
if terminal:
break
if time.monotonic() >= deadline:
raise TimeoutError("search fixtures did not reach terminal basic/visual/OCR states")
time.sleep(2)
positive = fixtures["positive"]["id"]
negative = fixtures["negative"]["id"]
subtitle = fixtures["subtitle"]["id"]
audio = fixtures["audio"]["id"]
report: dict[str, Any] = {
"fixtures": {name: item["id"] for name, item in fixtures.items()},
"index_state": {name: item.get("index_state") for name, item in fixtures.items()},
"recognition": {},
}
report["recognition"]["visual_positive"] = search(
client,
text="海边日落与蓝色海洋",
recognition_type="visual",
expected_video_id=positive,
expected_source="semantic",
)
report["recognition"]["visual_negative_fixture"] = search(
client,
text="绿色山脉与森林",
recognition_type="visual",
expected_video_id=negative,
expected_source="semantic",
)
report["recognition"]["ocr_positive"] = search(
client,
text="星河测试 729",
recognition_type="ocr",
expected_video_id=positive,
expected_source="ocr",
forbidden_video_id=negative,
)
report["recognition"]["ocr_negative_fixture"] = search(
client,
text="山谷样本 314",
recognition_type="ocr",
expected_video_id=negative,
expected_source="ocr",
forbidden_video_id=positive,
)
report["recognition"]["subtitle_positive"] = search(
client,
text="字幕验证 星河测试729",
recognition_type="subtitle",
expected_video_id=subtitle,
expected_source="subtitle",
forbidden_video_id=negative,
)
report["recognition"]["subtitle_non_match"] = search(
client,
text="quartz zeppelin 9834721",
recognition_type="subtitle",
expected_video_id=None,
expected_source="subtitle",
forbidden_video_id=subtitle,
)
report["recognition"]["metadata_positive"] = search(
client,
text="IF-E2E-729",
recognition_type="metadata",
expected_video_id=positive,
expected_source="metadata",
forbidden_video_id=negative,
)
report["recognition"]["metadata_non_match"] = search(
client,
text="不存在的资料验收词 9834721",
recognition_type="metadata",
expected_video_id=None,
expected_source="metadata",
forbidden_video_id=positive,
)
report["recognition"]["audio_positive"] = search(
client,
text="image",
recognition_type="audio",
expected_video_id=audio,
expected_source="audio",
forbidden_video_id=negative,
)
report["recognition"]["audio_non_match"] = search(
client,
text="pineapple submarine 9834721",
recognition_type="audio",
expected_video_id=None,
expected_source="audio",
forbidden_video_id=audio,
)
face_report_path = args.run_dir / "face-report.json"
if face_report_path.exists():
face = json.loads(face_report_path.read_text())
report["recognition"]["person_positive"] = search(
client,
text=face["person_name"],
recognition_type="person",
expected_video_id=face["video_id"],
expected_source="person",
)
report["recognition"]["person_non_match"] = search(
client,
text=f"{run_id}-不存在人物-9834721",
recognition_type="person",
expected_video_id=None,
expected_source="person",
forbidden_video_id=face["video_id"],
)
else:
report["recognition"]["person"] = {"status": "pending face acceptance"}
output = args.run_dir / "search-report.json"
output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n")
print(json.dumps(report, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()