From afbf18996ffb4625d42f677d51bd9c2c28454af3 Mon Sep 17 00:00:00 2001 From: iChubai <149308795+iChubai@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:46:20 +0800 Subject: [PATCH 1/2] [feat]Integrate WorldMemBench evaluation Signed-off-by: iChubai <149308795+iChubai@users.noreply.github.com> --- Makefile | 1 + .../docs/evaluation/benchmark-hub/meta.json | 1 + .../evaluation/benchmark-hub/meta.zh.json | 1 + .../benchmark-hub/worldmembench.mdx | 103 +++++ .../benchmark-hub/worldmembench.zh.mdx | 90 ++++ .../lib/benchmark-catalog-status.json | 21 + docs/fumadocs/lib/benchmark-page-data.json | 177 +++++++- docs/fumadocs/lib/catalog-coverage-data.json | 11 +- .../test_worldmembench_official_runner.py | 404 ++++++++++++++++++ .../catalog/video/worldmembench.yaml | 196 +++++++++ .../official/worldmembench.yaml | 67 +++ .../tasks/external/worldmembench.yaml | 172 ++++++++ .../evaluation/tasks/contracts/external.py | 93 +++- .../execution/framework/official_runner.py | 37 +- .../run_worldmembench_official_runner.py | 220 ++++++++++ .../worldmembench/worldmembench_metrics.py | 107 +++++ 16 files changed, 1660 insertions(+), 41 deletions(-) create mode 100644 docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.mdx create mode 100644 docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.zh.mdx create mode 100644 tests/eval_core/test_worldmembench_official_runner.py create mode 100644 worldfoundry/data/benchmarks/catalog/video/worldmembench.yaml create mode 100644 worldfoundry/data/benchmarks/runtime_profiles/official/worldmembench.yaml create mode 100644 worldfoundry/data/benchmarks/tasks/external/worldmembench.yaml create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/worldmembench_metrics.py diff --git a/Makefile b/Makefile index 2a6c4c268..9b1f67968 100644 --- a/Makefile +++ b/Makefile @@ -223,6 +223,7 @@ INFER_TENSOR_CONTRACTS = \ tests/eval_core/test_catalog_integration_status.py \ $(EVALUATION_CPU_CONTRACTS) EVALUATION_CPU_CONTRACTS = \ + tests/eval_core/test_worldmembench_official_runner.py \ tests/evaluation/execution/test_resume_options.py \ tests/evaluation/execution/test_result_integrity.py \ tests/evaluation/execution/test_helper_reuse.py \ diff --git a/docs/fumadocs/content/docs/evaluation/benchmark-hub/meta.json b/docs/fumadocs/content/docs/evaluation/benchmark-hub/meta.json index 47df0571e..e969c5e00 100644 --- a/docs/fumadocs/content/docs/evaluation/benchmark-hub/meta.json +++ b/docs/fumadocs/content/docs/evaluation/benchmark-hub/meta.json @@ -63,6 +63,7 @@ "rbench", "---World Models---", "worldmodelbench", + "worldmembench", "worldscore", "worldbench", "wbench", diff --git a/docs/fumadocs/content/docs/evaluation/benchmark-hub/meta.zh.json b/docs/fumadocs/content/docs/evaluation/benchmark-hub/meta.zh.json index 47df0571e..e969c5e00 100644 --- a/docs/fumadocs/content/docs/evaluation/benchmark-hub/meta.zh.json +++ b/docs/fumadocs/content/docs/evaluation/benchmark-hub/meta.zh.json @@ -63,6 +63,7 @@ "rbench", "---World Models---", "worldmodelbench", + "worldmembench", "worldscore", "worldbench", "wbench", diff --git a/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.mdx b/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.mdx new file mode 100644 index 000000000..479ec30aa --- /dev/null +++ b/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.mdx @@ -0,0 +1,103 @@ +--- +title: WorldMemBench +description: Official video-world-model memory metrics, per-model result import, and upstream runtime execution. +--- + +WorldMemBench evaluates video world models through human, object, environment, +and dynamic memory, alongside video quality and camera following. WorldFoundry +exposes all **12 official metrics**, imports their per-model summaries, and can +invoke the official metric runtime on staged videos. Model generation remains +separate from evaluation. + +Code: [ZJU-REAL/WorldMemBench](https://github.com/ZJU-REAL/WorldMemBench) (MIT). +Paper: [WorldMemBench](https://arxiv.org/abs/2610.19063). +Data: [rescuerz/WorldMemBench](https://huggingface.co/datasets/rescuerz/WorldMemBench). +The integration targets upstream revision `1b79840143a98e017891e893b5d3ef75b93078a3`. + + + +## Import official results + +Result import needs no GPU or metric checkpoints: + +```bash +worldfoundry-eval zoo benchmark-run \ + --benchmark-id worldmembench \ + --mode official-validation \ + --official-results-path /path/to/official_eval_output \ + --output-dir tmp/worldmembench/import \ + --json +``` + +The input may be the official evaluation output root, its `metrics` directory, +or one `metrics//summary.json` file. `cases.jsonl` files are imported +when present. A directory containing several model groups needs an explicit +selection through the runner: + +```bash +python -m worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner \ + --official-results-path /path/to/official_eval_output \ + --model-id model_a \ + --metrics human_identity,object_geometry,camera_rotation \ + --output-dir tmp/worldmembench/model_a \ + --json +``` + +Outputs are `scorecard.json`, `raw_metric_table.jsonl`, and +`per_sample_scores.jsonl`. The scorecard retains official score fields, grouped +category/trajectory results, sample counts, and case statuses. Models are not +averaged together. `camera_rotation` and `camera_translation` are errors: +**lower is better**. Metric values retain upstream units, including signed depth +correlations and unbounded camera errors; no clipping, inversion, or +cross-metric composite is applied. + +## Execute official metrics + +Provision the official checkout and selected metric dependencies following its +[installation instructions](https://github.com/ZJU-REAL/WorldMemBench#installation). +Use the unified WorldFoundry environment for orchestration; if upstream needs a +separate interpreter, select it with `--python` rather than replacing the +unified environment's Torch/CUDA packages. + +```bash +python -m worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner \ + --run-official \ + --repo-root /path/to/WorldMemBench \ + --python /path/to/upstream-env/bin/python \ + --input-path /path/to/manifest.jsonl \ + --config /path/to/worldmembench.yaml \ + --weights-dir /path/to/weights \ + --metrics imaging_quality,aesthetic_quality \ + --device cuda:0 \ + --output-dir tmp/worldmembench/official \ + --json +``` + +`--repo-root` loads the checkout's `src/worldmembench`; omit it when the selected +interpreter already has `worldmembench` installed. Input can also be a video or +an official video tree: `////video.mp4`. +Use upstream JSONL manifests to retain model IDs, camera protocols, and per-case +annotations. `--expected-revisits 1|2|3` selects the fixed human/object revisit +protocol; omitting it preserves upstream manifest/detection behavior. +`--gpus 0,1` delegates process sharding to upstream instead of using `--device`. +Each execution needs a fresh output directory; official results go under +`official/`, with stdout/stderr logs beside the scorecard. + +The unified CLI can also score staged videos using `--generated-artifact-dir` +and `--mode official-run`. Configure runtime paths through +`WORLDFOUNDRY_WORLDMEMBENCH_ROOT`, `WORLDFOUNDRY_WORLDMEMBENCH_PYTHON`, +`WORLDFOUNDRY_WORLDMEMBENCH_CONFIG`, and +`WORLDFOUNDRY_WORLDMEMBENCH_WEIGHTS_DIR`. For import, use +`WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH`. + +## Limitations + +Official execution needs the selected metric checkpoints, DA3/ViPE subprocess +environments, and a VLM endpoint for Dynamic Memory. Keep backend interpreters, +checkpoint paths, and VLM settings in the upstream YAML config; API keys stay +in the upstream-configured environment variable. Metric implementations and +weights are not vendored. GPU/checkpoint and complete-dataset replay remain +pending, so scorecards do not claim full-suite or leaderboard validity. +`model_failure` retains its official scored result; `not_applicable` is not +zero-filled. `eval_failed` or `data_invalid` cases retain their coverage and +make the runner exit nonzero, including when partial scores are available. diff --git a/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.zh.mdx b/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.zh.mdx new file mode 100644 index 000000000..c415ee297 --- /dev/null +++ b/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.zh.mdx @@ -0,0 +1,90 @@ +--- +title: WorldMemBench +description: 视频世界模型记忆评测:官方指标、按模型导入结果及调用上游评测运行时。 +--- + +WorldMemBench 从人物、物体、环境和动态记忆,以及视频质量与相机跟随能力评测视频世界模型。 +WorldFoundry 提供 **12 项官方指标**的统一入口,支持按模型导入官方结果,或调用官方指标运行时 +评测已生成的视频。模型生成与评测分开执行。 + +代码:[ZJU-REAL/WorldMemBench](https://github.com/ZJU-REAL/WorldMemBench)(MIT)。 +论文:[WorldMemBench](https://arxiv.org/abs/2610.19063)。 +数据:[rescuerz/WorldMemBench](https://huggingface.co/datasets/rescuerz/WorldMemBench)。 +集成对应上游版本 `1b79840143a98e017891e893b5d3ef75b93078a3`。 + + + +## 导入官方结果 + +结果导入不需要 GPU 或指标权重: + +```bash +worldfoundry-eval zoo benchmark-run \ + --benchmark-id worldmembench \ + --mode official-validation \ + --official-results-path /path/to/official_eval_output \ + --output-dir tmp/worldmembench/import \ + --json +``` + +输入可以是官方评测输出根目录、其中的 `metrics` 目录,或单个 +`metrics//summary.json`。存在 `cases.jsonl` 时同时导入逐样本结果。 +多模型结果必须通过 runner 显式选择一个模型: + +```bash +python -m worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner \ + --official-results-path /path/to/official_eval_output \ + --model-id model_a \ + --metrics human_identity,object_geometry,camera_rotation \ + --output-dir tmp/worldmembench/model_a \ + --json +``` + +输出为 `scorecard.json`、`raw_metric_table.jsonl` 和 `per_sample_scores.jsonl`。 +Scorecard 保留官方分数字段、类别/轨迹分组、样本数与状态,不对不同模型求平均。 +`camera_rotation` 和 `camera_translation` 是误差,**越低越好**。 +所有数值沿用官方单位,包括可为负的深度相关系数和无上界的相机误差;不裁剪、取倒数, +也不构造跨指标总分。 + +## 执行官方指标 + +先按上游的[安装说明](https://github.com/ZJU-REAL/WorldMemBench#installation) +准备官方代码与所选指标依赖。WorldFoundry 调度使用统一环境;上游需要独立解释器时通过 +`--python` 指定,不替换统一环境中的 Torch/CUDA。 + +```bash +python -m worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner \ + --run-official \ + --repo-root /path/to/WorldMemBench \ + --python /path/to/upstream-env/bin/python \ + --input-path /path/to/manifest.jsonl \ + --config /path/to/worldmembench.yaml \ + --weights-dir /path/to/weights \ + --metrics imaging_quality,aesthetic_quality \ + --device cuda:0 \ + --output-dir tmp/worldmembench/official \ + --json +``` + +`--repo-root` 使用该 checkout 的 `src/worldmembench`;指定解释器已安装 +`worldmembench` 时可省略。输入也可以是单个视频,或官方视频树: +`////video.mp4`。使用上游 JSONL manifest +可保留模型 ID、相机协议与逐样本标注。`--expected-revisits 1|2|3` 指定人物/物体的固定重访 +协议;省略时沿用官方 manifest 或自动检测行为。`--gpus 0,1` 将进程分片交给上游, +代替 `--device`。每次执行使用新的输出目录;官方结果位于 `official/`,标准输出与错误日志 +位于 scorecard 同级。 + +统一 CLI 也支持通过 `--generated-artifact-dir` 与 `--mode official-run` 评测已有视频。 +运行时路径使用 `WORLDFOUNDRY_WORLDMEMBENCH_ROOT`、 +`WORLDFOUNDRY_WORLDMEMBENCH_PYTHON`、`WORLDFOUNDRY_WORLDMEMBENCH_CONFIG` 和 +`WORLDFOUNDRY_WORLDMEMBENCH_WEIGHTS_DIR` 配置。 +结果导入可使用 `WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH`。 + +## 限制 + +官方计算需要所选指标权重、DA3/ViPE 子进程环境,以及 Dynamic Memory 的 VLM 服务。 +后端解释器、权重路径与 VLM 设置放在上游 YAML 中;API 密钥使用上游配置的环境变量。 +本仓库不复制指标实现或权重。GPU/真实权重与完整数据集回放尚待验证,因此 scorecard +不声明完整评测或排行榜有效性。`model_failure` 保留官方计分结果;`not_applicable` +不补零。`eval_failed` 或 `data_invalid` 保留覆盖率,且即使存在部分分数,runner +也返回非零退出码。 diff --git a/docs/fumadocs/lib/benchmark-catalog-status.json b/docs/fumadocs/lib/benchmark-catalog-status.json index 25ac9abd4..f63c0b1ba 100644 --- a/docs/fumadocs/lib/benchmark-catalog-status.json +++ b/docs/fumadocs/lib/benchmark-catalog-status.json @@ -1810,6 +1810,27 @@ "summaryZh": "仓内 evaluator 按 task protocol 评估 artifact bundle", "verificationStatus": "bounded_fixture_verified" }, + "worldmembench": { + "aliases": [], + "badges": ["integrated"], + "category": "World Models", + "categoryZh": "世界模型", + "domains": ["human-memory", "object-memory", "environment-memory", "dynamic-memory"], + "integrationStatus": "integrated", + "leaderboardValid": false, + "metrics": [ + "imaging_quality", "aesthetic_quality", "human_identity", "human_appearance", + "object_geometry", "object_appearance", "environment_visible_memory", + "environment_depth_fidelity", "environment_scene_reconstruction", "dynamic_memory", + "camera_rotation", "camera_translation" + ], + "name": "WorldMemBench", + "requiresApi": false, + "runtimeKind": "external_official_repo", + "summary": "Imports per-model memory metrics or invokes the official WorldMemBench runtime", + "summaryZh": "按模型导入记忆指标或调用官方 WorldMemBench 运行时", + "verificationStatus": "official_runtime_ready_real_data_validation_pending" + }, "worldmodelbench": { "aliases": [], "badges": [ diff --git a/docs/fumadocs/lib/benchmark-page-data.json b/docs/fumadocs/lib/benchmark-page-data.json index 99cfaf864..f2eb5b521 100644 --- a/docs/fumadocs/lib/benchmark-page-data.json +++ b/docs/fumadocs/lib/benchmark-page-data.json @@ -1,8 +1,8 @@ { "generatedFrom": "worldfoundry/data/benchmarks/catalog + tasks/external + recorded About facts + curated docs overrides", - "benchmarkCount": 74, + "benchmarkCount": 75, "leaderboardIngestedIds": [], - "leaderboardPlaceholderCount": 74, + "leaderboardPlaceholderCount": 75, "pages": { "ai2thor": { "id": "ai2thor", @@ -10324,6 +10324,179 @@ "entries": [] } }, + "worldmembench": { + "id": "worldmembench", + "name": "WorldMemBench", + "projectPage": "", + "paperUrl": "", + "dimensions": [ + { + "id": "imaging_quality", + "name": "Imaging Quality", + "nameZh": "Imaging Quality", + "kind": "axis", + "description": "Official MUSIQ-SPAQ frame scores, normalized upstream and averaged equally per video.", + "descriptionZh": "" + }, + { + "id": "aesthetic_quality", + "name": "Aesthetic Quality", + "nameZh": "Aesthetic Quality", + "kind": "axis", + "description": "Official CLIP/LAION aesthetic frame scores, normalized upstream and averaged per video.", + "descriptionZh": "" + }, + { + "id": "environment_scene_reconstruction", + "name": "Environment Scene Reconstruction", + "nameZh": "Environment Scene Reconstruction", + "kind": "axis", + "description": "Official scene reconstruction score at departure and revisit frames.", + "descriptionZh": "" + } + ], + "data": { + "summary": "", + "summaryZh": "", + "sizeLabel": "", + "sizeLabelZh": "", + "count": null, + "unit": "", + "sources": [ + { + "kind": "github", + "id": "https://github.com/ZJU-REAL/WorldMemBench", + "license": "" + } + ], + "splits": [], + "modalities": [ + "video", + "text" + ], + "license": "", + "notes": [], + "notesZh": [], + "notApplicable": false + }, + "metrics": [ + { + "id": "imaging_quality", + "name": "Imaging Quality", + "description": "Official MUSIQ-SPAQ frame scores, normalized upstream and averaged equally per video.", + "descriptionZh": "", + "higherIsBetter": true, + "primary": false, + "leaderboardKey": "imaging_quality" + }, + { + "id": "aesthetic_quality", + "name": "Aesthetic Quality", + "description": "Official CLIP/LAION aesthetic frame scores, normalized upstream and averaged per video.", + "descriptionZh": "", + "higherIsBetter": true, + "primary": false, + "leaderboardKey": "aesthetic_quality" + }, + { + "id": "human_identity", + "name": "Human Identity", + "description": "Human identity retention across departure and revisit segments.", + "descriptionZh": "", + "higherIsBetter": true, + "primary": false, + "leaderboardKey": "human_identity" + }, + { + "id": "human_appearance", + "name": "Human Appearance", + "description": "Human appearance retention under the official revisit protocol.", + "descriptionZh": "", + "higherIsBetter": true, + "primary": false, + "leaderboardKey": "human_appearance" + }, + { + "id": "object_geometry", + "name": "Object Geometry", + "description": "Object geometry retention under the official revisit protocol.", + "descriptionZh": "", + "higherIsBetter": true, + "primary": false, + "leaderboardKey": "object_geometry" + }, + { + "id": "object_appearance", + "name": "Object Appearance", + "description": "Object appearance retention under the official revisit protocol.", + "descriptionZh": "", + "higherIsBetter": true, + "primary": false, + "leaderboardKey": "object_appearance" + }, + { + "id": "environment_visible_memory", + "name": "Environment Visible Memory", + "description": "DINOv2 full-frame similarity at camera return frames.", + "descriptionZh": "", + "higherIsBetter": true, + "primary": false, + "leaderboardKey": "environment_visible_memory" + }, + { + "id": "environment_depth_fidelity", + "name": "Environment Depth Fidelity", + "description": "Pixel-aligned DA3 depth-order Spearman correlation at return frames.", + "descriptionZh": "", + "higherIsBetter": true, + "primary": false, + "leaderboardKey": "environment_depth_fidelity" + }, + { + "id": "environment_scene_reconstruction", + "name": "Environment Scene Reconstruction", + "description": "Official scene reconstruction score at departure and revisit frames.", + "descriptionZh": "", + "higherIsBetter": true, + "primary": false, + "leaderboardKey": "environment_scene_reconstruction" + }, + { + "id": "dynamic_memory", + "name": "Dynamic Memory", + "description": "Event-gated eight-question VLM checklist score with per-category details retained.", + "descriptionZh": "", + "higherIsBetter": true, + "primary": false, + "leaderboardKey": "dynamic_memory" + }, + { + "id": "camera_rotation", + "name": "Camera Rotation", + "description": "DA3 adjacent rotation RPE normalized by the reference rotation path.", + "descriptionZh": "", + "higherIsBetter": false, + "primary": false, + "leaderboardKey": "camera_rotation" + }, + { + "id": "camera_translation", + "name": "Camera Translation", + "description": "DA3/ViPE path-normalized translation RPE following Sim(3) alignment.", + "descriptionZh": "", + "higherIsBetter": false, + "primary": false, + "leaderboardKey": "camera_translation" + } + ], + "leaderboard": { + "ingested": false, + "source": "", + "sourceZh": "", + "officialUrl": "", + "entries": [] + } + }, "worldmodelbench": { "id": "worldmodelbench", "name": "WorldModelBench", diff --git a/docs/fumadocs/lib/catalog-coverage-data.json b/docs/fumadocs/lib/catalog-coverage-data.json index f9d95434b..bc4a9f4e3 100644 --- a/docs/fumadocs/lib/catalog-coverage-data.json +++ b/docs/fumadocs/lib/catalog-coverage-data.json @@ -1,7 +1,7 @@ { "modelsTotal": 275, "modelsListed": 265, - "benchmarksTotal": 74, + "benchmarksTotal": 75, "modelFamilies": [ { "id": "video", @@ -2214,7 +2214,7 @@ { "id": "video", "label": "Video & world", - "count": 52, + "count": 53, "entries": [ { "id": "4dworldbench", @@ -2545,6 +2545,13 @@ "worldbench" ] }, + { + "id": "worldmembench", + "name": "WorldMemBench", + "aliases": [ + "worldmembench" + ] + }, { "id": "worldmodelbench", "name": "WorldModelBench", diff --git a/tests/eval_core/test_worldmembench_official_runner.py b/tests/eval_core/test_worldmembench_official_runner.py new file mode 100644 index 000000000..80855cc42 --- /dev/null +++ b/tests/eval_core/test_worldmembench_official_runner.py @@ -0,0 +1,404 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 WorldFoundry contributors +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path +from typing import Any + +import pytest +import yaml + +from worldfoundry.evaluation.tasks.catalog.benchmark_catalog import formal_benchmark_ids +from worldfoundry.evaluation.tasks.catalog.dispatch import official_runner_spec +from worldfoundry.evaluation.tasks.catalog.registry import load_task_registry_from_paths +from worldfoundry.evaluation.tasks.contracts.external import get_external_benchmark_contract +from worldfoundry.evaluation.tasks.execution.orchestration.benchmark_runner import ( + build_benchmark_runner_registry, + run_benchmark_execution, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench import run_worldmembench_official_runner as runner +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.worldmembench_metrics import ( + METRIC_SPECS, + load_official_results, +) + +pytestmark = pytest.mark.unit + +REPO_ROOT = Path(__file__).resolve().parents[2] +MODULE = "worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner" + + +def _write_result( + root: Path, + metric: str = "human_identity", + *, + models: dict[str, dict[str, Any]] | None = None, + cases: list[dict[str, Any]] | None = None, +) -> Path: + """Write the compact summary schema emitted by upstream ResultWriter.""" + directory = root / "metrics" / metric + directory.mkdir(parents=True, exist_ok=True) + payload = { + "num_cases": 2, + "results": models or {"model_a": {"num_cases": 2, "scores": {"score": 0.75}}}, + "groups": [{"model_id": "model_a", "category": "human", "trajectory_id": "LRLR", "num_cases": 2}], + } + summary = directory / "summary.json" + summary.write_text(json.dumps(payload), encoding="utf-8") + if cases is not None: + (directory / "cases.jsonl").write_text("".join(json.dumps(case) + "\n" for case in cases), encoding="utf-8") + return summary + + +def _run(tmp_path: Path, *flags: str) -> tuple[dict[str, Any], int]: + args = runner.build_parser().parse_args(["--output-dir", str(tmp_path / "out"), *flags]) + return runner.run(args) + + +def test_catalog_contract_task_runtime_and_dispatch_agree() -> None: + assert "worldmembench" in formal_benchmark_ids() + contract = get_external_benchmark_contract("worldmembench") + assert contract.metric_ids == tuple(METRIC_SPECS) + spec = official_runner_spec("worldmembench") + assert spec is not None + assert spec.module == MODULE + assert spec.supports_official_runtime and spec.artifact_official_run + assert spec.model_arg == "--model-id" + assert not spec.supports_fixture + registry = build_benchmark_runner_registry() + entry = registry.zoo.get("worldmembench") + assert entry.integration_status == "integrated" + assert not entry.official_benchmark_verified + assert not entry.leaderboard_valid + task_path = REPO_ROOT / "worldfoundry/data/benchmarks/tasks/external/worldmembench.yaml" + task = load_task_registry_from_paths(task_path).get("worldmembench_standard", benchmark="worldmembench").task + assert task.metric_ids == contract.metric_ids + assert task.evaluation_protocol_names == ("external_official_runner",) + assert task.metadata["artifact_contract"] + for relative in [ + "tasks/external/worldmembench.yaml", + "runtime_profiles/official/worldmembench.yaml", + ]: + payload = yaml.safe_load((REPO_ROOT / "worldfoundry/data/benchmarks" / relative).read_text()) + if "metadata" in payload: + assert set(payload["metrics"]) == set(contract.metric_ids) + else: + assert payload["benchmark_ids"] == ["worldmembench"] + for item in payload["required_paths"]: + assert (REPO_ROOT / item["path"]).exists() + + +@pytest.mark.parametrize("input_kind", ["root", "metrics", "summary"]) +def test_import_official_summary_layouts(tmp_path: Path, input_kind: str) -> None: + root = tmp_path / "official" + cases = [{"item_id": "case_001", "model_id": "model_a", "status": "ok", "scores": {"score": 0.75}}] + summary = _write_result(root, cases=cases) + source = {"root": root, "metrics": root / "metrics", "summary": summary}[input_kind] + model, extracted, imported_cases = load_official_results(source) + assert model == "model_a" + assert extracted["human_identity"]["raw_score"] == 0.75 + assert extracted["human_identity"]["sample_count"] == 2 + assert imported_cases == [{**cases[0], "metric_id": "human_identity"}] + + +def test_models_are_explicitly_selected_without_pooling(tmp_path: Path) -> None: + _write_result( + tmp_path, + models={ + "model_a": {"num_cases": 1, "scores": {"score": 0.2}}, + "model_b": {"num_cases": 1, "scores": {"score": 0.9}}, + }, + cases=[ + {"item_id": "a", "model_id": "model_a", "scores": {"score": 0.2}}, + {"item_id": "b", "model_id": "model_b", "scores": {"score": 0.9}}, + ], + ) + with pytest.raises(ValueError, match="--model-id"): + load_official_results(tmp_path) + model, metrics, cases = load_official_results(tmp_path, model_id="model_b") + assert model == "model_b" + assert metrics["human_identity"]["raw_score"] == 0.9 + assert [case["item_id"] for case in cases] == ["b"] + with pytest.raises(ValueError, match="not present"): + load_official_results(tmp_path, model_id="absent") + + +def test_ungrouped_model_uses_official_all_group(tmp_path: Path) -> None: + _write_result( + tmp_path, + models={"all": {"num_cases": 1, "scores": {"score": 0.4}}}, + cases=[{"item_id": "video.mp4", "model_id": None, "status": "ok"}], + ) + model, _, cases = load_official_results(tmp_path) + assert model == "all" + assert len(cases) == 1 + + +def test_all_metrics_keep_official_units_without_a_composite(tmp_path: Path) -> None: + root = tmp_path / "official" + for metric, (key, _) in METRIC_SPECS.items(): + value = -0.25 if metric == "environment_depth_fidelity" else 3.5 if metric.startswith("camera_") else 0.6 + _write_result(root, metric, models={"model_a": {"num_cases": 2, "scores": {key: value}}}) + scorecard, code = _run(tmp_path, "--official-results-path", str(root)) + assert code == 0 + metrics = scorecard["metrics"]["per_metric"] + assert set(metrics) == set(METRIC_SPECS) + assert scorecard["metrics"]["summary"]["available_metrics"] == 12 + assert metrics["camera_rotation"]["raw_score"] == metrics["camera_rotation"]["normalized_score"] == 3.5 + assert metrics["camera_translation"]["higher_is_better"] is False + assert metrics["environment_depth_fidelity"]["normalized_score"] == -0.25 + assert not any("average" in key for key in scorecard["metrics"]["leaderboard"]) + assert scorecard["normalizer_only"] is True + assert scorecard["integration_evidence"] is False + assert scorecard["official_benchmark_verified"] is False + assert scorecard["eligibility"]["leaderboard_valid"] is False + + +@pytest.mark.parametrize( + "status,expected_code", + [ + ("model_failure", 0), + ("not_applicable", 0), + ("eval_failed", 1), + ("data_invalid", 1), + ], +) +def test_case_statuses_and_coverage_are_preserved(tmp_path: Path, status: str, expected_code: int) -> None: + root = tmp_path / "official" + scored = 2 if status == "model_failure" else 1 + _write_result( + root, + models={ + "model_a": { + "num_cases": 2, + "num_scored": scored, + "status_counts": {"ok": 1, status: 1}, + "scores": {"score": 0.25, "identity_score": 0.2}, + } + }, + cases=[{"item_id": "a", "model_id": "model_a", "status": status, "error": "metric detail"}], + ) + card, code = _run(tmp_path, "--official-results-path", str(root)) + assert code == expected_code + row = card["metrics"]["per_metric"]["human_identity"] + assert row["sample_count"] == scored + assert row["coverage"]["status_counts"][status] == 1 + assert row["official_scores"]["identity_score"] == 0.2 + assert row["official_groups"][0]["trajectory_id"] == "LRLR" + case = json.loads((tmp_path / "out/per_sample_scores.jsonl").read_text()) + assert case["status"] == status + assert case["error"] == "metric detail" + if expected_code: + assert card["run"]["status"] == "failed" + assert card["evaluation"]["available"] is False + assert card["metrics"]["leaderboard"] == {} + + +def test_all_failed_cases_keep_coverage_without_fabricated_scores(tmp_path: Path) -> None: + root = tmp_path / "official" + _write_result( + root, + models={ + "model_a": { + "num_cases": 2, + "num_scored": 0, + "status_counts": {"eval_failed": 2}, + } + }, + ) + card, code = _run(tmp_path, "--official-results-path", str(root)) + assert code == 1 + row = card["metrics"]["per_metric"]["human_identity"] + assert row["coverage"]["num_scored"] == 0 + assert row["raw_score"] is None + assert not row["available"] + assert card["metrics"]["leaderboard"] == {} + + +def test_requested_metrics_must_exist_for_the_selected_model(tmp_path: Path) -> None: + _write_result(tmp_path) + with pytest.raises(ValueError, match="Missing requested"): + load_official_results(tmp_path, metrics=("camera_rotation",)) + with pytest.raises(ValueError, match="Unknown"): + load_official_results(tmp_path, metrics=("invented_metric",)) + _write_result(tmp_path, "object_geometry", models={"model_b": {"num_cases": 1, "scores": {"score": 0.9}}}) + with pytest.raises(ValueError, match="missing from"): + load_official_results(tmp_path, model_id="model_a") + + +def test_missing_inputs_write_standard_failed_artifacts(tmp_path: Path) -> None: + card, code = _run(tmp_path) + assert code == 1 + assert card["run"]["status"] == "failed" + for filename in ("scorecard.json", "raw_metric_table.jsonl", "per_sample_scores.jsonl"): + assert (tmp_path / "out" / filename).is_file() + + +@pytest.mark.parametrize("upstream_code", [0, 1]) +def test_official_execution_delegates_to_upstream_and_respects_exit_status( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + upstream_code: int, +) -> None: + root = tmp_path / "upstream" + module = root / "src/worldmembench/__main__.py" + module.parent.mkdir(parents=True) + module.touch() + manifest = tmp_path / "inputs.jsonl" + manifest.touch() + config = tmp_path / "config.yaml" + config.write_text("metrics: {}\n") + calls: list[list[str]] = [] + + def execute(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[str]: + calls.append(command) + assert kwargs["cwd"] == root + assert kwargs["env"]["PYTHONPATH"].split(":")[0] == str(root / "src") + assert kwargs["timeout"] == 7200 + result_dir = Path(command[command.index("--output") + 1]) + _write_result(result_dir) + return subprocess.CompletedProcess(command, upstream_code) + + monkeypatch.setattr(runner.subprocess, "run", execute) + card, code = _run( + tmp_path, + "--run-official", + "--repo-root", + str(root), + "--input-path", + str(manifest), + "--config", + str(config), + "--metrics", + "human_identity", + "--gpus", + "0,1", + "--expected-revisits", + "2", + ) + assert code == upstream_code + command = calls[0] + assert command[:5] == [sys.executable, "-m", "worldmembench", "eval", str(manifest)] + assert command[command.index("--metrics") + 1] == "human_identity" + assert "--gpus" in command and "--device" not in command + assert command[command.index("--config") + 1] == str(config) + assert card["validation"]["official_runtime_executed"] is True + assert card["integration_evidence"] is (upstream_code == 0) + assert card["evaluation"]["available"] is (upstream_code == 0) + assert card["validation"]["full_suite_complete"] is False + assert (tmp_path / "out/upstream_stdout.log").is_file() + + +def test_official_execution_does_not_reuse_stale_outputs(tmp_path: Path) -> None: + _write_result(tmp_path / "out/official") + card, code = _run(tmp_path, "--run-official") + assert code == 1 + assert "fresh output" in card["run"]["error"] + + +def test_official_timeout_writes_failure_scorecard(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + video = tmp_path / "video.mp4" + video.touch() + + def execute(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[str]: + raise subprocess.TimeoutExpired(command, kwargs["timeout"]) + + monkeypatch.setattr(runner.subprocess, "run", execute) + card, code = _run(tmp_path, "--run-official", "--input-path", str(video), "--timeout", "1") + assert code == 1 + assert card["run"]["status"] == "failed" + assert "timed out" in card["run"]["error"] + + +def test_standard_benchmark_contract_and_result_import(tmp_path: Path) -> None: + contract = run_benchmark_execution("worldmembench", output_dir=tmp_path / "contract", mode="contract") + assert contract.ok + assert contract.scorecard_path.is_file() + root = tmp_path / "official" + _write_result(root) + imported = run_benchmark_execution( + "worldmembench", + output_dir=tmp_path / "import", + mode="official-validation", + official_results_path=root, + ) + assert imported.ok, imported.to_dict() + card = json.loads(imported.scorecard_path.read_text()) + assert card["benchmark"]["benchmark_id"] == "worldmembench" + assert card["normalizer_only"] is True + assert card["metrics"]["per_metric"]["human_identity"]["raw_score"] == 0.75 + + +def test_standard_official_run_dispatches_generated_videos(tmp_path: Path) -> None: + upstream = tmp_path / "upstream/src/worldmembench" + upstream.mkdir(parents=True) + (upstream / "__init__.py").touch() + payload = {"num_cases": 1, "results": {"all": {"num_cases": 1, "scores": {"score": 0.5}}}} + (upstream / "__main__.py").write_text( + "import json, pathlib, sys\n" + "assert sys.argv[1] == 'eval'\n" + "output = pathlib.Path(sys.argv[sys.argv.index('--output') + 1])\n" + "directory = output / 'metrics' / 'imaging_quality'\n" + "directory.mkdir(parents=True)\n" + f"(directory / 'summary.json').write_text({json.dumps(payload)!r})\n", + encoding="utf-8", + ) + videos = tmp_path / "videos" + videos.mkdir() + (videos / "case_001.mp4").touch() + result = run_benchmark_execution( + "worldmembench", + output_dir=tmp_path / "official-run", + mode="official-run", + generated_artifact_dir=videos, + env_overrides={"WORLDFOUNDRY_WORLDMEMBENCH_ROOT": str(upstream.parent.parent)}, + ) + assert result.ok, result.to_dict() + card = json.loads(result.scorecard_path.read_text()) + assert card["validation"]["official_runtime_executed"] is True + assert card["integration_evidence"] is True + assert card["validation"]["full_suite_complete"] is False + assert card["metrics"]["leaderboard"] == {"imaging_quality": 0.5} + + +@pytest.mark.parametrize("select_model", [False, True]) +def test_public_cli_imports_results(tmp_path: Path, select_model: bool) -> None: + root = tmp_path / "official" + _write_result( + root, + models={ + "model_a": {"num_cases": 2, "scores": {"score": 0.75}}, + **({"model_b": {"num_cases": 2, "scores": {"score": 0.2}}} if select_model else {}), + }, + ) + completed = subprocess.run( + [ + sys.executable, + "-m", + "worldfoundry.cli", + "zoo", + "benchmark-run", + "--benchmark-id", + "worldmembench", + "--mode", + "official-validation", + "--official-results-path", + str(root), + "--output-dir", + str(tmp_path / "cli"), + *(["--result-model-id", "model_a"] if select_model else []), + "--json", + ], + cwd=REPO_ROOT, + capture_output=True, + text=True, + check=False, + ) + assert completed.returncode == 0, completed.stderr + completed.stdout + card = json.loads((tmp_path / "cli/scorecard.json").read_text()) + assert card["metrics"]["leaderboard"] == {"human_identity": 0.75} + assert card["eligibility"]["leaderboard_valid"] is False diff --git a/worldfoundry/data/benchmarks/catalog/video/worldmembench.yaml b/worldfoundry/data/benchmarks/catalog/video/worldmembench.yaml new file mode 100644 index 000000000..a086a73d7 --- /dev/null +++ b/worldfoundry/data/benchmarks/catalog/video/worldmembench.yaml @@ -0,0 +1,196 @@ +id: worldmembench +name: WorldMemBench +publisher: Zhejiang University +publisher_kind: university +benchmark_kind: +- video-world-model +- visual-memory +- camera-control +domains: +- world-model +- human-memory +- object-memory +- environment-memory +- dynamic-memory +modalities: +- video +- text +status: official_runtime_ready_real_data_validation_pending +maturity: contract_ready +release_status: experimental +open_source_status: experimental +official_benchmark_verified: false +integration_evidence: false +leaderboard_valid: false +checkpoint_refs: [] +requires: +- Official result summaries, or generated videos with a provisioned upstream runtime and selected metric assets. +blockers: +- Complete-dataset and checkpoint replay remain pending. +source: + status: open_source +official_sources: + github: + url: https://github.com/ZJU-REAL/WorldMemBench + revision: 1b79840143a98e017891e893b5d3ef75b93078a3 + paper: + url: https://arxiv.org/abs/2610.19063 + huggingface_dataset: + repo_id: rescuerz/WorldMemBench +source_provenance: + official_repo: + url: https://github.com/ZJU-REAL/WorldMemBench + revision: 1b79840143a98e017891e893b5d3ef75b93078a3 + license: MIT + role: external_runtime_dependency + provenance: Official metric implementations are invoked without vendoring or reimplementation. +metrics: +- id: imaging_quality + name: Imaging Quality + description: Official MUSIQ-SPAQ frame scores, normalized upstream and averaged equally per video. + higher_is_better: true + output_unit: score + aggregator: mean +- id: aesthetic_quality + name: Aesthetic Quality + description: Official CLIP/LAION aesthetic frame scores, normalized upstream and averaged per video. + higher_is_better: true + output_unit: score + aggregator: mean +- id: human_identity + name: Human Identity + description: Human identity retention across departure and revisit segments. + higher_is_better: true + output_unit: score + aggregator: mean +- id: human_appearance + name: Human Appearance + description: Human appearance retention under the official revisit protocol. + higher_is_better: true + output_unit: score + aggregator: mean +- id: object_geometry + name: Object Geometry + description: Object geometry retention under the official revisit protocol. + higher_is_better: true + output_unit: score + aggregator: mean +- id: object_appearance + name: Object Appearance + description: Object appearance retention under the official revisit protocol. + higher_is_better: true + output_unit: score + aggregator: mean +- id: environment_visible_memory + name: Environment Visible Memory + description: DINOv2 full-frame similarity at camera return frames. + higher_is_better: true + output_unit: score + aggregator: mean +- id: environment_depth_fidelity + name: Environment Depth Fidelity + description: Pixel-aligned DA3 depth-order Spearman correlation at return frames. + higher_is_better: true + output_unit: correlation + aggregator: mean +- id: environment_scene_reconstruction + name: Environment Scene Reconstruction + description: Official scene reconstruction score at departure and revisit frames. + higher_is_better: true + output_unit: score + aggregator: mean +- id: dynamic_memory + name: Dynamic Memory + description: Event-gated eight-question VLM checklist score with per-category details retained. + higher_is_better: true + output_unit: score + aggregator: mean +- id: camera_rotation + name: Camera Rotation + description: DA3 adjacent rotation RPE normalized by the reference rotation path. + higher_is_better: false + output_unit: error + aggregator: mean +- id: camera_translation + name: Camera Translation + description: DA3/ViPE path-normalized translation RPE following Sim(3) alignment. + higher_is_better: false + output_unit: error + aggregator: mean +integration: + status: integrated + tier: in_tree + evidence: + runtime_surface: external_official_runtime_with_result_normalizer + full_suite_verified: false +runner: + runner_target: worldfoundry.evaluation.tasks.contracts.external:WorldMemBenchContract + official_script: worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py + workspace: + module: worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner + supports_official_runtime: true + accepts_generated_artifacts: true + supports_fixture: false + model_arg: --model-id + metric_arg: --metric + normalization: + artifact_official_run: true + pass_generated_artifact_dir: true + runtime: + kind: external_official_repo + repo_url: https://github.com/ZJU-REAL/WorldMemBench + repo_revision: 1b79840143a98e017891e893b5d3ef75b93078a3 + root_env: WORLDFOUNDRY_WORLDMEMBENCH_ROOT + root_env_required: false + results_path_env: WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH + expected_artifacts: + - scorecard.json + - raw_metric_table.jsonl + - per_sample_scores.jsonl + verification_status: pending + required_env: [] + optional_env: + - WORLDFOUNDRY_WORLDMEMBENCH_ROOT + - WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH + - WORLDFOUNDRY_WORLDMEMBENCH_PYTHON + - WORLDFOUNDRY_WORLDMEMBENCH_CONFIG + - WORLDFOUNDRY_WORLDMEMBENCH_WEIGHTS_DIR + - WORLDFOUNDRY_GENERATED_ARTIFACT_DIR + notes: + - Import the evaluation output root, its metrics directory, or a metric summary.json. + - Multi-model outputs require an explicit model selection; scores are never pooled across models. + - Evaluation and invalid-data failures retain coverage and produce a nonzero runner exit. + - Model failures remain scored according to the official protocol; not-applicable cases are not zero-filled. + - Raw metric units and score directions are retained without a cross-metric average. +data_refs: + task_yaml: worldfoundry/data/benchmarks/tasks/external/worldmembench.yaml + hf_dataset_ids: + - rescuerz/WorldMemBench +artifact_layout: + output_artifact: benchmark_scorecard + filename_template: scorecard.json +runner_availability: + available: true + runner_target: worldfoundry.evaluation.tasks.contracts.external:WorldMemBenchContract + runner_script_available: true + task_yaml_available: true + runtime_verified: false + full_suite_verified: false + surface: external_official_runtime_with_result_normalizer +validation_command: +- worldfoundry-eval +- zoo +- benchmark-run +- --benchmark-id +- worldmembench +- --mode +- official-validation +- --official-results-path +- +- --output-dir +- +- --json +notes: +- Source and metric protocol are attributed to ZJU-REAL/WorldMemBench under MIT. +- Result import is weight-free; official metric execution requires a separately provisioned upstream runtime and assets. +- GPU and full-checkpoint replay are pending; result import does not establish leaderboard validity. diff --git a/worldfoundry/data/benchmarks/runtime_profiles/official/worldmembench.yaml b/worldfoundry/data/benchmarks/runtime_profiles/official/worldmembench.yaml new file mode 100644 index 000000000..ce275af0f --- /dev/null +++ b/worldfoundry/data/benchmarks/runtime_profiles/official/worldmembench.yaml @@ -0,0 +1,67 @@ +id: worldmembench +name: WorldMemBench +benchmark_ids: +- worldmembench +environment_id: worldfoundry-unified-cu128 +needs_new_env: false +python_path: ${WORLDFOUNDRY_UNIFIED_PYTHON:-python} +pythonpath_roots: +- . +required_env: [] +optional_env: +- WORLDFOUNDRY_WORLDMEMBENCH_ROOT +- WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH +- WORLDFOUNDRY_WORLDMEMBENCH_PYTHON +- WORLDFOUNDRY_WORLDMEMBENCH_CONFIG +- WORLDFOUNDRY_WORLDMEMBENCH_WEIGHTS_DIR +- WORLDFOUNDRY_GENERATED_ARTIFACT_DIR +required_imports: [] +required_packages: [] +required_paths: +- id: official_runner_script + path: worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py +- id: metrics_module + path: worldfoundry/evaluation/tasks/execution/runners/worldmembench/worldmembench_metrics.py +- id: task_yaml + path: worldfoundry/data/benchmarks/tasks/external/worldmembench.yaml +required_assets: + task_yaml: worldfoundry/data/benchmarks/tasks/external/worldmembench.yaml + hf_dataset_ids: + - rescuerz/WorldMemBench +requires_cuda_visibility: false +status: official_runtime_ready_real_data_validation_pending +notes: +- Result import uses the unified environment without CUDA, metric checkpoints, or upstream Python dependencies. +- Official execution uses an installed worldmembench package or WORLDFOUNDRY_WORLDMEMBENCH_ROOT/src. +- WORLDFOUNDRY_WORLDMEMBENCH_PYTHON selects a provisioned upstream interpreter without changing the unified Torch/CUDA ABI. +- DA3 and ViPE subprocess interpreters and metric checkpoint paths belong in the upstream YAML configuration. +- Dynamic Memory uses the upstream VLM configuration and API-key environment variable; credentials are not runner arguments. +- Official dataset/checkpoint replay and leaderboard parity remain pending. +validation_command: &worldmembench_validation_command +- worldfoundry-eval +- zoo +- benchmark-run +- --benchmark-id +- worldmembench +- --mode +- official-validation +- --official-results-path +- +- --output-dir +- +- --json +normalization_command: *worldmembench_validation_command +full_runtime_command: +- python +- -m +- worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner +- --run-official +- --repo-root +- +- --input-path +- +- --config +- +- --output-dir +- +- --json diff --git a/worldfoundry/data/benchmarks/tasks/external/worldmembench.yaml b/worldfoundry/data/benchmarks/tasks/external/worldmembench.yaml new file mode 100644 index 000000000..e7b714a88 --- /dev/null +++ b/worldfoundry/data/benchmarks/tasks/external/worldmembench.yaml @@ -0,0 +1,172 @@ +benchmark: worldmembench +name: WorldMemBench +description: Official WorldMemBench memory and camera-control evaluation for video world models. +tags: +- benchmark-zoo +- external-official +- worldmembench +- video-world-model +- visual-memory +splits: +- standard +metrics: + imaging_quality: + aggregation: mean + aesthetic_quality: + aggregation: mean + human_identity: + aggregation: mean + human_appearance: + aggregation: mean + object_geometry: + aggregation: mean + object_appearance: + aggregation: mean + environment_visible_memory: + aggregation: mean + environment_depth_fidelity: + aggregation: mean + environment_scene_reconstruction: + aggregation: mean + dynamic_memory: + aggregation: mean + camera_rotation: + aggregation: mean + camera_translation: + aggregation: mean +metadata: + source_kind: benchmark_zoo + benchmark_id: worldmembench + official_runner_target: worldfoundry.evaluation.tasks.contracts.external:WorldMemBenchContract + official_runner_script: worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py + official_runner_command: worldfoundry-eval zoo benchmark-run + contract_validation_command: worldfoundry-eval zoo benchmark-run --benchmark-id worldmembench --mode contract --output-dir tmp/benchmark_zoo/benchmark_contract/worldmembench --json + task_type: video_world_model_memory_evaluation + status: integrated + maturity: contract_ready + release_status: experimental + open_source_status: experimental + checkpoint_refs: [] + requires: + - Official result summaries, or generated videos with a provisioned upstream runtime and selected metric assets. + blockers: + - Complete-dataset and checkpoint replay remain pending. + validation_command: + - worldfoundry-eval + - zoo + - benchmark-run + - --benchmark-id + - worldmembench + - --mode + - official-validation + - --official-results-path + - + - --output-dir + - + - --json + official_benchmark_verified: false + integration_evidence: false + leaderboard_valid: false + source_provenance: + official_repo: + url: https://github.com/ZJU-REAL/WorldMemBench + revision: 1b79840143a98e017891e893b5d3ef75b93078a3 + license: MIT + role: external_runtime_dependency + paper: + url: https://arxiv.org/abs/2610.19063 + huggingface_dataset: + repo_id: rescuerz/WorldMemBench + input_contract: + - generated_video_dir + - official_results_path + - worldmembench_config + - weights_dir + output_contract: + - scorecard + - raw_metric_table + - per_sample_scores + - official_metric_summaries + evaluation_requirements: + - Use official JSONL manifests or the upstream model/category/trajectory/case video layout. + - Import official metrics//summary.json outputs or execute a provisioned upstream runtime. + - Keep one model group per scorecard and retain official per-case statuses and camera-error directions. + - Official metric computation requires the selected checkpoint stacks and DA3/ViPE environments; Dynamic Memory requires a VLM endpoint. + data_refs: + task_yaml: worldfoundry/data/benchmarks/tasks/external/worldmembench.yaml + hf_dataset_ids: + - rescuerz/WorldMemBench + runner_availability: + available: true + runner_target: worldfoundry.evaluation.tasks.contracts.external:WorldMemBenchContract + runner_script_available: true + task_yaml_available: true + runtime_verified: false + full_suite_verified: false + surface: external_official_runtime_with_result_normalizer + runtime: + kind: external_official_repo + root_env: WORLDFOUNDRY_WORLDMEMBENCH_ROOT + results_path_env: WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH + artifact_layout: + generated_artifact_dir_env: WORLDFOUNDRY_GENERATED_ARTIFACT_DIR + output_artifact: benchmark_scorecard + filename_template: scorecard.json +tasks: + worldmembench_standard: + benchmark_name: worldmembench + protocol: external_official_benchmark + capability_track: video_world_model_memory + schema_type: generated_artifact_set + input_keys: + - generated_video_dir + - official_results_path + - worldmembench_config + - weights_dir + output_keys: + - scorecard + - raw_metric_table + - per_sample_scores + - official_metric_summaries + metric_ids: &worldmembench_metric_ids + - imaging_quality + - aesthetic_quality + - human_identity + - human_appearance + - object_geometry + - object_appearance + - environment_visible_memory + - environment_depth_fidelity + - environment_scene_reconstruction + - dynamic_memory + - camera_rotation + - camera_translation + tags: + - benchmark-zoo + - external-official + - worldmembench + - video-world-model + metadata: + source_kind: benchmark_zoo + benchmark_id: worldmembench + official_benchmark_verified: false + integration_evidence: false + leaderboard_valid: false + dataset: + hf_dataset_id: rescuerz/WorldMemBench + split: standard + runtime_requirements: + generated_artifact_dir_env: WORLDFOUNDRY_GENERATED_ARTIFACT_DIR + results_path_env: WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH + artifact_contract: + required: + generated_video: + kind: video + optional: + official_results_path: + kind: directory + evaluation_protocol: + - name: external_official_runner + metric_ids: *worldmembench_metric_ids + runner_target: worldfoundry.evaluation.tasks.contracts.external:WorldMemBenchContract + runner_script: worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py diff --git a/worldfoundry/evaluation/tasks/contracts/external.py b/worldfoundry/evaluation/tasks/contracts/external.py index 148626fef..ad5602907 100644 --- a/worldfoundry/evaluation/tasks/contracts/external.py +++ b/worldfoundry/evaluation/tasks/contracts/external.py @@ -22,7 +22,6 @@ register_external_benchmark_contract, ) - ROBOTWIN_BENCHMARK_ID = "robotwin" ROBOTWIN_BENCHMARK_NAME = "RoboTwin 2.0" ROBOTWIN_ROOT_ENV = "WORLDFOUNDRY_ROBOTWIN_ROOT" @@ -856,6 +855,33 @@ def build_robotwin_contract_payload(benchmark_id: str = ROBOTWIN_BENCHMARK_ID) - ), ) +WorldMemBenchContract = ExternalBenchmarkContract( + benchmark_id="worldmembench", + display_name="WorldMemBench", + input_keys=("generated_video_dir", "official_results_path", "worldmembench_config", "weights_dir"), + output_keys=("scorecard", "raw_metric_table", "per_sample_scores", "official_metric_summaries"), + metric_ids=( + "imaging_quality", + "aesthetic_quality", + "human_identity", + "human_appearance", + "object_geometry", + "object_appearance", + "environment_visible_memory", + "environment_depth_fidelity", + "environment_scene_reconstruction", + "dynamic_memory", + "camera_rotation", + "camera_translation", + ), + requires_upstream_runtime=True, + notes=( + "Imports official per-model summaries or invokes the installed ZJU-REAL/WorldMemBench runtime.", + "Scores retain official units, case statuses, and camera-error directions; no composite score is defined.", + "Full evaluation needs metric checkpoints, DA3/ViPE environments, and a Dynamic Memory VLM endpoint.", + ), +) + MemoBenchContract = ExternalBenchmarkContract( benchmark_id="memobench", display_name="MemoBench", @@ -1061,7 +1087,12 @@ def build_robotwin_contract_payload(benchmark_id: str = ROBOTWIN_BENCHMARK_ID) - IPVBenchContract = ExternalBenchmarkContract( benchmark_id="ipv-bench", display_name="IPV-Bench", - input_keys=("impossible_prompt_manifest", "generated_video_dir", "understanding_results_path", "official_results_path"), + input_keys=( + "impossible_prompt_manifest", + "generated_video_dir", + "understanding_results_path", + "official_results_path", + ), output_keys=("scorecard", "raw_metric_table", "per_sample_scores", "benchmark_contract"), metric_ids=( "visual_quality", @@ -1155,7 +1186,12 @@ def build_robotwin_contract_payload(benchmark_id: str = ROBOTWIN_BENCHMARK_ID) - PhyEduVideoContract = ExternalBenchmarkContract( benchmark_id="phyeduvideo", display_name="PhyEduVideo", - input_keys=("physics_education_prompt_manifest", "generated_video_dir", "judge_results_path", "official_results_path"), + input_keys=( + "physics_education_prompt_manifest", + "generated_video_dir", + "judge_results_path", + "official_results_path", + ), output_keys=("scorecard", "raw_metric_table", "per_sample_scores", "benchmark_contract"), metric_ids=( "semantic_adherence", @@ -1213,7 +1249,12 @@ def build_robotwin_contract_payload(benchmark_id: str = ROBOTWIN_BENCHMARK_ID) - WorldInWorldContract = ExternalBenchmarkContract( benchmark_id="world-in-world", display_name="World-in-World", - input_keys=("closed_loop_task_manifest", "world_model_rollouts_dir", "planner_results_path", "official_results_path"), + input_keys=( + "closed_loop_task_manifest", + "world_model_rollouts_dir", + "planner_results_path", + "official_results_path", + ), output_keys=("scorecard", "raw_metric_table", "rollout_logs", "benchmark_contract"), metric_ids=( "active_recognition_success_rate", @@ -1604,15 +1645,30 @@ def _embodied_action_contract( Behavior1KContract = _embodied_action_contract( "behavior1k", "BEHAVIOR-1K", - input_keys=("behavior_dataset_root", "simulator_root", "task_manifest", "policy_results_path", "official_results_path"), + input_keys=( + "behavior_dataset_root", + "simulator_root", + "task_manifest", + "policy_results_path", + "official_results_path", + ), metric_ids=("success_rate", "task_success", "episode_success", "subgoal_success", "reward"), notes=("BEHAVIOR-1K requires OmniGibson/Isaac Sim assets and license-gated simulator setup.",), ) KinetixContract = _embodied_action_contract( "kinetix", "Kinetix", - metric_ids=("success_rate", "task_success", "episode_success", "reward", "normalized_return", "real_time_success_rate"), - notes=("Kinetix uses a JAX 2D physics runtime; keep it in a simulator-specific environment when the unified PyTorch environment is insufficient.",), + metric_ids=( + "success_rate", + "task_success", + "episode_success", + "reward", + "normalized_return", + "real_time_success_rate", + ), + notes=( + "Kinetix uses a JAX 2D physics runtime; keep it in a simulator-specific environment when the unified PyTorch environment is insufficient.", + ), ) LIBEROMemContract = _embodied_action_contract( "libero-mem", @@ -1651,7 +1707,13 @@ def _embodied_action_contract( MolmoSpacesContract = _embodied_action_contract( "molmospaces", "MolmoSpaces-Bench", - input_keys=("molmospaces_asset_root", "simulator_root", "task_manifest", "policy_results_path", "official_results_path"), + input_keys=( + "molmospaces_asset_root", + "simulator_root", + "task_manifest", + "policy_results_path", + "official_results_path", + ), metric_ids=("success_rate", "pick_success_rate", "task_success", "episode_success", "reward"), notes=("MolmoSpaces assets are large; stage only the benchmark subset required by the selected suite.",), ) @@ -1664,7 +1726,13 @@ def _embodied_action_contract( "robomme", "RoboMME", input_keys=("dataset_root", "simulator_root", "memory_manifest", "policy_results_path", "official_results_path"), - metric_ids=("success_rate", "counting_success_rate", "permanence_success_rate", "reference_success_rate", "imitation_success_rate"), + metric_ids=( + "success_rate", + "counting_success_rate", + "permanence_success_rate", + "reference_success_rate", + "imitation_success_rate", + ), ) VLABenchContract = _embodied_action_contract( "vlabench", @@ -1792,6 +1860,7 @@ def _embodied_action_contract( WorldArenaContract, WorldBenchContract, WorldInWorldContract, + WorldMemBenchContract, WorldModelBenchContract, HarnessEvalWContract, WorldOlympiadContract, @@ -1834,8 +1903,4 @@ def list_external_benchmark_contracts() -> tuple[ExternalBenchmarkContract, ...] return _CONTRACT_REGISTRY.list() -__all__ = sorted( - name - for name in globals() - if not name.startswith("_") and name not in {"Path", "annotations"} -) +__all__ = sorted(name for name in globals() if not name.startswith("_") and name not in {"Path", "annotations"}) diff --git a/worldfoundry/evaluation/tasks/execution/framework/official_runner.py b/worldfoundry/evaluation/tasks/execution/framework/official_runner.py index 0cf2c483c..249f37b13 100644 --- a/worldfoundry/evaluation/tasks/execution/framework/official_runner.py +++ b/worldfoundry/evaluation/tasks/execution/framework/official_runner.py @@ -29,8 +29,8 @@ from types import MappingProxyType from typing import Any, Callable, Mapping -from worldfoundry.core.observability.logging_setup import configure_logging, get_logger from worldfoundry.core.execution.process import run_logged_subprocess +from worldfoundry.core.observability.logging_setup import configure_logging, get_logger from worldfoundry.evaluation.reporting.scorecard import SCORECARD_SCHEMA_VERSION from worldfoundry.evaluation.tasks.catalog.zoo_registry import load_benchmark_zoo_registry from worldfoundry.evaluation.tasks.execution.framework.io import ( @@ -49,7 +49,6 @@ from worldfoundry.evaluation.tasks.execution.framework.result_normalizer import OfficialResultsNormalizer from worldfoundry.evaluation.utils import BENCHMARK_ZOO_DIR, REPO_ROOT - ExtractMetricsFn = Callable[[Any, Path], dict[str, dict[str, Any]]] DiscoverResultsFn = Callable[[Path, Path | None], Path | None] BuildCommandFn = Callable[..., list[str] | None] @@ -69,7 +68,7 @@ class BenchRunnerConfig: metric_order: tuple[str, ...] metric_specs: dict[str, dict[str, Any]] metric_aliases: dict[str, str] - average_metric_id: str + average_metric_id: str | None official_entry: str | None = None official_output_globs: tuple[str, ...] = () generated_video_envs: tuple[str, ...] = ("WORLDFOUNDRY_GENERATED_ARTIFACT_DIR",) @@ -248,9 +247,7 @@ def metric_id_from_key(key: Any, config: BenchRunnerConfig) -> str | None: # only trusted when every fuzzy hit agrees on one metric id; ambiguous keys # are dropped instead of being attached to whichever alias iterates first. fuzzy_hits = { - metric_id - for alias, metric_id in config.metric_aliases.items() - if alias in normalized or normalized in alias + metric_id for alias, metric_id in config.metric_aliases.items() if alias in normalized or normalized in alias } if len(fuzzy_hits) != 1: if len(fuzzy_hits) > 1 and (config.benchmark_id, normalized) not in _FUZZY_ALIAS_WARNED: @@ -291,9 +288,7 @@ def declared_metric_normalizers(benchmark_id: str) -> Mapping[str, str]: try: entry = load_benchmark_zoo_registry(BENCHMARK_ZOO_DIR).get(benchmark_id) declared = { - metric.metric_id: str(metric.normalizer) - for metric in entry.metrics - if getattr(metric, "normalizer", None) + metric.metric_id: str(metric.normalizer) for metric in entry.metrics if getattr(metric, "normalizer", None) } except Exception: # noqa: BLE001 - uncataloged benchmarks fall back to the legacy heuristic. declared = {} @@ -417,6 +412,8 @@ def catalog_fallback( def compute_average(config: BenchRunnerConfig, extracted: dict[str, dict[str, Any]]) -> None: + if config.average_metric_id is None: + return existing = extracted.get(config.average_metric_id) if existing is not None and existing.get("normalized_score") is not None: return @@ -447,9 +444,7 @@ def metric_row( config: BenchRunnerConfig | None = None, ) -> dict[str, Any]: normalized = ( - normalized_metric_score(config, metric_id, raw_score) - if config is not None - else normalize_unit_score(raw_score) + normalized_metric_score(config, metric_id, raw_score) if config is not None else normalize_unit_score(raw_score) ) return { "metric_id": metric_id, @@ -479,7 +474,7 @@ def merge_metric_id_score_rows( for metric_id, values in buckets.items(): extracted[metric_id] = metric_row( metric_id, - mean_numeric(values), + sum(values) / len(values), source=str(results_path), sample_count=len(values), config=config, @@ -609,9 +604,7 @@ def metric_group(metric_id: str) -> str: for metric_id in metric_order } sample_display_name = ( - f"sample_results{Path(sample_filename).suffix}" - if sample_filename is not None - else "sample_results.*" + f"sample_results{Path(sample_filename).suffix}" if sample_filename is not None else "sample_results.*" ) return BenchRunnerConfig( benchmark_id=benchmark_id, @@ -831,9 +824,7 @@ def run_official_pipeline( if results_path is not None and hooks.prepare_upstream_results is not None: results_path = hooks.prepare_upstream_results(config, results_path, args, output_dir) - extract_fn = hooks.extract_metrics or ( - lambda payload, path: generic_extract_metrics(payload, config, str(path)) - ) + extract_fn = hooks.extract_metrics or (lambda payload, path: generic_extract_metrics(payload, config, str(path))) if results_path is not None: payload, _fmt = load_upstream_payload(results_path) @@ -1126,7 +1117,9 @@ def run_main(config: BenchRunnerConfig, hooks: RunnerHooks, argv: list[str] | No sample_path = benchmark_task_sample_path(config.benchmark_id) if sample_path is None: logger.event("ERROR", "official_runner.fixture_missing", "Official runner fixture is unavailable") - reason = f"no checked-in sample for {config.benchmark_id} under {BENCHMARK_ASSETS_ROOT / config.benchmark_id}" + reason = ( + f"no checked-in sample for {config.benchmark_id} under {BENCHMARK_ASSETS_ROOT / config.benchmark_id}" + ) write_failed_scorecard( benchmark_id=config.benchmark_id, display_name=config.display_name, @@ -1144,9 +1137,7 @@ def run_main(config: BenchRunnerConfig, hooks: RunnerHooks, argv: list[str] | No env_results = resolve_results_path(config, None) if env_results is None: logger.event("ERROR", "official_runner.results_missing", "Official runner results path is unavailable") - reason = ( - f"--official-results-path or {config.results_path_env} is required unless --run-official is set" - ) + reason = f"--official-results-path or {config.results_path_env} is required unless --run-official is set" write_failed_scorecard( benchmark_id=config.benchmark_id, display_name=config.display_name, diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py new file mode 100644 index 000000000..18475834d --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 WorldFoundry contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Run ZJU-REAL/WorldMemBench or import its official per-model results.""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +from worldfoundry.evaluation.tasks.execution.framework.io import write_json, write_jsonl +from worldfoundry.evaluation.tasks.execution.framework.official_runner import ( + BenchRunnerConfig, + build_scorecard, + write_failed_scorecard, +) +from worldfoundry.evaluation.tasks.execution.framework.runner_common import resolve_env_path as _env_path +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.worldmembench_metrics import ( + METRIC_SPECS, + load_official_results, +) + +CONFIG = BenchRunnerConfig( + benchmark_id="worldmembench", + display_name="WorldMemBench", + root_env="WORLDFOUNDRY_WORLDMEMBENCH_ROOT", + results_path_env="WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH", + default_repo_subdir="", + metric_order=tuple(METRIC_SPECS), + metric_specs={ + name: {"name": name.replace("_", " ").title(), "group": name.split("_")[0], "higher_is_better": higher} + for name, (_, higher) in METRIC_SPECS.items() + }, + metric_aliases={}, + average_metric_id=None, +) + + +def build_parser() -> argparse.ArgumentParser: + """Expose the standard benchmark runner flags and upstream runtime inputs.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--benchmark-id", choices=["worldmembench"], default="worldmembench") + parser.add_argument( + "--official-results-path", "--results-path", type=Path, default=_env_path(CONFIG.results_path_env) + ) + parser.add_argument( + "--generated-video-dir", + "--generated-artifact-dir", + "--input-path", + type=Path, + default=_env_path("WORLDFOUNDRY_GENERATED_ARTIFACT_DIR"), + help="Official video tree, video file, or JSONL input manifest", + ) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--run-official", action="store_true") + parser.add_argument( + "--repo-root", + type=Path, + default=_env_path(CONFIG.root_env), + help="Optional official checkout; otherwise use the installed worldmembench package", + ) + parser.add_argument("--python", default=os.environ.get("WORLDFOUNDRY_WORLDMEMBENCH_PYTHON", sys.executable)) + parser.add_argument("--config", type=Path, default=_env_path("WORLDFOUNDRY_WORLDMEMBENCH_CONFIG")) + parser.add_argument("--weights-dir", type=Path, default=_env_path("WORLDFOUNDRY_WORLDMEMBENCH_WEIGHTS_DIR")) + parser.add_argument("--device", default="cuda") + parser.add_argument("--gpus", help="Official multi-GPU process sharding, e.g. 0,1") + parser.add_argument("--model-id", help="Select one upstream model group; required for multi-model imports") + parser.add_argument( + "--metric", + "--metrics", + action="append", + default=[], + help="Official metric IDs; repeat the flag or separate IDs with commas", + ) + parser.add_argument("--expected-revisits", type=int, choices=[1, 2, 3]) + parser.add_argument( + "--timeout", type=float, default=float(os.environ.get("WORLDFOUNDRY_BENCHMARK_TIMEOUT", "7200")) + ) + parser.add_argument("--json", action="store_true") + return parser + + +def build_official_command(args: argparse.Namespace, results: Path, metrics: tuple[str, ...]) -> list[str]: + """Delegate metric computation and native multi-GPU execution to upstream.""" + if args.generated_video_dir is None or not args.generated_video_dir.exists(): + raise ValueError("--run-official requires an existing --generated-video-dir or JSONL --input-path") + command = [ + args.python, + "-m", + "worldmembench", + "eval", + str(args.generated_video_dir.resolve()), + "--output", + str(results), + ] + if metrics: + command.extend(["--metrics", *metrics]) + if args.model_id and args.model_id != "all": + command.extend(["--models", args.model_id]) + if args.gpus: + command.extend(["--gpus", args.gpus]) + else: + command.extend(["--device", args.device]) + for flag, value in [("--config", args.config), ("--weights-dir", args.weights_dir)]: + if value is not None: + command.extend([flag, str(value.expanduser().resolve())]) + if args.expected_revisits is not None: + command.extend(["--expected-revisits", str(args.expected_revisits)]) + return command + + +def run(args: argparse.Namespace) -> tuple[dict[str, Any], int]: + """Write standard scorecard artifacts, retaining upstream coverage and failures.""" + output = args.output_dir.expanduser().resolve() + output.mkdir(parents=True, exist_ok=True) + command: list[str] | None = None + started = time.monotonic() + try: + metrics = tuple( + dict.fromkeys(name.strip() for value in args.metric for name in value.split(",") if name.strip()) + ) + if set(metrics) - METRIC_SPECS.keys(): + raise ValueError("Unknown WorldMemBench metric: " + ", ".join(sorted(set(metrics) - METRIC_SPECS.keys()))) + results = args.official_results_path + returncode = 0 + if args.run_official: + if results is not None: + raise ValueError("Use --run-official or --official-results-path, not both") + results = output / "official" + if results.exists(): + raise ValueError(f"Use a fresh output directory; official results already exist: {results}") + command = build_official_command(args, results, metrics) + environment = dict(os.environ) + if args.repo_root is not None: + args.repo_root = args.repo_root.expanduser().resolve() + source = args.repo_root / "src" + if not (source / "worldmembench" / "__main__.py").is_file(): + raise ValueError(f"Not a WorldMemBench checkout: {args.repo_root}") + environment["PYTHONPATH"] = os.pathsep.join( + filter(None, [str(source), environment.get("PYTHONPATH", "")]) + ) + with ( + (output / "upstream_stdout.log").open("w", encoding="utf-8") as stdout, + (output / "upstream_stderr.log").open("w", encoding="utf-8") as stderr, + ): + completed = subprocess.run( + command, + cwd=args.repo_root, + env=environment, + stdout=stdout, + stderr=stderr, + timeout=args.timeout, + check=False, + ) + returncode = completed.returncode + elif results is None: + raise ValueError("Pass --official-results-path or --run-official with generated videos") + results = results.expanduser().resolve() + model_id, extracted, cases = load_official_results(results, model_id=args.model_id, metrics=metrics) + evaluation_failed = any( + row["coverage"]["status_counts"].get(status, 0) + for row in extracted.values() + for status in ("eval_failed", "data_invalid") + ) + if evaluation_failed or not any(row["raw_score"] is not None for row in extracted.values()): + returncode = returncode or 1 + scorecard = build_scorecard( + config=CONFIG, + output_dir=output, + results_path=results, + extracted=extracted, + command=command, + duration_seconds=time.monotonic() - started, + returncode=returncode, + repo_root=args.repo_root, + generated_video_dir=args.generated_video_dir, + run_status_override="failed" if returncode else None, + ) + scorecard["dataset"]["model_id"] = model_id + scorecard["evaluation"]["kind"] = "worldmembench_official" + scorecard["evaluation"]["available"] = returncode == 0 + if returncode: + scorecard["metrics"]["leaderboard"] = {} + scorecard["evaluation"]["leaderboard_metrics"] = {} + scorecard["validation"]["upstream_case_failures"] = evaluation_failed + for name, item in extracted.items(): + scorecard["metrics"]["per_metric"][name].update( + {key: item[key] for key in ("coverage", "official_scores", "official_groups")} + ) + write_jsonl(output / "raw_metric_table.jsonl", list(scorecard["metrics"]["per_metric"].values())) + write_jsonl(output / "per_sample_scores.jsonl", cases) + write_json(output / "scorecard.json", scorecard) + return scorecard, 0 if returncode == 0 else 1 + except (ValueError, OSError, subprocess.SubprocessError) as error: + scorecard = write_failed_scorecard( + benchmark_id=CONFIG.benchmark_id, + display_name=CONFIG.display_name, + output_dir=output, + reason=str(error), + command=command, + ) + return scorecard, 1 + + +def main(argv: list[str] | None = None) -> int: + """Run the benchmark CLI without importing metric models in this process.""" + args = build_parser().parse_args(argv) + scorecard, returncode = run(args) + print(json.dumps(scorecard, ensure_ascii=False) if args.json else str(args.output_dir / "scorecard.json")) + return returncode + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/worldmembench_metrics.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/worldmembench_metrics.py new file mode 100644 index 000000000..7f108622a --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/worldmembench_metrics.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 WorldFoundry contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Import the official ZJU-REAL/WorldMemBench per-model metric summaries. + +Protocol: https://github.com/ZJU-REAL/WorldMemBench +Scores retain upstream units and directions; there is no cross-metric average. +""" + +from __future__ import annotations + +import json +import math +from pathlib import Path +from typing import Any + +# Metric directory -> (official score field, higher-is-better). +METRIC_SPECS: dict[str, tuple[str, bool]] = { + "imaging_quality": ("score", True), + "aesthetic_quality": ("score", True), + "human_identity": ("score", True), + "human_appearance": ("score", True), + "object_geometry": ("score", True), + "object_appearance": ("score", True), + "environment_visible_memory": ("score", True), + "environment_depth_fidelity": ("depth_spearman", True), + "environment_scene_reconstruction": ("reconstruction_score", True), + "dynamic_memory": ("score", True), + "camera_rotation": ("camera_rotation_rpe_path", False), + "camera_translation": ("camera_translation_rpe", False), +} + + +def load_official_results( + path: Path, + *, + model_id: str | None = None, + metrics: tuple[str, ...] = (), +) -> tuple[str, dict[str, dict[str, Any]], list[dict[str, Any]]]: + """Read compact official summaries without pooling generation models. + + Accept an evaluation output root, its ``metrics`` directory, or one + ``metrics//summary.json`` file. Official case records, when present, + are copied with their status and detailed scores intact. + """ + unknown = set(metrics) - METRIC_SPECS.keys() + if unknown: + raise ValueError(f"Unknown WorldMemBench metrics: {', '.join(sorted(unknown))}") + if path.is_file(): + paths = [path] + else: + root = path / "metrics" if (path / "metrics").is_dir() else path + paths = sorted(root.glob("*/summary.json")) + summaries: dict[str, tuple[Path, dict[str, Any]]] = {} + for summary_path in paths: + metric_id = summary_path.parent.name + if metric_id not in METRIC_SPECS or (metrics and metric_id not in metrics): + continue + summaries[metric_id] = (summary_path, json.loads(summary_path.read_text(encoding="utf-8"))) + missing = set(metrics) - summaries.keys() + if missing: + raise ValueError(f"Missing requested metric summaries: {', '.join(sorted(missing))}") + if not summaries: + raise ValueError(f"No selected WorldMemBench metric summaries under {path}") + models = {key for _, summary in summaries.values() for key in summary.get("results", {})} + if model_id is None: + if len(models) != 1: + raise ValueError("Pass --model-id to select one official model group; models: " + ", ".join(sorted(models))) + model_id = next(iter(models)) + if model_id not in models: + raise ValueError(f"Model {model_id!r} is not present in the selected summaries") + + extracted: dict[str, dict[str, Any]] = {} + cases: list[dict[str, Any]] = [] + for metric_id, (summary_path, summary) in summaries.items(): + group = summary.get("results", {}).get(model_id) + if group is None: + raise ValueError(f"Model {model_id!r} is missing from {summary_path}") + count = int(group["num_cases"]) + scored = int(group.get("num_scored", count)) + statuses = group.get("status_counts", {"ok": count}) + score_key, _ = METRIC_SPECS[metric_id] + raw = group.get("scores", {}).get(score_key) + if raw is not None: + if isinstance(raw, bool) or not isinstance(raw, (int, float)) or not math.isfinite(raw): + raise ValueError(f"Non-finite or non-numeric {metric_id} score in {summary_path}") + raw = float(raw) + extracted[metric_id] = { + "metric_id": metric_id, + "raw_score": raw, + "normalized_score": raw, + "sample_count": scored, + "source": str(summary_path), + "model_id": model_id, + "coverage": {"num_cases": count, "num_scored": scored, "status_counts": statuses}, + "official_scores": group.get("scores", {}), + "official_groups": [g for g in summary.get("groups", []) if str(g.get("model_id") or "all") == model_id], + } + case_path = summary_path.with_name("cases.jsonl") + if case_path.is_file(): + for line in case_path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + case = json.loads(line) + if str(case.get("model_id") or "all") == model_id: + cases.append({**case, "metric_id": metric_id}) + return model_id, extracted, cases From a2b940bc371c7a093a7d893122e5499246ce5e35 Mon Sep 17 00:00:00 2001 From: iChubai <149308795+iChubai@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:47:28 +0800 Subject: [PATCH 2/2] [refactor]Run WorldMemBench in tree with shared metrics Execute all twelve official metric protocols in tree. Reuse canonical CLIP, LAION heads and DINO loaders, and share compatible DINO backbones per run. Centralize WorldScore and 4DWorldBench IQA preprocessing and preserve benchmark-specific sampling, alignment and aggregation. Bind WBench imaging quality to its official KonIQ checkpoint profile. Signed-off-by: iChubai <149308795+iChubai@users.noreply.github.com> --- Makefile | 3 + THIRD-PARTY-NOTICES | 35 ++ .../docs/evaluation/benchmark-hub/wbench.mdx | 6 + .../evaluation/benchmark-hub/wbench.zh.mdx | 4 + .../benchmark-hub/worldmembench.mdx | 167 +++--- .../benchmark-hub/worldmembench.zh.mdx | 139 +++-- .../lib/benchmark-catalog-status.json | 6 +- .../eval_core/test_wbench_official_runner.py | 26 +- .../test_worldmembench_official_runner.py | 33 +- tests/eval_core/test_worldmembench_runtime.py | 135 +++++ .../execution/test_shared_quality.py | 269 +++++++++ .../execution/test_worldmembench_camera.py | 73 +++ worldfoundry/base_models/capabilities.py | 200 +++---- .../general_perception/dino_embeddings.py | 19 + .../catalog/video/worldmembench.yaml | 20 +- .../benchmarks/requirements/worldmembench.txt | 13 + .../official/worldmembench.yaml | 16 +- .../runtime_profiles/official/worldscore.yaml | 26 +- .../evaluation/tasks/contracts/external.py | 4 +- .../four_d_worldbench/metric/base_metrics.py | 114 +--- .../runners/mind/runtime/mind/src/process.py | 347 +++++++---- .../metrics/video_quality/imaging_quality.py | 17 +- .../wbench/src/metrics/weight_utils.py | 8 +- .../run_worldmembench_official_runner.py | 30 +- .../runners/worldmembench/runtime/LICENSE | 21 + .../runners/worldmembench/runtime/__init__.py | 4 + .../runners/worldmembench/runtime/__main__.py | 6 + .../runners/worldmembench/runtime/cli.py | 106 ++++ .../runners/worldmembench/runtime/context.py | 60 ++ .../runners/worldmembench/runtime/data.py | 125 ++++ .../runners/worldmembench/runtime/errors.py | 24 + .../runners/worldmembench/runtime/evaluate.py | 121 ++++ .../worldmembench/runtime/execution.py | 156 +++++ .../worldmembench/runtime/metrics/__init__.py | 45 ++ .../runtime/metrics/aggregation.py | 47 ++ .../worldmembench/runtime/metrics/base.py | 36 ++ .../runtime/metrics/camera/__init__.py | 8 + .../runtime/metrics/camera/metric.py | 151 +++++ .../runtime/metrics/camera/pose_io.py | 40 ++ .../runtime/metrics/camera/preparation.py | 93 +++ .../runtime/metrics/camera/rpe.py | 126 ++++ .../runtime/metrics/camera/schemas.py | 49 ++ .../runtime/metrics/camera/trajectory.py | 93 +++ .../runtime/metrics/dynamic/__init__.py | 6 + .../runtime/metrics/dynamic/checklist.py | 81 +++ .../runtime/metrics/dynamic/metric.py | 97 +++ .../runtime/metrics/dynamic/preparation.py | 64 ++ .../runtime/metrics/dynamic/protocol.py | 227 +++++++ .../runtime/metrics/dynamic/scoring.py | 49 ++ .../runtime/metrics/dynamic/vlm.py | 396 ++++++++++++ .../runtime/metrics/environment/__init__.py | 8 + .../runtime/metrics/environment/artifacts.py | 86 +++ .../runtime/metrics/environment/backends.py | 95 +++ .../runtime/metrics/environment/geometry.py | 146 +++++ .../runtime/metrics/environment/metric.py | 293 +++++++++ .../metrics/environment/preparation.py | 67 +++ .../runtime/metrics/environment/protocol.py | 107 ++++ .../worldmembench/runtime/metrics/features.py | 130 ++++ .../metrics/general_video_quality/__init__.py | 9 + .../general_video_quality/aesthetic.py | 118 ++++ .../aesthetic_quality.py | 43 ++ .../metrics/general_video_quality/common.py | 72 +++ .../general_video_quality/imaging_quality.py | 32 + .../metrics/general_video_quality/musiq.py | 86 +++ .../runtime/metrics/human/__init__.py | 8 + .../metrics/human/appearance_features.py | 183 ++++++ .../runtime/metrics/human/body_tracking.py | 486 +++++++++++++++ .../runtime/metrics/human/evaluator.py | 314 ++++++++++ .../runtime/metrics/human/face_identity.py | 400 +++++++++++++ .../runtime/metrics/human/metric.py | 565 ++++++++++++++++++ .../metrics/human/person_mask_tracker.py | 224 +++++++ .../runtime/metrics/human/phase_protocol.py | 61 ++ .../runtime/metrics/human/preparation.py | 423 +++++++++++++ .../runtime/metrics/human/schemas.py | 155 +++++ .../metrics/human/target_association.py | 85 +++ .../runtime/metrics/object/__init__.py | 8 + .../metrics/object/appearance_features.py | 110 ++++ .../runtime/metrics/object/evaluator.py | 142 +++++ .../metrics/object/geometry_consistency.py | 180 ++++++ .../runtime/metrics/object/mask_alignment.py | 131 ++++ .../runtime/metrics/object/metric.py | 222 +++++++ .../metrics/object/object_redetector.py | 527 ++++++++++++++++ .../runtime/metrics/object/object_tracker.py | 317 ++++++++++ .../runtime/metrics/object/preparation.py | 72 +++ .../runtime/metrics/object/protocol.py | 28 + .../runtime/metrics/object/revisit.py | 46 ++ .../runtime/metrics/object/schemas.py | 95 +++ .../worldmembench/runtime/metrics/registry.py | 29 + .../runtime/metrics/revisit_protocol.py | 426 +++++++++++++ .../runners/worldmembench/runtime/paths.py | 13 + .../runtime/preparation/__init__.py | 4 + .../worldmembench/runtime/preparation/da3.py | 226 +++++++ .../worldmembench/runtime/preparation/vipe.py | 109 ++++ .../runners/worldmembench/runtime/results.py | 119 ++++ .../runners/worldmembench/runtime/video.py | 93 +++ .../benchmark/metrics/base_metrics.py | 109 +--- .../clip_mlp_aesthetic_metrics.py | 117 +--- .../evaluation/tasks/metrics/_shared/pyiqa.py | 78 +++ 98 files changed, 10577 insertions(+), 761 deletions(-) create mode 100644 tests/eval_core/test_worldmembench_runtime.py create mode 100644 tests/evaluation/execution/test_shared_quality.py create mode 100644 tests/evaluation/execution/test_worldmembench_camera.py create mode 100644 worldfoundry/data/benchmarks/requirements/worldmembench.txt create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/LICENSE create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/__init__.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/__main__.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/cli.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/context.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/data.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/errors.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/evaluate.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/execution.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/__init__.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/aggregation.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/base.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/__init__.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/metric.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/pose_io.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/preparation.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/rpe.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/schemas.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/trajectory.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/__init__.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/checklist.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/metric.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/preparation.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/protocol.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/scoring.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/vlm.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/__init__.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/artifacts.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/backends.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/geometry.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/metric.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/preparation.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/protocol.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/features.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/__init__.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/aesthetic.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/aesthetic_quality.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/common.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/imaging_quality.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/musiq.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/__init__.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/appearance_features.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/body_tracking.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/evaluator.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/face_identity.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/metric.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/person_mask_tracker.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/phase_protocol.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/preparation.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/schemas.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/target_association.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/__init__.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/appearance_features.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/evaluator.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/geometry_consistency.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/mask_alignment.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/metric.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/object_redetector.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/object_tracker.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/preparation.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/protocol.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/revisit.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/schemas.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/registry.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/revisit_protocol.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/paths.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/__init__.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/da3.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/vipe.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/results.py create mode 100644 worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/video.py create mode 100644 worldfoundry/evaluation/tasks/metrics/_shared/pyiqa.py diff --git a/Makefile b/Makefile index 9b1f67968..b7d8869e9 100644 --- a/Makefile +++ b/Makefile @@ -224,12 +224,15 @@ INFER_TENSOR_CONTRACTS = \ $(EVALUATION_CPU_CONTRACTS) EVALUATION_CPU_CONTRACTS = \ tests/eval_core/test_worldmembench_official_runner.py \ + tests/eval_core/test_worldmembench_runtime.py \ tests/evaluation/execution/test_resume_options.py \ tests/evaluation/execution/test_result_integrity.py \ tests/evaluation/execution/test_helper_reuse.py \ tests/evaluation/execution/test_pipeline_boundaries.py \ tests/evaluation/execution/test_generation_artifact_materialization.py \ tests/evaluation/execution/test_shared_scoring_backend.py \ + tests/evaluation/execution/test_shared_quality.py \ + tests/evaluation/execution/test_worldmembench_camera.py \ tests/evaluation/execution/test_vbench_appearance_reduction.py \ tests/evaluation/execution/test_worldscore_evaluator_lifecycle.py \ tests/evaluation/metrics/test_scalar_aggregation.py \ diff --git a/THIRD-PARTY-NOTICES b/THIRD-PARTY-NOTICES index d8fd06b2e..59bee2f13 100644 --- a/THIRD-PARTY-NOTICES +++ b/THIRD-PARTY-NOTICES @@ -6331,3 +6331,38 @@ Wan2.2 support assets retain their Apache-2.0 license. Original-File: WorldPlay2/WorldPlay2/LICENSE.txt license: cc-by-nc-4.0 + + +WorldMemBench evaluation protocols +--------------------------------- +Source: https://github.com/ZJU-REAL/WorldMemBench +Source revision: 1b79840143a98e017891e893b5d3ef75b93078a3 +License: MIT +Files: worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/ + +WorldFoundry retains the official sampling, revisit, alignment, aggregation, +VLM judging and result protocols. CLIP, LAION aesthetic heads and DINO model +loading use shared WorldFoundry components. Human and object masked DINO +features share one implementation. Model weights are distributed separately. + +MIT License + +Copyright (c) 2026 WorldMemBench + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/docs/fumadocs/content/docs/evaluation/benchmark-hub/wbench.mdx b/docs/fumadocs/content/docs/evaluation/benchmark-hub/wbench.mdx index 6e8e8943b..6ab70d36f 100644 --- a/docs/fumadocs/content/docs/evaluation/benchmark-hub/wbench.mdx +++ b/docs/fumadocs/content/docs/evaluation/benchmark-hub/wbench.mdx @@ -17,6 +17,12 @@ WorldFoundry **checks in the metric runtime** (HPSv3, TransNetV2, PAVRM, DreamSi +Imaging quality uses the official PyIQA `musiq` variant (KonIQ), with its +checkpoint resolved explicitly through `wbench_quality_metric_assets`. +`WORLDFOUNDRY_WBENCH_MUSIQ_KONIQ_CKPT` can select an already staged checkpoint. +Its preprocessing and per-frame normalization retain the official protocol; +WorldMemBench's MUSIQ-SPAQ is a different metric backend profile. + ## How to run ```bash diff --git a/docs/fumadocs/content/docs/evaluation/benchmark-hub/wbench.zh.mdx b/docs/fumadocs/content/docs/evaluation/benchmark-hub/wbench.zh.mdx index 6f7891668..0db067f10 100644 --- a/docs/fumadocs/content/docs/evaluation/benchmark-hub/wbench.zh.mdx +++ b/docs/fumadocs/content/docs/evaluation/benchmark-hub/wbench.zh.mdx @@ -17,6 +17,10 @@ WorldFoundry **检入指标 runtime**(HPSv3、TransNetV2、PAVRM、DreamSim、 +成像质量采用官方 PyIQA `musiq` 版本(KonIQ),通过 `wbench_quality_metric_assets` +显式解析权重。可用 `WORLDFOUNDRY_WBENCH_MUSIQ_KONIQ_CKPT` 指向已准备的权重。 +预处理和逐帧归一化保留官方协议;WorldMemBench 的 MUSIQ-SPAQ 属于不同的后端配置。 + ## 如何运行 ```bash diff --git a/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.mdx b/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.mdx index 479ec30aa..17ab271cb 100644 --- a/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.mdx +++ b/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.mdx @@ -1,103 +1,104 @@ --- title: WorldMemBench -description: Official video-world-model memory metrics, per-model result import, and upstream runtime execution. +description: In-tree visual-memory and camera-control evaluation with shared model components. --- -WorldMemBench evaluates video world models through human, object, environment, -and dynamic memory, alongside video quality and camera following. WorldFoundry -exposes all **12 official metrics**, imports their per-model summaries, and can -invoke the official metric runtime on staged videos. Model generation remains -separate from evaluation. - -Code: [ZJU-REAL/WorldMemBench](https://github.com/ZJU-REAL/WorldMemBench) (MIT). -Paper: [WorldMemBench](https://arxiv.org/abs/2610.19063). -Data: [rescuerz/WorldMemBench](https://huggingface.co/datasets/rescuerz/WorldMemBench). -The integration targets upstream revision `1b79840143a98e017891e893b5d3ef75b93078a3`. - -## Import official results - -Result import needs no GPU or metric checkpoints: - -```bash -worldfoundry-eval zoo benchmark-run \ - --benchmark-id worldmembench \ - --mode official-validation \ - --official-results-path /path/to/official_eval_output \ - --output-dir tmp/worldmembench/import \ - --json -``` - -The input may be the official evaluation output root, its `metrics` directory, -or one `metrics//summary.json` file. `cases.jsonl` files are imported -when present. A directory containing several model groups needs an explicit -selection through the runner: +## Scope + +[WorldMemBench](https://github.com/ZJU-REAL/WorldMemBench) evaluates human, +object, environment and dynamic memory, camera control, and video quality. +WorldFoundry runs all twelve official metric protocols in tree, adapted from +revision `1b79840143a98e017891e893b5d3ef75b93078a3` under MIT. An upstream +WorldMemBench checkout or Python package is not required. + +| Metric ID | Official score field | Direction | +| --- | --- | --- | +| `imaging_quality` | `score` | Higher | +| `aesthetic_quality` | `score` | Higher | +| `human_identity` | `score` | Higher | +| `human_appearance` | `score` | Higher | +| `object_geometry` | `score` | Higher | +| `object_appearance` | `score` | Higher | +| `environment_visible_memory` | `score` | Higher | +| `environment_depth_fidelity` | `depth_spearman` | Higher | +| `environment_scene_reconstruction` | `reconstruction_score` | Higher | +| `dynamic_memory` | `score` | Higher | +| `camera_rotation` | `camera_rotation_rpe_path` | Lower | +| `camera_translation` | `camera_translation_rpe` | Lower | + +Scores retain their official units, including negative depth correlations and +unbounded camera errors. There is no cross-metric average or composite score. + +## Shared components and protocol + +OpenAI CLIP, LAION aesthetic heads and DINO checkpoint loading use canonical +WorldFoundry components. Human and object appearance share one masked DINOv2 +extractor. Within a run, identical DINO checkpoint/device pairs share the model; +preprocessing and batch-size state remain protocol-specific. +WorldScore and MemoBench use the shared LAION MLP; MIND, WBench and +WorldArena use the shared linear head. These heads are not interchangeable. +WorldScore and 4DWorldBench share their 512-square IQA preprocessing and base +implementation; that transform is not applied to WorldMemBench. + +WorldMemBench retains MUSIQ-SPAQ, its aspect-preserving OpenCV resize, stride-4 +quality sampling, score normalization, and equal-weight video aggregation. +Revisit segmentation, frozen target association, masked crops, camera alignment, +scene reconstruction and Dynamic Memory's event-gated VLM checklist retain +the official protocols. + +## Evaluate generated outputs + +Use the official `model/category/trajectory/case/video.mp4` layout or an official +JSONL manifest. Keep annotations, reference images, trajectory files and target +poses with each case. The dataset is +[`rescuerz/WorldMemBench`](https://huggingface.co/datasets/rescuerz/WorldMemBench). + +Install the selected metric dependencies into the prepared GPU environment; +`worldfoundry/data/benchmarks/requirements/worldmembench.txt` lists the complete +set without replacing its Torch/CUDA installation. Stage the official +checkpoints explicitly. `--weights-dir` accepts the upstream metric asset +layout; YAML checkpoint overrides can point to already staged shared assets. ```bash python -m worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner \ - --official-results-path /path/to/official_eval_output \ - --model-id model_a \ - --metrics human_identity,object_geometry,camera_rotation \ - --output-dir tmp/worldmembench/model_a \ - --json + --run-official --input-path \ + --weights-dir --config \ + --output-dir --json ``` -Outputs are `scorecard.json`, `raw_metric_table.jsonl`, and -`per_sample_scores.jsonl`. The scorecard retains official score fields, grouped -category/trajectory results, sample counts, and case statuses. Models are not -averaged together. `camera_rotation` and `camera_translation` are errors: -**lower is better**. Metric values retain upstream units, including signed depth -correlations and unbounded camera errors; no clipping, inversion, or -cross-metric composite is applied. +Repeat `--metric` to select metrics. `--gpus 0,1` shards videos across visible +GPUs. `--expected-revisits 1|2|3` selects the official fixed revisit protocol; +otherwise manifest values or official return detection apply. -## Execute official metrics +DA3-Streaming and ViPE preparation reuse existing `da3/results.npz`, +`da3/scene_geometry.npz` and `vipe/results.npz` artifacts. To generate missing +artifacts, configure `da3.source`, `da3.dinov2_source` and `vipe.source` for the +official model runtimes, with optional `python` and `weights` overrides in each +section. DA3-Streaming is not replaced with single-batch DA3, nor ViPE's default +pipeline with WorldFoundry's reduced pose-only profile. Dynamic Memory requires +the official VLM configuration; API keys stay in its configured environment +variable, not CLI arguments. -Provision the official checkout and selected metric dependencies following its -[installation instructions](https://github.com/ZJU-REAL/WorldMemBench#installation). -Use the unified WorldFoundry environment for orchestration; if upstream needs a -separate interpreter, select it with `--python` rather than replacing the -unified environment's Torch/CUDA packages. +## Import official results ```bash -python -m worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner \ - --run-official \ - --repo-root /path/to/WorldMemBench \ - --python /path/to/upstream-env/bin/python \ - --input-path /path/to/manifest.jsonl \ - --config /path/to/worldmembench.yaml \ - --weights-dir /path/to/weights \ - --metrics imaging_quality,aesthetic_quality \ - --device cuda:0 \ - --output-dir tmp/worldmembench/official \ - --json +worldfoundry-eval zoo benchmark-run --benchmark-id worldmembench \ + --mode official-validation --official-results-path \ + --result-model-id --output-dir --json ``` -`--repo-root` loads the checkout's `src/worldmembench`; omit it when the selected -interpreter already has `worldmembench` installed. Input can also be a video or -an official video tree: `////video.mp4`. -Use upstream JSONL manifests to retain model IDs, camera protocols, and per-case -annotations. `--expected-revisits 1|2|3` selects the fixed human/object revisit -protocol; omitting it preserves upstream manifest/detection behavior. -`--gpus 0,1` delegates process sharding to upstream instead of using `--device`. -Each execution needs a fresh output directory; official results go under -`official/`, with stdout/stderr logs beside the scorecard. - -The unified CLI can also score staged videos using `--generated-artifact-dir` -and `--mode official-run`. Configure runtime paths through -`WORLDFOUNDRY_WORLDMEMBENCH_ROOT`, `WORLDFOUNDRY_WORLDMEMBENCH_PYTHON`, -`WORLDFOUNDRY_WORLDMEMBENCH_CONFIG`, and -`WORLDFOUNDRY_WORLDMEMBENCH_WEIGHTS_DIR`. For import, use -`WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH`. +Accepts an evaluation output root, its `metrics/` directory, or a metric's +`summary.json`. Multiple model groups require explicit selection. Summary +coverage, group breakdowns and adjacent `cases.jsonl` records are retained in +`scorecard.json`, `raw_metric_table.jsonl` and `per_sample_scores.jsonl`. ## Limitations -Official execution needs the selected metric checkpoints, DA3/ViPE subprocess -environments, and a VLM endpoint for Dynamic Memory. Keep backend interpreters, -checkpoint paths, and VLM settings in the upstream YAML config; API keys stay -in the upstream-configured environment variable. Metric implementations and -weights are not vendored. GPU/checkpoint and complete-dataset replay remain -pending, so scorecards do not claim full-suite or leaderboard validity. -`model_failure` retains its official scored result; `not_applicable` is not -zero-filled. `eval_failed` or `data_invalid` cases retain their coverage and -make the runner exit nonzero, including when partial scores are available. +GPU/checkpoint and complete-dataset replay remain pending; in-tree execution +and weight-free tests do not establish full-suite or leaderboard validity. +DA3-Streaming and full-profile ViPE inference require their configured model +runtimes unless artifacts are already available. `model_failure` retains its +official scored result; `not_applicable` is not zero-filled. `eval_failed` and +`data_invalid` retain coverage and cause a nonzero exit, even with partial scores. diff --git a/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.zh.mdx b/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.zh.mdx index c415ee297..101c5d406 100644 --- a/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.zh.mdx +++ b/docs/fumadocs/content/docs/evaluation/benchmark-hub/worldmembench.zh.mdx @@ -1,90 +1,89 @@ --- title: WorldMemBench -description: 视频世界模型记忆评测:官方指标、按模型导入结果及调用上游评测运行时。 +description: 复用共享模型组件的仓内视觉记忆与相机控制评测。 --- -WorldMemBench 从人物、物体、环境和动态记忆,以及视频质量与相机跟随能力评测视频世界模型。 -WorldFoundry 提供 **12 项官方指标**的统一入口,支持按模型导入官方结果,或调用官方指标运行时 -评测已生成的视频。模型生成与评测分开执行。 - -代码:[ZJU-REAL/WorldMemBench](https://github.com/ZJU-REAL/WorldMemBench)(MIT)。 -论文:[WorldMemBench](https://arxiv.org/abs/2610.19063)。 -数据:[rescuerz/WorldMemBench](https://huggingface.co/datasets/rescuerz/WorldMemBench)。 -集成对应上游版本 `1b79840143a98e017891e893b5d3ef75b93078a3`。 - -## 导入官方结果 - -结果导入不需要 GPU 或指标权重: - -```bash -worldfoundry-eval zoo benchmark-run \ - --benchmark-id worldmembench \ - --mode official-validation \ - --official-results-path /path/to/official_eval_output \ - --output-dir tmp/worldmembench/import \ - --json -``` - -输入可以是官方评测输出根目录、其中的 `metrics` 目录,或单个 -`metrics//summary.json`。存在 `cases.jsonl` 时同时导入逐样本结果。 -多模型结果必须通过 runner 显式选择一个模型: +## 范围 + +[WorldMemBench](https://github.com/ZJU-REAL/WorldMemBench) 评测人物、物体、环境和 +动态记忆,以及相机控制和视频质量。WorldFoundry 的 12 项官方指标协议均在仓内执行, +基于 MIT 许可的上游版本 `1b79840143a98e017891e893b5d3ef75b93078a3` 适配, +不需要额外检出或安装 WorldMemBench。 + +| 指标 ID | 官方分数字段 | 方向 | +| --- | --- | --- | +| `imaging_quality` | `score` | 越高越好 | +| `aesthetic_quality` | `score` | 越高越好 | +| `human_identity` | `score` | 越高越好 | +| `human_appearance` | `score` | 越高越好 | +| `object_geometry` | `score` | 越高越好 | +| `object_appearance` | `score` | 越高越好 | +| `environment_visible_memory` | `score` | 越高越好 | +| `environment_depth_fidelity` | `depth_spearman` | 越高越好 | +| `environment_scene_reconstruction` | `reconstruction_score` | 越高越好 | +| `dynamic_memory` | `score` | 越高越好 | +| `camera_rotation` | `camera_rotation_rpe_path` | 越低越好 | +| `camera_translation` | `camera_translation_rpe` | 越低越好 | + +分数保留官方单位,包括负的深度相关系数和无上界的相机误差,不构造跨指标平均分或综合分。 + +## 共享组件与协议 + +OpenAI CLIP、LAION 美学头和 DINO 权重加载复用 WorldFoundry 的规范组件;人物和物体 +外观共用一份掩码 DINOv2 特征提取实现。同次评测中,相同权重路径与设备共享模型, +预处理和批大小状态保持独立。WorldScore、MemoBench 使用共享 LAION MLP, +MIND、WBench、WorldArena 使用共享线性头,两种头不能互换。WorldScore 与 +4DWorldBench 共用其 512 方形 IQA 预处理和基类,该变换不应用于 WorldMemBench。 + +WorldMemBench 保留 MUSIQ-SPAQ、保持宽高比的 OpenCV 缩放、每 4 帧采样、分数归一化 +和视频等权聚合。重访分段、冻结目标关联、掩码裁剪、相机对齐、场景重建,以及动态记忆 +按事件门控的 VLM 问答均遵循官方协议。 + +## 评测生成输出 + +使用官方 `model/category/trajectory/case/video.mp4` 目录结构或官方 JSONL 清单。 +每个 case 保留标注、参考图、轨迹文件及目标相机位姿。数据集为 +[`rescuerz/WorldMemBench`](https://huggingface.co/datasets/rescuerz/WorldMemBench)。 + +在已准备的 GPU 环境中安装所选指标的依赖; +`worldfoundry/data/benchmarks/requirements/worldmembench.txt` 列出完整依赖, +不替换环境中的 Torch/CUDA。显式准备官方权重;`--weights-dir` 接受上游指标资产布局, +YAML 中的 checkpoint 覆盖项可以指向已经准备的共享资产。 ```bash python -m worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner \ - --official-results-path /path/to/official_eval_output \ - --model-id model_a \ - --metrics human_identity,object_geometry,camera_rotation \ - --output-dir tmp/worldmembench/model_a \ - --json + --run-official --input-path \ + --weights-dir --config \ + --output-dir --json ``` -输出为 `scorecard.json`、`raw_metric_table.jsonl` 和 `per_sample_scores.jsonl`。 -Scorecard 保留官方分数字段、类别/轨迹分组、样本数与状态,不对不同模型求平均。 -`camera_rotation` 和 `camera_translation` 是误差,**越低越好**。 -所有数值沿用官方单位,包括可为负的深度相关系数和无上界的相机误差;不裁剪、取倒数, -也不构造跨指标总分。 +重复 `--metric` 可选择指标,`--gpus 0,1` 按视频划分到可见 GPU。 +`--expected-revisits 1|2|3` 选择官方固定重访协议;否则使用清单中的值或官方返回检测。 -## 执行官方指标 +DA3-Streaming 和 ViPE 预处理复用已有的 `da3/results.npz`、 +`da3/scene_geometry.npz` 和 `vipe/results.npz`。如需生成缺失资产,配置 +`da3.source`、`da3.dinov2_source` 和 `vipe.source` 指向官方模型运行时, +各节还可指定 `python` 和 `weights`。不将 DA3-Streaming 替换为单批 DA3, +也不将 ViPE 默认流水线替换为 WorldFoundry 的精简位姿配置。动态记忆需要官方 VLM +配置;API 密钥保留在配置指定的环境变量中,不通过命令行传递。 -先按上游的[安装说明](https://github.com/ZJU-REAL/WorldMemBench#installation) -准备官方代码与所选指标依赖。WorldFoundry 调度使用统一环境;上游需要独立解释器时通过 -`--python` 指定,不替换统一环境中的 Torch/CUDA。 +## 导入官方结果 ```bash -python -m worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner \ - --run-official \ - --repo-root /path/to/WorldMemBench \ - --python /path/to/upstream-env/bin/python \ - --input-path /path/to/manifest.jsonl \ - --config /path/to/worldmembench.yaml \ - --weights-dir /path/to/weights \ - --metrics imaging_quality,aesthetic_quality \ - --device cuda:0 \ - --output-dir tmp/worldmembench/official \ - --json +worldfoundry-eval zoo benchmark-run --benchmark-id worldmembench \ + --mode official-validation --official-results-path \ + --result-model-id --output-dir --json ``` -`--repo-root` 使用该 checkout 的 `src/worldmembench`;指定解释器已安装 -`worldmembench` 时可省略。输入也可以是单个视频,或官方视频树: -`////video.mp4`。使用上游 JSONL manifest -可保留模型 ID、相机协议与逐样本标注。`--expected-revisits 1|2|3` 指定人物/物体的固定重访 -协议;省略时沿用官方 manifest 或自动检测行为。`--gpus 0,1` 将进程分片交给上游, -代替 `--device`。每次执行使用新的输出目录;官方结果位于 `official/`,标准输出与错误日志 -位于 scorecard 同级。 - -统一 CLI 也支持通过 `--generated-artifact-dir` 与 `--mode official-run` 评测已有视频。 -运行时路径使用 `WORLDFOUNDRY_WORLDMEMBENCH_ROOT`、 -`WORLDFOUNDRY_WORLDMEMBENCH_PYTHON`、`WORLDFOUNDRY_WORLDMEMBENCH_CONFIG` 和 -`WORLDFOUNDRY_WORLDMEMBENCH_WEIGHTS_DIR` 配置。 -结果导入可使用 `WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH`。 +接受评测输出根目录、其中的 `metrics/` 目录或单指标的 `summary.json`。 +多模型结果必须显式选择模型。汇总覆盖率、分组明细和相邻的 `cases.jsonl` 记录 +保留在 `scorecard.json`、`raw_metric_table.jsonl` 和 `per_sample_scores.jsonl` 中。 ## 限制 -官方计算需要所选指标权重、DA3/ViPE 子进程环境,以及 Dynamic Memory 的 VLM 服务。 -后端解释器、权重路径与 VLM 设置放在上游 YAML 中;API 密钥使用上游配置的环境变量。 -本仓库不复制指标实现或权重。GPU/真实权重与完整数据集回放尚待验证,因此 scorecard -不声明完整评测或排行榜有效性。`model_failure` 保留官方计分结果;`not_applicable` -不补零。`eval_failed` 或 `data_invalid` 保留覆盖率,且即使存在部分分数,runner -也返回非零退出码。 +GPU/真实权重和完整数据集回放尚未完成;仓内执行与免权重测试不代表全套验证或排行榜 +有效性。没有预计算资产时,DA3-Streaming 和完整配置 ViPE 推理仍需要配置对应模型 +运行时。`model_failure` 保留官方赋分,`not_applicable` 不填零。 +`eval_failed` 和 `data_invalid` 保留覆盖率,并使运行器返回非零,即使已有部分分数。 diff --git a/docs/fumadocs/lib/benchmark-catalog-status.json b/docs/fumadocs/lib/benchmark-catalog-status.json index f63c0b1ba..58e88ec5f 100644 --- a/docs/fumadocs/lib/benchmark-catalog-status.json +++ b/docs/fumadocs/lib/benchmark-catalog-status.json @@ -1826,9 +1826,9 @@ ], "name": "WorldMemBench", "requiresApi": false, - "runtimeKind": "external_official_repo", - "summary": "Imports per-model memory metrics or invokes the official WorldMemBench runtime", - "summaryZh": "按模型导入记忆指标或调用官方 WorldMemBench 运行时", + "runtimeKind": "in_tree_official_runtime", + "summary": "Runs twelve in-tree memory and camera metrics with shared model components", + "summaryZh": "复用共享模型组件,在仓内执行十二项记忆与相机指标", "verificationStatus": "official_runtime_ready_real_data_validation_pending" }, "worldmodelbench": { diff --git a/tests/eval_core/test_wbench_official_runner.py b/tests/eval_core/test_wbench_official_runner.py index 9906dcaed..0c765914e 100644 --- a/tests/eval_core/test_wbench_official_runner.py +++ b/tests/eval_core/test_wbench_official_runner.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 WorldFoundry contributors +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations import json @@ -7,6 +10,8 @@ from worldfoundry.evaluation.tasks.execution.runners.wbench import run_wbench_official_runner +pytestmark = pytest.mark.unit + def test_wbench_official_runner_normalizes_report_json(tmp_path: Path) -> None: report_path = tmp_path / "report.json" @@ -42,11 +47,24 @@ def test_wbench_official_runner_normalizes_report_json(tmp_path: Path) -> None: scorecard = json.loads((output_dir / "scorecard.json").read_text(encoding="utf-8")) raw_rows = [ - json.loads(line) - for line in (output_dir / "raw_metric_table.jsonl").read_text(encoding="utf-8").splitlines() + json.loads(line) for line in (output_dir / "raw_metric_table.jsonl").read_text(encoding="utf-8").splitlines() ] assert exit_code == 0 assert scorecard["evaluation"]["kind"] == "wbench_result_normalizer" - assert scorecard["metrics"]["leaderboard"]["quality_score"] == pytest.approx(0.7) - assert scorecard["metrics"]["leaderboard"]["wbench_average"] == pytest.approx(0.64) + assert scorecard["metrics"]["leaderboard"] == pytest.approx( + { + "aesthetic_quality": 0.8, + "imaging_quality": 0.6, + "scene_adherence": 0.7, + "navigation_trajectory": 0.5, + "background_consistency": 0.9, + "visual_plausibility": 0.4, + } + ) + assert scorecard["metrics"]["per_metric"]["quality_score"]["raw_score"] is None + assert scorecard["metrics"]["per_metric"]["wbench_average"]["raw_score"] is None + assert len(scorecard["evaluation"]["missing_component_metrics"]) == 16 + assert not scorecard["evaluation"]["benchmark_comparable"] + assert not scorecard["leaderboard_valid"] assert raw_rows[-1]["metric_id"] == "wbench_average" + assert raw_rows[-1]["raw_score"] is None diff --git a/tests/eval_core/test_worldmembench_official_runner.py b/tests/eval_core/test_worldmembench_official_runner.py index 80855cc42..2918b27bd 100644 --- a/tests/eval_core/test_worldmembench_official_runner.py +++ b/tests/eval_core/test_worldmembench_official_runner.py @@ -239,15 +239,11 @@ def test_missing_inputs_write_standard_failed_artifacts(tmp_path: Path) -> None: @pytest.mark.parametrize("upstream_code", [0, 1]) -def test_official_execution_delegates_to_upstream_and_respects_exit_status( +def test_official_execution_uses_in_tree_runtime_and_respects_exit_status( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, upstream_code: int, ) -> None: - root = tmp_path / "upstream" - module = root / "src/worldmembench/__main__.py" - module.parent.mkdir(parents=True) - module.touch() manifest = tmp_path / "inputs.jsonl" manifest.touch() config = tmp_path / "config.yaml" @@ -256,8 +252,7 @@ def test_official_execution_delegates_to_upstream_and_respects_exit_status( def execute(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[str]: calls.append(command) - assert kwargs["cwd"] == root - assert kwargs["env"]["PYTHONPATH"].split(":")[0] == str(root / "src") + assert "cwd" not in kwargs assert kwargs["timeout"] == 7200 result_dir = Path(command[command.index("--output") + 1]) _write_result(result_dir) @@ -267,8 +262,6 @@ def execute(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[st card, code = _run( tmp_path, "--run-official", - "--repo-root", - str(root), "--input-path", str(manifest), "--config", @@ -282,7 +275,13 @@ def execute(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[st ) assert code == upstream_code command = calls[0] - assert command[:5] == [sys.executable, "-m", "worldmembench", "eval", str(manifest)] + assert command[:5] == [ + sys.executable, + "-m", + "worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime", + "eval", + str(manifest), + ] assert command[command.index("--metrics") + 1] == "human_identity" assert "--gpus" in command and "--device" not in command assert command[command.index("--config") + 1] == str(config) @@ -290,7 +289,7 @@ def execute(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[st assert card["integration_evidence"] is (upstream_code == 0) assert card["evaluation"]["available"] is (upstream_code == 0) assert card["validation"]["full_suite_complete"] is False - assert (tmp_path / "out/upstream_stdout.log").is_file() + assert (tmp_path / "out/runtime_stdout.log").is_file() def test_official_execution_does_not_reuse_stale_outputs(tmp_path: Path) -> None: @@ -334,19 +333,19 @@ def test_standard_benchmark_contract_and_result_import(tmp_path: Path) -> None: def test_standard_official_run_dispatches_generated_videos(tmp_path: Path) -> None: - upstream = tmp_path / "upstream/src/worldmembench" - upstream.mkdir(parents=True) - (upstream / "__init__.py").touch() + interpreter = tmp_path / "metric-python" payload = {"num_cases": 1, "results": {"all": {"num_cases": 1, "scores": {"score": 0.5}}}} - (upstream / "__main__.py").write_text( + interpreter.write_text( + f"#!{sys.executable}\n" "import json, pathlib, sys\n" - "assert sys.argv[1] == 'eval'\n" + "assert sys.argv[1:4] == ['-m', 'worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime', 'eval']\n" "output = pathlib.Path(sys.argv[sys.argv.index('--output') + 1])\n" "directory = output / 'metrics' / 'imaging_quality'\n" "directory.mkdir(parents=True)\n" f"(directory / 'summary.json').write_text({json.dumps(payload)!r})\n", encoding="utf-8", ) + interpreter.chmod(0o755) videos = tmp_path / "videos" videos.mkdir() (videos / "case_001.mp4").touch() @@ -355,7 +354,7 @@ def test_standard_official_run_dispatches_generated_videos(tmp_path: Path) -> No output_dir=tmp_path / "official-run", mode="official-run", generated_artifact_dir=videos, - env_overrides={"WORLDFOUNDRY_WORLDMEMBENCH_ROOT": str(upstream.parent.parent)}, + env_overrides={"WORLDFOUNDRY_WORLDMEMBENCH_PYTHON": str(interpreter)}, ) assert result.ok, result.to_dict() card = json.loads(result.scorecard_path.read_text()) diff --git a/tests/eval_core/test_worldmembench_runtime.py b/tests/eval_core/test_worldmembench_runtime.py new file mode 100644 index 000000000..f491d342a --- /dev/null +++ b/tests/eval_core/test_worldmembench_runtime.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 WorldFoundry contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Weight-free checks of the native WorldMemBench protocol and runner.""" + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.evaluate import evaluate +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics import registry +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.dynamic.checklist import CATEGORIES +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.dynamic.scoring import ( + aggregate_scores, + score_answers, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.general_video_quality.aesthetic_quality import ( + AestheticQualityMetric, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.general_video_quality.imaging_quality import ( + ImagingQualityMetric, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.human.evaluator import ( + MaskedDinoV2Features as HumanFeatures, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.object.evaluator import ( + MaskedDinoV2Features as ObjectFeatures, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.registry import MetricRegistry +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.video import DecodedFrame, VideoInfo +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.worldmembench_metrics import ( + METRIC_SPECS, + load_official_results, +) + +pytestmark = pytest.mark.unit + + +class _Reader: + def probe(self, path: str) -> VideoInfo: + return VideoInfo(9 if Path(path).stem == "long" else 5, 24.0, 2, 2) + + def read_indices(self, path: str, indices: list[int]) -> list[DecodedFrame]: + value = 80 if Path(path).stem == "long" else 20 + return [DecodedFrame(index, np.full((2, 2, 3), value, dtype=np.uint8)) for index in indices] + + +class _Backend: + def score(self, frames: list[np.ndarray]) -> list[float]: + return [float(frame[0, 0, 0]) for frame in frames] + + +def test_all_official_metrics_are_available_without_upstream_package() -> None: + assert set(registry.names()) == set(METRIC_SPECS) + completed = subprocess.run( + [sys.executable, "-m", "worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime", "list-metrics"], + capture_output=True, + text=True, + check=True, + ) + assert set(completed.stdout.splitlines()) == set(METRIC_SPECS) + + +def test_human_and_object_share_masked_feature_extraction() -> None: + assert HumanFeatures is ObjectFeatures + frame = np.arange(4 * 5 * 3, dtype=np.uint8).reshape(4, 5, 3) + mask = np.zeros((4, 5), dtype=bool) + mask[1:3, 1:4] = True + mask[1, 2] = False + crop = HumanFeatures._masked_crop(frame, mask) + assert crop is not None + expected = frame[1:3, 1:4].copy() + expected[0, 1] = 0 + np.testing.assert_array_equal(crop, expected) + assert HumanFeatures._masked_crop(frame, np.zeros_like(mask)) is None + + +@pytest.mark.parametrize("metric_type,normalizer", [(ImagingQualityMetric, 100.0), (AestheticQualityMetric, 10.0)]) +def test_quality_uses_stride_four_clipping_and_equal_video_weights( + tmp_path: Path, metric_type: Any, normalizer: float +) -> None: + manifest = tmp_path / "inputs.jsonl" + rows = [] + for name in ("long", "short"): + (tmp_path / f"{name}.mp4").touch() + rows.append({"item_id": name, "model_id": "model_a", "video": f"{name}.mp4"}) + manifest.write_text("".join(json.dumps(row) + "\n" for row in rows)) + metrics = MetricRegistry() + metric = metric_type(backend=_Backend(), reader=_Reader()) + metrics.register(metric) + output, success = evaluate( + manifest, + metrics=[metric.name], + metric_registry=metrics, + device="cpu", + weights_dir=tmp_path / "assets", + output=tmp_path / "native", + ) + assert success + model, extracted, cases = load_official_results(output) + assert model == "model_a" + expected = (min(1.0, 80.0 / normalizer) + min(1.0, 20.0 / normalizer)) / 2 + assert extracted[metric.name]["raw_score"] == pytest.approx(expected) + assert [row["metadata"]["sampled_frame_indices"] for row in cases] == [[0, 4, 8], [0, 4]] + assert [row["num_units"] for row in cases] == [3, 2] + + +@pytest.mark.parametrize("started,expected", [("yes", 1.0), ("no", 0.0)]) +def test_dynamic_memory_event_gate_applies_to_all_eight_questions(started: str, expected: float) -> None: + categories = [category for category in CATEGORIES for _ in range(8 // len(CATEGORIES))] + assert len(categories) == 8 + checklist = SimpleNamespace( + questions=[ + SimpleNamespace( + id=f"q{index}", expected_answer="yes", category=category, question="question", evidence_window="return" + ) + for index, category in enumerate(categories) + ] + ) + response = SimpleNamespace( + event_started=SimpleNamespace(answer=started), + answers=[SimpleNamespace(answer="yes", reason="visible", evidence_image_numbers=(1,)) for _ in range(8)], + ) + units = score_answers(checklist, response, [4]) + assert aggregate_scores(units)["score"] == expected + assert all(unit.metadata["ungated_score"] == 1.0 for unit in units) + assert all(unit.metadata["evidence_frame_indices"] == [4] for unit in units) diff --git a/tests/evaluation/execution/test_shared_quality.py b/tests/evaluation/execution/test_shared_quality.py new file mode 100644 index 000000000..7fef352a9 --- /dev/null +++ b/tests/evaluation/execution/test_shared_quality.py @@ -0,0 +1,269 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 WorldFoundry contributors +# SPDX-License-Identifier: Apache-2.0 + +"""CPU numerical contracts for shared quality components; no pretrained assets.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +from PIL import Image + +from worldfoundry.evaluation.tasks.metrics._shared.aesthetic import ( + load_laion_aesthetic_linear_head, + load_laion_aesthetic_mlp_head, +) +from worldfoundry.evaluation.tasks.metrics._shared.pyiqa import ImageQualityMetric, PyIQAMetric + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def torch_cpu() -> Any: + torch = pytest.importorskip("torch") + previous = torch.get_num_threads() + torch.set_num_threads(2) + yield torch + torch.set_num_threads(previous) + + +@pytest.mark.parametrize("head", ["linear", "mlp"]) +def test_shared_laion_head_matches_reference_weights_and_predictions(tmp_path: Path, torch_cpu: Any, head: str) -> None: + torch, nn = torch_cpu, torch_cpu.nn + if head == "linear": + reference = nn.Linear(768, 1).eval() + state = reference.state_dict() + load = load_laion_aesthetic_linear_head + else: + reference = nn.Sequential( + nn.Linear(768, 1024), + nn.Dropout(0.2), + nn.Linear(1024, 128), + nn.Dropout(0.2), + nn.Linear(128, 64), + nn.Dropout(0.1), + nn.Linear(64, 16), + nn.Linear(16, 1), + ).eval() + state = {f"layers.{key}": value for key, value in reference.state_dict().items()} + load = load_laion_aesthetic_mlp_head + checkpoint = tmp_path / "head.pth" + torch.save(state, checkpoint) + actual = load(checkpoint) + features = torch.arange(3 * 768, dtype=torch.float32).reshape(3, 768) / 768 + with torch.inference_mode(): + torch.testing.assert_close(actual(features), reference(features), rtol=0, atol=0) + assert not actual.training + + +def test_iqa_preprocessing_matches_official_resize_for_paths_and_pil(tmp_path: Path, torch_cpu: Any) -> None: + from torchvision import transforms + + class Metric(ImageQualityMetric): + def _compute_scores(self, *args: Any) -> None: + pass + + image = Image.fromarray(np.arange(19 * 31 * 3, dtype=np.uint8).reshape(19, 31, 3)) + path = tmp_path / "image.png" + image.save(path) + metric = Metric(device="cpu") + transform = transforms.Compose([transforms.Resize((512, 512)), transforms.ToTensor()]) + expected = torch_cpu.stack([transform(image), transform(image)]) + torch_cpu.testing.assert_close(metric._process_image([str(path), image]), expected, rtol=0, atol=0) + for reference in (str(path), image): + rendered, references = metric._process_images([str(path), image], reference) + torch_cpu.testing.assert_close(rendered, expected, rtol=0, atol=0) + torch_cpu.testing.assert_close(references, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("variant", ["musiq", "musiq-spaq", "musiq-paq2piq"]) +def test_shared_iqa_preserves_the_selected_checkpoint_variant(monkeypatch: pytest.MonkeyPatch, variant: str) -> None: + class Metric(PyIQAMetric): + def _compute_scores(self, *args: Any) -> None: + pass + + calls: list[str] = [] + model = SimpleNamespace(to=lambda device: device) + monkeypatch.setitem(sys.modules, "pyiqa", SimpleNamespace(create_metric=lambda name: calls.append(name) or model)) + metric = Metric(variant, device="cpu") + assert calls == [variant] + assert metric._metric == "cpu" + + +def test_dinov2_loader_preserves_explicit_cls_and_pooler_semantics( + monkeypatch: pytest.MonkeyPatch, torch_cpu: Any +) -> None: + from worldfoundry.base_models.perception_core.general_perception.dino_embeddings import load_dinov2_feature_model + + class Model(torch_cpu.nn.Module): + def forward(self, pixel_values: Any) -> Any: + return SimpleNamespace(last_hidden_state=pixel_values, pooler_output=pixel_values[:, 1]) + + calls: list[tuple[str, bool]] = [] + + def load(source: str, *, local_files_only: bool) -> Any: + calls.append((source, local_files_only)) + return Model() + + monkeypatch.setitem(sys.modules, "transformers", SimpleNamespace(AutoModel=SimpleNamespace(from_pretrained=load))) + features = torch_cpu.arange(12, dtype=torch_cpu.float32).reshape(2, 3, 2) + cls_model = load_dinov2_feature_model("local-dinov2-large", device="cpu") + pooler_model = load_dinov2_feature_model("local-dinov2-base", use_pooler_output=True) + torch_cpu.testing.assert_close(cls_model(features), features[:, 0], rtol=0, atol=0) + torch_cpu.testing.assert_close(pooler_model(features), features[:, 1], rtol=0, atol=0) + assert calls == [("local-dinov2-large", True), ("local-dinov2-base", True)] + assert not cls_model.training + + +def test_masked_dino_batch_and_single_paths_keep_the_official_transform(torch_cpu: Any) -> None: + from torchvision import transforms + + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.features import ( + MaskedDinoV2Features, + ) + + class Model(torch_cpu.nn.Module): + def forward(self, tensor: Any) -> Any: + assert tensor.shape[-2:] == (14, 14) + return tensor.mean(dim=(-2, -1)) + + model = Model().eval() + extractor = MaskedDinoV2Features(model=model, device="cpu", input_size=14, batch_size=2) + frame = np.arange(7 * 9 * 3, dtype=np.uint8).reshape(7, 9, 3) + mask = np.zeros((7, 9), dtype=bool) + mask[1:6, 2:8] = True + mask[2, 3] = False + crop = frame[1:6, 2:8].copy() + crop[1, 1] = 0 + transform = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Resize((14, 14), antialias=True), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ] + ) + reference = model(transform(crop).unsqueeze(0))[0].numpy() + reference = reference / float(np.linalg.norm(reference)) + result = extractor.embed_batch([(frame, np.zeros_like(mask)), (frame, mask)]) + assert result[0] is None + assert result[1] is not None + single = extractor.embed(frame, mask) + assert single is not None + np.testing.assert_array_equal(result[1], reference) + np.testing.assert_array_equal(single, reference) + + +def test_memory_protocols_share_backbone_but_not_batch_state( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, torch_cpu: Any +) -> None: + from worldfoundry.base_models.perception_core.general_perception import dino_embeddings + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.context import EvaluationContext + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.features import get_dino_model + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.human.evaluator import ( + _dino_extractor, + ) + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.object.evaluator import ( + _feature_extractor, + ) + + checkpoint = tmp_path / "dinov2-large" + calls: list[Path] = [] + + def load(path: Path, *, device: str) -> Any: + assert device == "cpu" + calls.append(path) + return torch_cpu.nn.Identity() + + monkeypatch.setattr(dino_embeddings, "load_dinov2_feature_model", load) + context = EvaluationContext( + device="cpu", + cache_dir=tmp_path / "assets", + config={ + "metrics": { + "human_appearance": {"dino_checkpoint": str(checkpoint)}, + "object_appearance": {"dino_checkpoint": str(checkpoint)}, + } + }, + ) + try: + human = _dino_extractor(context) + obj = _feature_extractor(context) + assert human is not obj + assert human.model is obj.model is get_dino_model(context, checkpoint) + assert human.effective_batch_size > obj.effective_batch_size == 1 + assert len(calls) == 1 + assert get_dino_model(context, tmp_path / "other-checkpoint") is not human.model + assert len(calls) == 2 + finally: + context.close() + + +def test_wbench_binds_the_official_koniq_asset_without_global_cache_changes( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, torch_cpu: Any +) -> None: + from worldfoundry.base_models.capabilities import BASE_MODEL_CAPABILITIES + from worldfoundry.evaluation.tasks.execution.runners.wbench.runtime.wbench.src.metrics.video_quality import ( + imaging_quality, + ) + + checkpoint = tmp_path / "musiq_koniq_ckpt-e95806b9.pth" + asset_ids: list[str] = [] + calls: list[tuple[str, dict[str, Any]]] = [] + model = object() + + def create(name: str, **kwargs: Any) -> Any: + calls.append((name, kwargs)) + return model + + monkeypatch.setitem(sys.modules, "pyiqa", SimpleNamespace(create_metric=create)) + monkeypatch.setattr(imaging_quality, "wbench_asset_path", lambda name: asset_ids.append(name) or checkpoint) + metric = imaging_quality.ImagingQualityMetric(device="cpu") + assert metric.model is model + assert calls == [("musiq", {"device": "cpu", "pretrained_model_path": str(checkpoint)})] + assert asset_ids == ["wbench_musiq_koniq_checkpoint"] + asset = next( + asset for asset in BASE_MODEL_CAPABILITIES["wbench_quality_metric_assets"].assets if asset.id == asset_ids[0] + ) + assert str(asset.local_path).endswith(checkpoint.name) + + +@pytest.mark.parametrize("height,width", [(180, 320), (720, 1280)]) +def test_worldmembench_spaq_preserves_opencv_aspect_resize( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, torch_cpu: Any, height: int, width: int +) -> None: + cv2 = pytest.importorskip("cv2") + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.general_video_quality.musiq import ( + MusiqSpaqBackend, + ) + + captured: list[Any] = [] + + class Model(torch_cpu.nn.Module): + def forward(self, tensor: Any) -> Any: + captured.append(tensor.clone()) + return torch_cpu.full((len(tensor),), 50.0) + + checkpoint = tmp_path / "musiq_spaq_ckpt-358bb6af.pth" + checkpoint.touch() + model = Model() + + def load(*, pretrained_model_path: str) -> Any: + assert pretrained_model_path == str(checkpoint) + return model + + monkeypatch.setitem(sys.modules, "pyiqa.archs.musiq_arch", SimpleNamespace(MUSIQ=load)) + backend = MusiqSpaqBackend(checkpoint=checkpoint, device="cpu") + assert not backend.model.training + frame = np.arange(height * width * 3, dtype=np.uint8).reshape(height, width, 3) + expected = frame + if max(height, width) > 512: + scale = 512 / max(height, width) + expected = cv2.resize(frame, (int(width * scale), int(height * scale)), interpolation=cv2.INTER_LINEAR) + tensor = torch_cpu.from_numpy(expected.copy()).permute(2, 0, 1).float().unsqueeze(0) / 255 + assert backend.score([frame]) == [50.0] + torch_cpu.testing.assert_close(captured[0], tensor, rtol=0, atol=0) diff --git a/tests/evaluation/execution/test_worldmembench_camera.py b/tests/evaluation/execution/test_worldmembench_camera.py new file mode 100644 index 000000000..1f7a3127f --- /dev/null +++ b/tests/evaluation/execution/test_worldmembench_camera.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 WorldFoundry contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Weight-free camera replay with OpenCV and staged synthetic pose artifacts.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +import pytest + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench import run_worldmembench_official_runner as runner + +pytestmark = pytest.mark.unit + + +def test_native_camera_runner_scores_staged_poses_without_model_runtimes(tmp_path: Path) -> None: + cv2 = pytest.importorskip("cv2") + video = tmp_path / "video.mp4" + writer = cv2.VideoWriter(str(video), cv2.VideoWriter_fourcc(*"mp4v"), 10.0, (24, 16)) + assert writer.isOpened() + for _ in range(9): + writer.write(np.zeros((16, 24, 3), dtype=np.uint8)) + writer.release() + poses = np.repeat(np.eye(4)[None], 9, axis=0) + angles = np.linspace(0, 0.5, 9) + poses[:, 0, 0] = poses[:, 1, 1] = np.cos(angles) + poses[:, 1, 0] = np.sin(angles) + poses[:, 0, 1] = -np.sin(angles) + poses[:, :3, 3] = np.column_stack((angles, angles**2, np.sin(angles))) + np.savez(tmp_path / "target.npz", c2w=poses) + np.savez(tmp_path / "da3.npz", extrinsics=np.linalg.inv(poses), frame_indices=np.arange(9)) + np.savez(tmp_path / "vipe.npz", c2w=poses, frame_indices=np.arange(9)) + manifest = tmp_path / "manifest.jsonl" + manifest.write_text( + json.dumps( + { + "item_id": "case", + "model_id": "model_a", + "video": "video.mp4", + "annotations": {"camera_motion_type": "mixed"}, + "artifacts": {"target_poses": "target.npz", "da3": "da3.npz", "vipe": "vipe.npz"}, + } + ) + + "\n" + ) + args = runner.build_parser().parse_args( + [ + "--run-official", + "--input-path", + str(manifest), + "--device", + "cpu", + "--metric", + "camera_rotation,camera_translation", + "--python", + sys.executable, + "--weights-dir", + str(tmp_path / "assets"), + "--output-dir", + str(tmp_path / "out"), + ] + ) + scorecard, code = runner.run(args) + assert code == 0, (tmp_path / "out/runtime_stderr.log").read_text() + for name in ("camera_rotation", "camera_translation"): + metric = scorecard["metrics"]["per_metric"][name] + assert metric["raw_score"] == pytest.approx(0.0, abs=1e-6) + assert not metric["higher_is_better"] + assert len((tmp_path / "out/per_sample_scores.jsonl").read_text().splitlines()) == 2 diff --git a/worldfoundry/base_models/capabilities.py b/worldfoundry/base_models/capabilities.py index efaa25c1b..2c56b88e0 100644 --- a/worldfoundry/base_models/capabilities.py +++ b/worldfoundry/base_models/capabilities.py @@ -15,7 +15,9 @@ from worldfoundry.core.io.paths import worldfoundry_path_tokens -REPO_ROOT = Path(os.environ.get("WORLDFOUNDRY_REPO_ROOT") or worldfoundry_path_tokens()["WORLDFOUNDRY_REPO_ROOT"]).expanduser() +REPO_ROOT = Path( + os.environ.get("WORLDFOUNDRY_REPO_ROOT") or worldfoundry_path_tokens()["WORLDFOUNDRY_REPO_ROOT"] +).expanduser() def _path_tokens(env: Mapping[str, str] | None = None) -> dict[str, str]: @@ -44,8 +46,7 @@ def _path_tokens(env: Mapping[str, str] | None = None) -> dict[str, str]: ).expanduser() ) worldscore_asset_dir = Path( - environ.get("WORLDFOUNDRY_WORLDSCORE_ASSET_DIR") - or cache_root / "assets" / "worldscore_official" + environ.get("WORLDFOUNDRY_WORLDSCORE_ASSET_DIR") or cache_root / "assets" / "worldscore_official" ).expanduser() return { "WORLDFOUNDRY_REPO_ROOT": repo_root, @@ -60,8 +61,10 @@ def _path_tokens(env: Mapping[str, str] | None = None) -> dict[str, str]: "WORLDFOUNDRY_MODEL_DIR": core_tokens["WORLDFOUNDRY_MODEL_DIR"], "WORLDFOUNDRY_WORLDSCORE_ASSET_DIR": str(worldscore_asset_dir), "WORLDFOUNDRY_WORLDSCORE_DATASET_DIR": str( - Path(environ.get("WORLDFOUNDRY_WORLDSCORE_DATASET_DIR") or worldscore_asset_dir / "data" / "WorldScore-Dataset") - .expanduser() + Path( + environ.get("WORLDFOUNDRY_WORLDSCORE_DATASET_DIR") + or worldscore_asset_dir / "data" / "WorldScore-Dataset" + ).expanduser() ), "WORLDFOUNDRY_WORLDSCORE_ASSET_CHECKPOINT_DIR": str( Path( @@ -557,11 +560,7 @@ def pip_package_by_import(self) -> dict[str, str]: return {} if self.install_packages and len(self.install_packages) == len(self.package_imports): pairs = zip(self.package_imports, self.install_packages) - return { - import_name: package - for import_name, package in pairs - if import_name not in non_pip and package - } + return {import_name: package for import_name, package in pairs if import_name not in non_pip and package} pip_imports = [name for name in self.package_imports if name not in non_pip] install_packages = self.install_packages or tuple(pip_imports) return { @@ -928,7 +927,17 @@ def _perception_data_capability( family="depth", canonical_owner="worldfoundry.base_models.three_dimensions.depth.depth_anything.depth_anything_v3", canonical_path="worldfoundry/base_models/three_dimensions/depth/depth_anything/depth_anything_v3", - package_imports=("torch", "torchvision", "cv2", "numpy", "PIL", "omegaconf", "einops", "huggingface_hub", "safetensors"), + package_imports=( + "torch", + "torchvision", + "cv2", + "numpy", + "PIL", + "omegaconf", + "einops", + "huggingface_hub", + "safetensors", + ), install_packages=( "torch", "torchvision", @@ -1300,8 +1309,24 @@ def _perception_data_capability( family="depth", canonical_owner="worldfoundry.base_models.three_dimensions.depth.unidepth", canonical_path="worldfoundry/base_models/three_dimensions/depth/unidepth", - package_imports=("torch", "torchvision", "numpy", "einops", "huggingface_hub", "safetensors", "xformers.components.attention"), - install_packages=("torch", "torchvision", "numpy", "einops", "huggingface-hub", "safetensors", "xformers==0.0.28.post3"), + package_imports=( + "torch", + "torchvision", + "numpy", + "einops", + "huggingface_hub", + "safetensors", + "xformers.components.attention", + ), + install_packages=( + "torch", + "torchvision", + "numpy", + "einops", + "huggingface-hub", + "safetensors", + "xformers==0.0.28.post3", + ), asset_env=("WORLDFOUNDRY_UNIDEPTH_V2_VITL14_MODEL_DIR",), assets=( BaseModelAsset( @@ -1548,9 +1573,7 @@ def _perception_data_capability( role="checkpoint", env=("WORLDFOUNDRY_DEAOT_R50_CKPT",), local_path="${WORLDFOUNDRY_CKPT_DIR}/R50_DeAOTL_PRE_YTB_DAV.pth", - alternate_paths=( - "${WORLDFOUNDRY_WORKSPACE_ROOT}/ckpt/R50_DeAOTL_PRE_YTB_DAV.pth", - ), + alternate_paths=("${WORLDFOUNDRY_WORKSPACE_ROOT}/ckpt/R50_DeAOTL_PRE_YTB_DAV.pth",), min_size_bytes=150_000_000, note="DeAOT-L R50 checkpoint used by Segment-and-Track-Anything object attribute metrics.", ), @@ -1618,7 +1641,16 @@ def _perception_data_capability( canonical_owner="worldfoundry.base_models.perception_core.frame_interpolation.vfimamba", canonical_path="worldfoundry/base_models/perception_core/frame_interpolation/vfimamba", package_imports=("torch", "cv2", "numpy", "timm", "einops", "mamba_ssm", "huggingface_hub", "transformers"), - install_packages=("torch", "opencv-python", "numpy", "timm", "einops", "mamba-ssm", "huggingface-hub", "transformers"), + install_packages=( + "torch", + "opencv-python", + "numpy", + "timm", + "einops", + "mamba-ssm", + "huggingface-hub", + "transformers", + ), asset_env=("WORLDFOUNDRY_VFIMAMBA_CKPT", "WORLDFOUNDRY_VFIMAMBA_S_CKPT"), assets=( BaseModelAsset( @@ -1742,17 +1774,17 @@ def _perception_data_capability( download_command_template=( "RAFT_CKPT=${WORLDFOUNDRY_RAFT_THINGS_CKPT:-${WORLDFOUNDRY_WORLDSCORE_ASSET_CHECKPOINT_DIR}/models/raft-things.pth}; " "RAFT_ZIP=${WORLDFOUNDRY_RAFT_MODELS_ZIP:-${WORLDFOUNDRY_WORLDSCORE_ASSET_CHECKPOINT_DIR}/models.zip}; " - "mkdir -p \"$(dirname \"$RAFT_CKPT\")\" \"$(dirname \"$RAFT_ZIP\")\"; " - "if [ ! -f \"$RAFT_CKPT\" ]; then " - "curl -L --fail --retry 5 --retry-delay 2 https://dl.dropboxusercontent.com/s/4j4z58wuv8o0mfz/models.zip -o \"$RAFT_ZIP\" " + 'mkdir -p "$(dirname "$RAFT_CKPT")" "$(dirname "$RAFT_ZIP")"; ' + 'if [ ! -f "$RAFT_CKPT" ]; then ' + 'curl -L --fail --retry 5 --retry-delay 2 https://dl.dropboxusercontent.com/s/4j4z58wuv8o0mfz/models.zip -o "$RAFT_ZIP" ' "&& python -c 'import pathlib, sys, zipfile; " - "member=\"models/raft-things.pth\"; " + 'member="models/raft-things.pth"; ' "archive=zipfile.ZipFile(sys.argv[1]); " "info=archive.getinfo(member); " "target=pathlib.Path(sys.argv[2]); " "target.parent.mkdir(parents=True, exist_ok=True); " "target.write_bytes(archive.read(info)); " - "print(target)' \"$RAFT_ZIP\" \"$RAFT_CKPT\"; " + 'print(target)\' "$RAFT_ZIP" "$RAFT_CKPT"; ' "fi" ), min_size_bytes=20_000_000, @@ -1803,7 +1835,7 @@ def _perception_data_capability( min_file_count=10, download_command_template=( "COCO_ROOT=${WORLDFOUNDRY_COCO2017_ROOT:-${WORLDFOUNDRY_DATA_DIR}/perception/coco2017}; " - "mkdir -p \"$COCO_ROOT\" && cd \"$COCO_ROOT\"; " + 'mkdir -p "$COCO_ROOT" && cd "$COCO_ROOT"; ' "if [ ! -d val2017 ]; then " "curl -L --fail --retry 5 --retry-delay 2 http://images.cocodataset.org/zips/val2017.zip -o val2017.zip " "&& python -m zipfile -e val2017.zip .; " @@ -1828,12 +1860,12 @@ def _perception_data_capability( min_file_count=20, download_command_template=( "DAVIS_ROOT=${WORLDFOUNDRY_DAVIS2017_ROOT:-${WORLDFOUNDRY_DATA_DIR}/perception/davis2017/DAVIS}; " - "DAVIS_PARENT=$(dirname \"$DAVIS_ROOT\"); " - "mkdir -p \"$DAVIS_PARENT\"; " - "if [ ! -d \"$DAVIS_ROOT/JPEGImages/480p\" ]; then " + 'DAVIS_PARENT=$(dirname "$DAVIS_ROOT"); ' + 'mkdir -p "$DAVIS_PARENT"; ' + 'if [ ! -d "$DAVIS_ROOT/JPEGImages/480p" ]; then ' "curl -L --fail --retry 5 --retry-delay 2 https://data.vision.ee.ethz.ch/csergi/share/davis/DAVIS-2017-trainval-480p.zip " - "-o \"$DAVIS_PARENT/DAVIS-2017-trainval-480p.zip\" " - "&& python -m zipfile -e \"$DAVIS_PARENT/DAVIS-2017-trainval-480p.zip\" \"$DAVIS_PARENT\"; " + '-o "$DAVIS_PARENT/DAVIS-2017-trainval-480p.zip" ' + '&& python -m zipfile -e "$DAVIS_PARENT/DAVIS-2017-trainval-480p.zip" "$DAVIS_PARENT"; ' "fi" ), note="DAVIS 2017 trainval 480p bundle for video object segmentation and tracking metric validation.", @@ -1851,12 +1883,12 @@ def _perception_data_capability( min_file_count=10, download_command_template=( "TUM_ROOT=${WORLDFOUNDRY_TUM_RGBD_SLAM_SAMPLE_ROOT:-${WORLDFOUNDRY_DATA_DIR}/perception/tum_rgbd/rgbd_dataset_freiburg1_desk}; " - "TUM_PARENT=$(dirname \"$TUM_ROOT\"); " - "mkdir -p \"$TUM_PARENT\"; " - "if [ ! -f \"$TUM_ROOT/rgb.txt\" ]; then " + 'TUM_PARENT=$(dirname "$TUM_ROOT"); ' + 'mkdir -p "$TUM_PARENT"; ' + 'if [ ! -f "$TUM_ROOT/rgb.txt" ]; then ' "curl -L --fail --retry 5 --retry-delay 2 https://vision.in.tum.de/rgbd/dataset/freiburg1/rgbd_dataset_freiburg1_desk.tgz " - "-o \"$TUM_PARENT/rgbd_dataset_freiburg1_desk.tgz\" " - "&& tar -xzf \"$TUM_PARENT/rgbd_dataset_freiburg1_desk.tgz\" -C \"$TUM_PARENT\"; " + '-o "$TUM_PARENT/rgbd_dataset_freiburg1_desk.tgz" ' + '&& tar -xzf "$TUM_PARENT/rgbd_dataset_freiburg1_desk.tgz" -C "$TUM_PARENT"; ' "fi" ), note="Small official TUM RGB-D Freiburg1 desk sequence for SLAM/depth validation and trajectory sanity checks.", @@ -1905,9 +1937,7 @@ def _perception_data_capability( role="metric_checkpoint_bundle", env=("WORLDFOUNDRY_WORLDSCORE_ASSET_CHECKPOINT_DIR",), local_path="${WORLDFOUNDRY_WORLDSCORE_ASSET_CHECKPOINT_DIR}", - alternate_paths=( - "${WORLDFOUNDRY_WORLDSCORE_ASSET_DIR}/metrics/checkpoints", - ), + alternate_paths=("${WORLDFOUNDRY_WORLDSCORE_ASSET_DIR}/metrics/checkpoints",), min_file_count=8, note="WorldScore metric checkpoint bundle containing DROID, SAM, SAM2, RAFT, SEA-RAFT, and quality checkpoints.", ), @@ -1981,9 +2011,7 @@ def _perception_data_capability( role="benchmark_dataset", env=("WORLDFOUNDRY_VIDEOSCORE_BENCH_DATASET_ROOT",), local_path="${WORLDFOUNDRY_HFD_DATASET_ROOT}/TIGER-Lab__VideoScore-Bench", - alternate_paths=( - "${WORLDFOUNDRY_WORKSPACE_ROOT}/data/hfd_datasets/TIGER-Lab__VideoScore-Bench", - ), + alternate_paths=("${WORLDFOUNDRY_WORKSPACE_ROOT}/data/hfd_datasets/TIGER-Lab__VideoScore-Bench",), hf_repo_id="TIGER-Lab/VideoScore-Bench", hf_repo_type="dataset", min_file_count=4, @@ -2269,20 +2297,19 @@ def _perception_data_capability( note="LAION aesthetic predictor linear head used by WBench aesthetic quality.", ), BaseModelAsset( - id="wbench_musiq_spaq_checkpoint", + id="wbench_musiq_koniq_checkpoint", kind="file", role="checkpoint", - env=("WORLDFOUNDRY_WBENCH_MUSIQ_SPAQ_CKPT",), - local_path="${WORLDFOUNDRY_CACHE_DIR}/models/wbench/pyiqa_model/musiq_spaq_ckpt-358bb6af.pth", + env=("WORLDFOUNDRY_WBENCH_MUSIQ_KONIQ_CKPT",), + local_path="${WORLDFOUNDRY_CACHE_DIR}/models/wbench/pyiqa_model/musiq_koniq_ckpt-e95806b9.pth", alternate_paths=( - "${WORLDFOUNDRY_CKPT_DIR}/WBench/pyiqa_model/musiq_spaq_ckpt-358bb6af.pth", - "${WORLDFOUNDRY_CKPT_DIR}/WBench/weights/pyiqa/musiq_spaq_ckpt-358bb6af.pth", - "${WORLDFOUNDRY_CKPT_DIR}/WBench/weights/pyiqa_model/musiq_spaq_ckpt-358bb6af.pth", - "${WORLDFOUNDRY_CKPT_DIR}/VBench/pyiqa_model/musiq_spaq_ckpt-358bb6af.pth", + "${WORLDFOUNDRY_CKPT_DIR}/WBench/pyiqa_model/musiq_koniq_ckpt-e95806b9.pth", + "${WORLDFOUNDRY_CKPT_DIR}/WBench/weights/pyiqa/musiq_koniq_ckpt-e95806b9.pth", + "${WORLDFOUNDRY_CKPT_DIR}/WBench/weights/pyiqa_model/musiq_koniq_ckpt-e95806b9.pth", ), - download_url="https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/musiq_spaq_ckpt-358bb6af.pth", + download_url="https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/musiq_koniq_ckpt-e95806b9.pth", min_size_bytes=100_000_000, - note="MUSIQ-SPAQ checkpoint used by WBench imaging quality.", + note="MUSIQ-KonIQ checkpoint selected by the official WBench pyiqa.create_metric('musiq') call.", ), BaseModelAsset( id="wbench_clip_vit_b16_model_dir", @@ -2605,9 +2632,7 @@ def _perception_data_capability( min_file_count=8, required_files=("constants.py", "models/model.py", "models/vqascore_models/__init__.py"), required_paths=("models",), - note=( - "In-tree VQAScore model backend shared by evaluation scorers." - ), + note=("In-tree VQAScore model backend shared by evaluation scorers."), ), ), gpu_required=False, @@ -2642,9 +2667,7 @@ def _perception_data_capability( role="metric_tokenizer_dir", env=("WORLDFOUNDRY_T2V_METRICS_FLAN_T5_XXL_DIR",), local_path="${WORLDFOUNDRY_HFD_ROOT}/google--flan-t5-xxl", - alternate_paths=( - "${WORLDFOUNDRY_WORKSPACE_ROOT}/ckpt/hfd/google--flan-t5-xxl", - ), + alternate_paths=("${WORLDFOUNDRY_WORKSPACE_ROOT}/ckpt/hfd/google--flan-t5-xxl",), hf_repo_id="google/flan-t5-xxl", min_file_count=2, required_files=("config.json", "tokenizer_config.json"), @@ -2667,9 +2690,7 @@ def _perception_data_capability( role="metric_model_dir", env=("WORLDFOUNDRY_T2V_METRICS_BLIP2_ITM_DIR",), local_path="${WORLDFOUNDRY_HFD_ROOT}/Salesforce--blip2-opt-2.7b", - alternate_paths=( - "${WORLDFOUNDRY_WORKSPACE_ROOT}/ckpt/hfd/Salesforce--blip2-opt-2.7b", - ), + alternate_paths=("${WORLDFOUNDRY_WORKSPACE_ROOT}/ckpt/hfd/Salesforce--blip2-opt-2.7b",), hf_repo_id="Salesforce/blip2-opt-2.7b", min_file_count=2, required_files=("config.json",), @@ -2749,9 +2770,7 @@ def _perception_data_capability( role="metric_model_file", env=("WORLDFOUNDRY_JEDI_VITH16_CHECKPOINT",), local_path="${WORLDFOUNDRY_JEDI_MODEL_DIR}/vith16.pth.tar", - alternate_paths=( - "${WORLDFOUNDRY_HFD_ROOT}/jedi-vjepa/vith16.pth.tar", - ), + alternate_paths=("${WORLDFOUNDRY_HFD_ROOT}/jedi-vjepa/vith16.pth.tar",), download_url="https://dl.fbaipublicfiles.com/jepa/vith16/vith16.pth.tar", required_files=("vith16.pth.tar",), note="V-JEPA ViT-H/16 encoder checkpoint used by VideoJEDi feature extraction.", @@ -2773,9 +2792,7 @@ def _perception_data_capability( role="metric_model_file", env=("WORLDFOUNDRY_JEDI_SSV2_PROBE_CHECKPOINT",), local_path="${WORLDFOUNDRY_JEDI_MODEL_DIR}/ssv2-probe.pth.tar", - alternate_paths=( - "${WORLDFOUNDRY_HFD_ROOT}/jedi-vjepa/ssv2-probe.pth.tar", - ), + alternate_paths=("${WORLDFOUNDRY_HFD_ROOT}/jedi-vjepa/ssv2-probe.pth.tar",), download_url="https://dl.fbaipublicfiles.com/jepa/vith16/ssv2-probe.pth.tar", required_files=("ssv2-probe.pth.tar",), note="Fine-tuned SSV2 attentive probe used by VideoJEDi feature aggregation.", @@ -3282,7 +3299,9 @@ def _perception_data_capability( role="checkpoint", env=("WORLDFOUNDRY_VBENCH_UMT_CKPT",), local_path="${WORLDFOUNDRY_CACHE_DIR}/models/vbench/umt_model/l16_ptk710_ftk710_ftk400_f16_res224.pth", - alternate_paths=("${WORLDFOUNDRY_CKPT_DIR}/VBench/umt_model/l16_ptk710_ftk710_ftk400_f16_res224.pth",), + alternate_paths=( + "${WORLDFOUNDRY_CKPT_DIR}/VBench/umt_model/l16_ptk710_ftk710_ftk400_f16_res224.pth", + ), hf_repo_id="OpenGVLab/VBench_Used_Models", hf_filename="l16_ptk710_ftk710_ftk400_f16_res224.pth", min_size_bytes=200_000_000, @@ -3382,7 +3401,9 @@ def _perception_data_capability( role="checkpoint", env=("WORLDFOUNDRY_VBENCH_AESTHETIC_LINEAR_CKPT",), local_path="${WORLDFOUNDRY_CACHE_DIR}/models/vbench/aesthetic_model/emb_reader/sa_0_4_vit_l_14_linear.pth", - alternate_paths=("${WORLDFOUNDRY_CKPT_DIR}/VBench/aesthetic_model/emb_reader/sa_0_4_vit_l_14_linear.pth",), + alternate_paths=( + "${WORLDFOUNDRY_CKPT_DIR}/VBench/aesthetic_model/emb_reader/sa_0_4_vit_l_14_linear.pth", + ), download_url="https://github.com/LAION-AI/aesthetic-predictor/raw/main/sa_0_4_vit_l_14_linear.pth", min_size_bytes=1_000, note="LAION aesthetic predictor linear head used by VBench aesthetic quality.", @@ -3404,7 +3425,9 @@ def _perception_data_capability( role="checkpoint", env=("WORLDFOUNDRY_VBENCH_RETINAFACE_CKPT",), local_path="${WORLDFOUNDRY_CACHE_DIR}/models/vbench/retina_face_model/retinaface_resnet50_2020-07-20-f168fae3c.zip", - alternate_paths=("${WORLDFOUNDRY_CKPT_DIR}/VBench/retina_face_model/retinaface_resnet50_2020-07-20-f168fae3c.zip",), + alternate_paths=( + "${WORLDFOUNDRY_CKPT_DIR}/VBench/retina_face_model/retinaface_resnet50_2020-07-20-f168fae3c.zip", + ), download_url="https://github.com/ternaus/retinaface/releases/download/0.01/retinaface_resnet50_2020-07-20-f168fae3c.zip", min_size_bytes=100_000_000, required=False, @@ -3490,7 +3513,9 @@ def _perception_data_capability( role="checkpoint", env=("WORLDFOUNDRY_VBENCH2_YOLO_WORLD_CKPT",), local_path="${WORLDFOUNDRY_CACHE_DIR}/models/vbench2/YOLO-World/yolo_world_v2_xl_obj365v1_goldg_cc3mlite_pretrain-5daf1395.pth", - alternate_paths=("${WORLDFOUNDRY_CKPT_DIR}/VBench2/YOLO-World/yolo_world_v2_xl_obj365v1_goldg_cc3mlite_pretrain-5daf1395.pth",), + alternate_paths=( + "${WORLDFOUNDRY_CKPT_DIR}/VBench2/YOLO-World/yolo_world_v2_xl_obj365v1_goldg_cc3mlite_pretrain-5daf1395.pth", + ), required=False, manual_action=( "Stage the official VBench-2.0 YOLO-World checkpoint under $WORLDFOUNDRY_VBENCH2_CACHE_DIR " @@ -4254,10 +4279,7 @@ def base_model_inventory() -> dict[str, Any]: "stack_count": len(BASE_MODEL_STACKS), "asset_count": asset_count, "required_asset_count": sum( - 1 - for capability in BASE_MODEL_CAPABILITIES.values() - for asset in capability.assets - if asset.required + 1 for capability in BASE_MODEL_CAPABILITIES.values() for asset in capability.assets if asset.required ), "capabilities_by_family": _capabilities_by_family(), "stacks_by_family": _stacks_by_family(), @@ -4268,13 +4290,9 @@ def base_model_inventory() -> dict[str, Any]: for benchmark_id, capability_ids in sorted(BENCHMARK_DATA_ASSET_CAPABILITIES.items()) }, "capabilities": [ - BASE_MODEL_CAPABILITIES[capability_id].to_dict() - for capability_id in sorted(BASE_MODEL_CAPABILITIES) - ], - "stacks": [ - BASE_MODEL_STACKS[stack_id].to_dict() - for stack_id in sorted(BASE_MODEL_STACKS) + BASE_MODEL_CAPABILITIES[capability_id].to_dict() for capability_id in sorted(BASE_MODEL_CAPABILITIES) ], + "stacks": [BASE_MODEL_STACKS[stack_id].to_dict() for stack_id in sorted(BASE_MODEL_STACKS)], } @@ -4607,33 +4625,15 @@ def base_model_materialization_plan( check_base_model_capability(capability_id, env=env, python_executable=python_executable) for capability_id in selected_ids ] - download_commands = [ - command - for check in checks - for command in check["repair_hints"]["download_commands"] - ] + download_commands = [command for check in checks for command in check["repair_hints"]["download_commands"]] download_command_argvs = [ - command - for check in checks - for command in check["repair_hints"]["download_command_argvs"] - ] - export_commands = [ - command - for check in checks - for command in check["repair_hints"]["export_commands"] + command for check in checks for command in check["repair_hints"]["download_command_argvs"] ] + export_commands = [command for check in checks for command in check["repair_hints"]["export_commands"]] pip_install_packages = sorted( - { - package - for check in checks - for package in check["repair_hints"]["pip_install_packages"] - } + {package for check in checks for package in check["repair_hints"]["pip_install_packages"]} ) - manual_actions = [ - action - for check in checks - for action in check["repair_hints"]["manual_actions"] - ] + manual_actions = [action for check in checks for action in check["repair_hints"]["manual_actions"]] return { "schema_version": "worldfoundry-base-model-assets-v1", "ok": all(check["ok"] for check in checks), diff --git a/worldfoundry/base_models/perception_core/general_perception/dino_embeddings.py b/worldfoundry/base_models/perception_core/general_perception/dino_embeddings.py index 639c73122..13256eaaf 100644 --- a/worldfoundry/base_models/perception_core/general_perception/dino_embeddings.py +++ b/worldfoundry/base_models/perception_core/general_perception/dino_embeddings.py @@ -61,6 +61,25 @@ def load_dino_vitb16_feature_model( return model.eval() +def load_dinov2_feature_model( + checkpoint: str | Path, + *, + device: str | torch.device | None = None, + use_pooler_output: bool = False, + local_files_only: bool = True, +) -> HFVisionEmbeddingModel: + """Load a selected DINOv2 checkpoint with explicit CLS/pooler semantics.""" + from transformers import AutoModel + + model = HFVisionEmbeddingModel( + AutoModel.from_pretrained(str(checkpoint), local_files_only=local_files_only), + use_pooler_output=use_pooler_output, + ) + if device is not None: + model = model.to(device) + return model.eval() + + def load_dinov2_base_feature_model( *, env: Mapping[str, str] | None = None, diff --git a/worldfoundry/data/benchmarks/catalog/video/worldmembench.yaml b/worldfoundry/data/benchmarks/catalog/video/worldmembench.yaml index a086a73d7..b1238df72 100644 --- a/worldfoundry/data/benchmarks/catalog/video/worldmembench.yaml +++ b/worldfoundry/data/benchmarks/catalog/video/worldmembench.yaml @@ -24,7 +24,7 @@ integration_evidence: false leaderboard_valid: false checkpoint_refs: [] requires: -- Official result summaries, or generated videos with a provisioned upstream runtime and selected metric assets. +- Official result summaries, or generated videos with staged metric checkpoints and annotations. blockers: - Complete-dataset and checkpoint replay remain pending. source: @@ -42,8 +42,8 @@ source_provenance: url: https://github.com/ZJU-REAL/WorldMemBench revision: 1b79840143a98e017891e893b5d3ef75b93078a3 license: MIT - role: external_runtime_dependency - provenance: Official metric implementations are invoked without vendoring or reimplementation. + role: adapted_official_protocol + provenance: Official metric protocols run in tree with shared CLIP, LAION and DINO model components. metrics: - id: imaging_quality name: Imaging Quality @@ -121,7 +121,7 @@ integration: status: integrated tier: in_tree evidence: - runtime_surface: external_official_runtime_with_result_normalizer + runtime_surface: in_tree_official_runtime_with_result_normalizer full_suite_verified: false runner: runner_target: worldfoundry.evaluation.tasks.contracts.external:WorldMemBenchContract @@ -137,11 +137,10 @@ runner: artifact_official_run: true pass_generated_artifact_dir: true runtime: - kind: external_official_repo + kind: in_tree_official_runtime + in_tree_path: worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime repo_url: https://github.com/ZJU-REAL/WorldMemBench repo_revision: 1b79840143a98e017891e893b5d3ef75b93078a3 - root_env: WORLDFOUNDRY_WORLDMEMBENCH_ROOT - root_env_required: false results_path_env: WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH expected_artifacts: - scorecard.json @@ -150,7 +149,6 @@ runner: verification_status: pending required_env: [] optional_env: - - WORLDFOUNDRY_WORLDMEMBENCH_ROOT - WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH - WORLDFOUNDRY_WORLDMEMBENCH_PYTHON - WORLDFOUNDRY_WORLDMEMBENCH_CONFIG @@ -176,7 +174,7 @@ runner_availability: task_yaml_available: true runtime_verified: false full_suite_verified: false - surface: external_official_runtime_with_result_normalizer + surface: in_tree_official_runtime_with_result_normalizer validation_command: - worldfoundry-eval - zoo @@ -192,5 +190,7 @@ validation_command: - --json notes: - Source and metric protocol are attributed to ZJU-REAL/WorldMemBench under MIT. -- Result import is weight-free; official metric execution requires a separately provisioned upstream runtime and assets. +- All twelve metric protocols execute in tree; result import remains weight-free. +- CLIP, LAION heads and DINO checkpoint loading reuse canonical components; human/object masked DINO features share one implementation. +- DA3-Streaming and ViPE pose preparation accept staged artifacts or explicitly configured model runtimes; the WorldMemBench checkout is not required. - GPU and full-checkpoint replay are pending; result import does not establish leaderboard validity. diff --git a/worldfoundry/data/benchmarks/requirements/worldmembench.txt b/worldfoundry/data/benchmarks/requirements/worldmembench.txt new file mode 100644 index 000000000..562d7fb7d --- /dev/null +++ b/worldfoundry/data/benchmarks/requirements/worldmembench.txt @@ -0,0 +1,13 @@ +# Install into the prepared GPU environment; retain its Torch/CUDA ABI. +# Quality: pyiqa. Human: ultralytics, insightface, onnxruntime-gpu. +# Object/environment: transformers. Dynamic Memory uses a configured VLM endpoint. +numpy>=1.24 +opencv-python-headless>=4.8,<5 +Pillow>=10 +PyYAML>=6 +pyiqa +transformers>=4.57 +ultralytics>=8.3 +insightface>=0.7 +onnxruntime-gpu>=1.24 +scipy>=1.10 diff --git a/worldfoundry/data/benchmarks/runtime_profiles/official/worldmembench.yaml b/worldfoundry/data/benchmarks/runtime_profiles/official/worldmembench.yaml index ce275af0f..d93e22bbc 100644 --- a/worldfoundry/data/benchmarks/runtime_profiles/official/worldmembench.yaml +++ b/worldfoundry/data/benchmarks/runtime_profiles/official/worldmembench.yaml @@ -9,7 +9,6 @@ pythonpath_roots: - . required_env: [] optional_env: -- WORLDFOUNDRY_WORLDMEMBENCH_ROOT - WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH - WORLDFOUNDRY_WORLDMEMBENCH_PYTHON - WORLDFOUNDRY_WORLDMEMBENCH_CONFIG @@ -20,6 +19,10 @@ required_packages: [] required_paths: - id: official_runner_script path: worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py +- id: in_tree_runtime + path: worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/cli.py +- id: metric_requirements + path: worldfoundry/data/benchmarks/requirements/worldmembench.txt - id: metrics_module path: worldfoundry/evaluation/tasks/execution/runners/worldmembench/worldmembench_metrics.py - id: task_yaml @@ -32,10 +35,11 @@ requires_cuda_visibility: false status: official_runtime_ready_real_data_validation_pending notes: - Result import uses the unified environment without CUDA, metric checkpoints, or upstream Python dependencies. -- Official execution uses an installed worldmembench package or WORLDFOUNDRY_WORLDMEMBENCH_ROOT/src. -- WORLDFOUNDRY_WORLDMEMBENCH_PYTHON selects a provisioned upstream interpreter without changing the unified Torch/CUDA ABI. -- DA3 and ViPE subprocess interpreters and metric checkpoint paths belong in the upstream YAML configuration. -- Dynamic Memory uses the upstream VLM configuration and API-key environment variable; credentials are not runner arguments. +- All twelve official metric protocols execute in tree without an installed worldmembench package or checkout. +- CLIP, LAION heads and DINO checkpoint loading use shared components; masked human/object DINO inference has one implementation. +- WORLDFOUNDRY_WORLDMEMBENCH_PYTHON selects the metric interpreter without changing the unified Torch/CUDA ABI. +- DA3-Streaming and ViPE preparation reuse staged artifacts or explicitly configured model runtimes and interpreters. +- Dynamic Memory uses the official VLM configuration and API-key environment variable; credentials are not runner arguments. - Official dataset/checkpoint replay and leaderboard parity remain pending. validation_command: &worldmembench_validation_command - worldfoundry-eval @@ -56,8 +60,6 @@ full_runtime_command: - -m - worldfoundry.evaluation.tasks.execution.runners.worldmembench.run_worldmembench_official_runner - --run-official -- --repo-root -- - --input-path - - --config diff --git a/worldfoundry/data/benchmarks/runtime_profiles/official/worldscore.yaml b/worldfoundry/data/benchmarks/runtime_profiles/official/worldscore.yaml index 8f125c952..7bfd58695 100644 --- a/worldfoundry/data/benchmarks/runtime_profiles/official/worldscore.yaml +++ b/worldfoundry/data/benchmarks/runtime_profiles/official/worldscore.yaml @@ -15,23 +15,23 @@ bounded_official_gpu_validation: verified: true scope: bounded_official_gpu_validation scorecard: tmp/local_open_eval/worldscore_bounded_gpu_validation_after_readiness_layers_20260525/scorecard.json -blocked_reason: Bounded WorldScore dynamic official GPU validation passed; - DROID/GroundingDINO/SAM/RAFT/SEA-RAFT/FlowFormer++/VFIMamba are tracked as reusable WorldFoundry base-model +blocked_reason: Bounded WorldScore dynamic official GPU validation passed; + DROID/GroundingDINO/SAM/RAFT/SEA-RAFT/FlowFormer++/VFIMamba are tracked as reusable WorldFoundry base-model dependencies. Full leaderboard parity still needs complete evaluation.json trees and full GPU metric execution. environment_id: worldfoundry-unified-cu128 id: worldscore name: WorldScore needs_new_env: false notes: -- DROID/GroundingDINO/SAM/RAFT/SEA-RAFT/FlowFormer++/VFIMamba checkpoints are now modeled as reusable WorldFoundry +- DROID/GroundingDINO/SAM/RAFT/SEA-RAFT/FlowFormer++/VFIMamba checkpoints are now modeled as reusable WorldFoundry base-model dependencies rather than benchmark-private assets. -- Full leaderboard parity still needs complete evaluation.json trees and full GPU metric execution. +- Full leaderboard parity still needs complete evaluation.json trees and full GPU metric execution. CUDA_VISIBLE_DEVICES=7; both cu118 stacks initialized CUDA and completed a one-element tensor sum. -preflight_command: +preflight_command: PYTHONPATH=${WORLDFOUNDRY_REPO_ROOT}:${WORLDFOUNDRY_REPO_ROOT}/worldfoundry/evaluation/tasks/execution/runners/worldscore/runtime/worldscore ${WORLDFOUNDRY_UNIFIED_PYTHON:-python} -m - worldfoundry.evaluation.tasks.execution.orchestration.runtime_preflight --profile worldscore --manifest - worldfoundry/data/benchmarks/runtime_profiles/official/worldscore.yaml --output-dir + worldfoundry.evaluation.tasks.execution.orchestration.runtime_preflight --profile worldscore --manifest + worldfoundry/data/benchmarks/runtime_profiles/official/worldscore.yaml --output-dir tmp/worldfoundry_preflight/worldscore --json python_path: ${WORLDFOUNDRY_UNIFIED_PYTHON:-python} pythonpath_roots: @@ -57,7 +57,6 @@ required_imports: - structlog - submitit - pyiqa -- pytorch_lightning - worldfoundry.base_models.perception_core.detection.grounding_dino - worldfoundry.base_models.perception_core.segment.sam2 - worldfoundry.base_models.perception_core.segment.sam_v1 @@ -95,7 +94,6 @@ required_packages: - structlog - submitit - pyiqa -- pytorch-lightning - lietorch - spacy - cvxpy @@ -142,20 +140,20 @@ base_model_dependency_preflight: - vfimamba ready_env: WORLDFOUNDRY_DROID_SLAM_CKPT: cache/worldfoundry/assets/worldscore_official/metrics/checkpoints/droid.pth - WORLDFOUNDRY_FLOWFORMERPLUSPLUS_CKPT: + WORLDFOUNDRY_FLOWFORMERPLUSPLUS_CKPT: cache/worldfoundry/assets/worldscore_official/metrics/checkpoints/things_288960.pth WORLDFOUNDRY_GROUNDING_DINO_CKPT: cache/hfd/ShilongLiu--GroundingDINO/groundingdino_swint_ogc.pth - WORLDFOUNDRY_RAFT_THINGS_CKPT: + WORLDFOUNDRY_RAFT_THINGS_CKPT: cache/worldfoundry/assets/worldscore_official/metrics/checkpoints/models/raft-things.pth WORLDFOUNDRY_SAM_VIT_H_CKPT: cache/worldfoundry/assets/worldscore_official/metrics/checkpoints/sam_vit_h_4b8939.pth WORLDFOUNDRY_SAM2_CKPT: cache/hfd/facebook--sam2.1-hiera-base-plus/sam2.1_hiera_base_plus.pt - WORLDFOUNDRY_SEA_RAFT_CKPT: + WORLDFOUNDRY_SEA_RAFT_CKPT: cache/worldfoundry/assets/worldscore_official/metrics/checkpoints/Tartan-C-T-TSKH-spring540x960-M.pth WORLDFOUNDRY_VFIMAMBA_CKPT: ckpt/WorldScore/VFIMamba.pkl runtime_env: - WORLDFOUNDRY_WORLDSCORE_ROOT: + WORLDFOUNDRY_WORLDSCORE_ROOT: ${WORLDFOUNDRY_REPO_ROOT}/worldfoundry/evaluation/tasks/execution/runners/worldscore/runtime/worldscore - WORLDFOUNDRY_WORLDSCORE_CONFIG_ROOT: + WORLDFOUNDRY_WORLDSCORE_CONFIG_ROOT: ${WORLDFOUNDRY_REPO_ROOT}/worldfoundry/data/benchmarks/assets/worldscore/config WORLDFOUNDRY_WORLDSCORE_ASSET_CHECKPOINT_DIR: cache/worldfoundry/assets/worldscore_official/metrics/checkpoints missing: [] diff --git a/worldfoundry/evaluation/tasks/contracts/external.py b/worldfoundry/evaluation/tasks/contracts/external.py index ad5602907..41189a843 100644 --- a/worldfoundry/evaluation/tasks/contracts/external.py +++ b/worldfoundry/evaluation/tasks/contracts/external.py @@ -876,9 +876,9 @@ def build_robotwin_contract_payload(benchmark_id: str = ROBOTWIN_BENCHMARK_ID) - ), requires_upstream_runtime=True, notes=( - "Imports official per-model summaries or invokes the installed ZJU-REAL/WorldMemBench runtime.", + "Imports official per-model summaries or runs all twelve metric protocols in tree with shared model components.", "Scores retain official units, case statuses, and camera-error directions; no composite score is defined.", - "Full evaluation needs metric checkpoints, DA3/ViPE environments, and a Dynamic Memory VLM endpoint.", + "Full evaluation needs metric checkpoints, staged or configured DA3-Streaming/ViPE artifacts, and a Dynamic Memory VLM endpoint.", ), ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/four_d_worldbench/runtime/four_d_worldbench/metric/base_metrics.py b/worldfoundry/evaluation/tasks/execution/runners/four_d_worldbench/runtime/four_d_worldbench/metric/base_metrics.py index 2a20e2411..d98f77736 100644 --- a/worldfoundry/evaluation/tasks/execution/runners/four_d_worldbench/runtime/four_d_worldbench/metric/base_metrics.py +++ b/worldfoundry/evaluation/tasks/execution/runners/four_d_worldbench/runtime/four_d_worldbench/metric/base_metrics.py @@ -1,114 +1,30 @@ -from abc import ABC, abstractmethod -from typing import List, Union +"""4DWorldBench image metrics backed by the shared IQA implementation.""" -import pyiqa -import pyiqa.models -import pyiqa.models.inference_model -import torch -from PIL import Image -from pyiqa.models.inference_model import InferenceModel -from torchvision import transforms -from torchvision.transforms import ToTensor -from typing import Tuple -import numpy as np +from typing import Any + +from worldfoundry.evaluation.tasks.metrics._shared.pyiqa import ImageQualityMetric, PyIQAMetric + + +def get_torch_device() -> Any: + """Retain the benchmark's CUDA, MPS, then CPU device preference.""" + import torch -def get_torch_device(): - """Return a torch device string preferring CUDA when available.""" if torch.cuda.is_available(): return torch.device("cuda") if torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") -def open_image(path: str) -> Image.Image: - """Open image from disk as RGB using PIL.""" - with Image.open(path) as im: - return im.convert('RGB') - -class BaseMetric(ABC): - """BaseMetric Class.""" +class BaseMetric(ImageQualityMetric): + """Use the benchmark's device preference for image comparisons.""" def __init__(self) -> None: - self._metric = None - self._device = get_torch_device() - - def _process_image( - self, - rendered_images: List[Union[str, Image.Image]], - ) -> float: - preprocessing = transforms.Compose( - [ - transforms.Resize((512, 512)), - transforms.ToTensor(), - ] - ) - - rendered_images_: List[torch.Tensor] = [] - for image in rendered_images: - # Handle the rendered image input - if isinstance(image, str): - image = preprocessing(open_image(image)) - else: - image = preprocessing(image) - rendered_images_.append(image) - - img: torch.Tensor = torch.stack(rendered_images_).to(self._device) - - return img + super().__init__(device=get_torch_device()) - def _process_images( - self, - rendered_images: List[Union[str, Image.Image]], - reference_image: Union[str, Image.Image], - ) -> float: - preprocessing = transforms.Compose( - [ - transforms.Resize((512, 512)), - transforms.ToTensor(), - ] - ) - # Handle the reference image input - if isinstance(reference_image, str): - reference_image = preprocessing(open_image(reference_image)) +class IQAPytorchMetric(PyIQAMetric): + """Apply 4DWorldBench's input policy to the selected PyIQA metric.""" - rendered_images_: List[torch.Tensor] = [] - reference_images_: List[torch.Tensor] = [] - for image in rendered_images: - # Handle the rendered image input - if isinstance(image, str): - image = preprocessing(open_image(image)) - else: - image = preprocessing(image) - rendered_images_.append(image) - - reference_images_.append(reference_image) - - img1: torch.Tensor = torch.stack(rendered_images_).to(self._device) - img2: torch.Tensor = torch.stack(reference_images_).to(self._device) - - return img1, img2 - - def _process_np_to_tensor( - self, - rendered_image: np.ndarray, - reference_image: np.ndarray, - ) -> Tuple[torch.Tensor, torch.Tensor]: - img1 = ToTensor()(rendered_image).unsqueeze(0).to(self._device) - img2 = ToTensor()(reference_image).unsqueeze(0).to(self._device) - return img1, img2 - - @abstractmethod - def _compute_scores(self, *args): - pass - - -class IQAPytorchMetric(BaseMetric): def __init__(self, metric_name: str) -> None: - super().__init__() - self._metric = self._create_metric(metric_name).to(self._device) - - def _create_metric(self, metric: str) -> InferenceModel: - metric = pyiqa.create_metric(metric) - return metric + super().__init__(metric_name, device=get_torch_device()) diff --git a/worldfoundry/evaluation/tasks/execution/runners/mind/runtime/mind/src/process.py b/worldfoundry/evaluation/tasks/execution/runners/mind/runtime/mind/src/process.py index 1b8ea79cd..d748135a7 100644 --- a/worldfoundry/evaluation/tasks/execution/runners/mind/runtime/mind/src/process.py +++ b/worldfoundry/evaluation/tasks/execution/runners/mind/runtime/mind/src/process.py @@ -1,30 +1,54 @@ -import lpips -import torch -from pyiqa.archs.musiq_arch import MUSIQ -from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure -from utils.utils import load_time_from_json, print_gpu_memory, get_musiq_spaq_path, get_vitl_path, get_aes_path, extract_actions_from_json, ensure_all_models_downloaded, CACHE_DIR, VideoStreamReader, get_video_length -from utils.dino_utils import load_dinov3_model, extract_dinov3_features -from tqdm import tqdm -import threading -import torch.nn.functional as F import json -import clip import os -from pathlib import Path -import torch.multiprocessing as mp -import warnings -from metrics.action import action_accuracy_metric -from metrics.lcm import lcm_metric, merge_lcm_results -from metrics.visual_quality import visual_quality_metric, merge_visual_results -from metrics.dino import dino_mse_metric, merge_dino_results +import threading import time +import warnings from collections import namedtuple +from queue import Empty -Task = namedtuple('Task', ['data', 'retry_count']) - -def compute_metrics_single_gpu(task_queue, result_list, gt_root, test_root, dino_path, - requested_metrics, video_max_time, process_batch_size, device, gpu_id, stop_event): +import lpips +import torch +import torch.multiprocessing as mp +from metrics.action import action_accuracy_metric +from metrics.dino import dino_mse_metric, merge_dino_results +from metrics.lcm import lcm_metric, merge_lcm_results +from metrics.visual_quality import merge_visual_results, visual_quality_metric +from pyiqa.archs.musiq_arch import MUSIQ +from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure +from tqdm import tqdm +from utils.dino_utils import load_dinov3_model +from utils.utils import ( + VideoStreamReader, + ensure_all_models_downloaded, + extract_actions_from_json, + get_aes_path, + get_musiq_spaq_path, + get_video_length, + get_vitl_path, + load_time_from_json, +) + +from worldfoundry.base_models.perception_core.general_perception import openai_clip as clip +from worldfoundry.evaluation.tasks.metrics._shared.aesthetic import load_laion_aesthetic_linear_head + +Task = namedtuple("Task", ["data", "retry_count"]) + + +def compute_metrics_single_gpu( + task_queue, + result_list, + gt_root, + test_root, + dino_path, + requested_metrics, + video_max_time, + process_batch_size, + device, + gpu_id, + stop_event, +): import warnings + warnings.filterwarnings("ignore", message=".*pretrained.*") warnings.filterwarnings("ignore", message=".*Weights.*") warnings.filterwarnings("ignore", message=".*video.*deprecated.*") @@ -34,28 +58,23 @@ def compute_metrics_single_gpu(task_queue, result_list, gt_root, test_root, dino 在单个GPU上从任务队列获取并处理任务 """ # 初始化模型 - if 'lcm' in requested_metrics or 'gsc' in requested_metrics: + if "lcm" in requested_metrics or "gsc" in requested_metrics: tqdm.write(f"GPU[{gpu_id}]: loading lcm model") - lpips_metric = lpips.LPIPS(net='alex', spatial=False).to(device) - ssim_metric = StructuralSimilarityIndexMeasure(data_range=1.0, reduction='none').to(device) - psnr_metric = PeakSignalNoiseRatio(data_range=1.0, reduction='none', dim=[1, 2, 3]).to(device) + lpips_metric = lpips.LPIPS(net="alex", spatial=False).to(device) + ssim_metric = StructuralSimilarityIndexMeasure(data_range=1.0, reduction="none").to(device) + psnr_metric = PeakSignalNoiseRatio(data_range=1.0, reduction="none", dim=[1, 2, 3]).to(device) - if 'visual' in requested_metrics: + if "visual" in requested_metrics: tqdm.write(f"GPU[{gpu_id}]: loading visual quality model") imaging_model = MUSIQ(pretrained_model_path=get_musiq_spaq_path()) - aesthetic_model = torch.nn.Linear(768, 1) + aesthetic_model = load_laion_aesthetic_linear_head(get_aes_path()).to(device) clip_model, preprocess = clip.load(get_vitl_path(), device=device) - s = torch.load(get_aes_path(), weights_only=False) - aesthetic_model.load_state_dict(s) - aesthetic_model.to(device) - aesthetic_model.eval() - imaging_model.to(device) imaging_model.training = False # 加载DINOv3模型 - if 'dino' in requested_metrics: + if "dino" in requested_metrics: tqdm.write(f"GPU[{gpu_id}]: loading dinov3 model") dino_model, dino_processor = load_dinov3_model(dino_path, device) @@ -72,72 +91,85 @@ def compute_metrics_single_gpu(task_queue, result_list, gt_root, test_root, dino task = task_queue.get(timeout=5) data = task.data retry_count = task.retry_count - except: + except Empty: continue - data_path = data['path'] - gt_dir = os.path.join(gt_root, data['perspective'], 'test', data['test_type']) - test_dir = os.path.join(test_root, data['perspective'], data['test_type']) + data_path = data["path"] + gt_dir = os.path.join(gt_root, data["perspective"], "test", data["test_type"]) + test_dir = os.path.join(test_root, data["perspective"], data["test_type"]) prefix = f"[GPU{gpu_id}] {data_path}" try: - if data['test_type'] == 'mirror_test': - if 'gsc' not in requested_metrics: + if data["test_type"] == "mirror_test": + if "gsc" not in requested_metrics: continue result = data.copy() - path_videos = [v for v in os.listdir(os.path.join(test_dir, data_path)) if v.endswith('.mp4')] - result['video_results'] = [] - result['videos'] = path_videos - result['error'] = None + path_videos = [v for v in os.listdir(os.path.join(test_dir, data_path)) if v.endswith(".mp4")] + result["video_results"] = [] + result["videos"] = path_videos + result["error"] = None for video in path_videos: - vid_result = {'video_name': video, 'error': None} + vid_result = {"video_name": video, "error": None} sample_frames = get_video_length(os.path.join(test_dir, data_path, video)) - vid_result['mark_time'] = sample_frames // 2 - vid_result['sample_frames'] = sample_frames - if sample_frames%2 == 1: - vid_result['error'] = 'frame count is not an even number!' + vid_result["mark_time"] = sample_frames // 2 + vid_result["sample_frames"] = sample_frames + if sample_frames % 2 == 1: + vid_result["error"] = "frame count is not an even number!" tqdm.write(f"{prefix}: Task failed because frame count is not an even number!") - result['video_results'].append(vid_result) + result["video_results"].append(vid_result) continue tqdm.write(f"{prefix}: [1/2] Reading videos...") - sample_reader = VideoStreamReader(os.path.join(test_dir, data_path, video), start_frame=0, total_frames=sample_frames) + sample_reader = VideoStreamReader( + os.path.join(test_dir, data_path, video), start_frame=0, total_frames=sample_frames + ) _, imgs = sample_reader.read_batch(sample_frames) - origin_pred = imgs[:sample_frames // 2] - mirror_pred = torch.flip(imgs[sample_frames // 2:], dims=[0]) + origin_pred = imgs[: sample_frames // 2] + mirror_pred = torch.flip(imgs[sample_frames // 2 :], dims=[0]) tqdm.write(f"{prefix}: part len: {len(origin_pred)}, {len(mirror_pred)}") - + tqdm.write(f"{prefix}: [2/2] Computing GSC metrics (MSE/PSNR/SSIM/LPIPS)...") - gsc = lcm_metric(origin_pred, mirror_pred, lpips_metric, ssim_metric, psnr_metric, process_batch_size, device) - vid_result['gsc'] = gsc - result['video_results'].append(vid_result) + gsc = lcm_metric( + origin_pred, mirror_pred, lpips_metric, ssim_metric, psnr_metric, process_batch_size, device + ) + vid_result["gsc"] = gsc + result["video_results"].append(vid_result) - del sample_reader + del sample_reader torch.cuda.empty_cache() else: - mark_time, total_time = load_time_from_json(os.path.join(gt_dir, data_path, 'action.json')) + mark_time, total_time = load_time_from_json(os.path.join(gt_dir, data_path, "action.json")) result = data.copy() - result['error'] = None - result['mark_time'] = mark_time - result['total_time'] = total_time + result["error"] = None + result["mark_time"] = mark_time + result["total_time"] = total_time - gt_frames = get_video_length(os.path.join(gt_dir, data_path, 'video.mp4')) - sample_frames = get_video_length(os.path.join(test_dir, data_path, 'video.mp4')) - tqdm.write(f"{prefix}: gt_frames={gt_frames}, sample={sample_frames}, total_time={total_time}, mark_time={mark_time}") - result['sample_frames'] = sample_frames + gt_frames = get_video_length(os.path.join(gt_dir, data_path, "video.mp4")) + sample_frames = get_video_length(os.path.join(test_dir, data_path, "video.mp4")) + tqdm.write( + f"{prefix}: gt_frames={gt_frames}, sample={sample_frames}, total_time={total_time}, mark_time={mark_time}" + ) + result["sample_frames"] = sample_frames if sample_frames != total_time - mark_time: real_time = min(total_time - mark_time, sample_frames) - else: real_time = sample_frames + else: + real_time = sample_frames if video_max_time is not None: real_time = min(real_time, video_max_time) - gt_reader = VideoStreamReader(os.path.join(gt_dir, data_path, 'video.mp4'), start_frame=mark_time, total_frames=mark_time + real_time) - sample_reader = VideoStreamReader(os.path.join(test_dir, data_path, 'video.mp4'), start_frame=0, total_frames=real_time) + gt_reader = VideoStreamReader( + os.path.join(gt_dir, data_path, "video.mp4"), + start_frame=mark_time, + total_frames=mark_time + real_time, + ) + sample_reader = VideoStreamReader( + os.path.join(test_dir, data_path, "video.mp4"), start_frame=0, total_frames=real_time + ) while True: tqdm.write(f"{prefix}: [1/5] Reading videos...") @@ -147,42 +179,50 @@ def compute_metrics_single_gpu(task_queue, result_list, gt_root, test_root, dino if is_ended or gt_imgs is None or sample_imgs is None: break - if 'lcm' in requested_metrics: + if "lcm" in requested_metrics: tqdm.write(f"{prefix}: [2/5] Computing LCM metrics (MSE/PSNR/SSIM/LPIPS)...") - lcm = lcm_metric(sample_imgs, gt_imgs, lpips_metric, ssim_metric, psnr_metric, process_batch_size, device) - result['lcm'] = merge_lcm_results(result.get('lcm'), lcm) + lcm = lcm_metric( + sample_imgs, gt_imgs, lpips_metric, ssim_metric, psnr_metric, process_batch_size, device + ) + result["lcm"] = merge_lcm_results(result.get("lcm"), lcm) - if 'visual' in requested_metrics: + if "visual" in requested_metrics: tqdm.write(f"{prefix}: [3/5] Computing visual quality metrics...") - vq = visual_quality_metric(sample_imgs, imaging_model, aesthetic_model, clip_model, process_batch_size, device) - result['visual_quality'] = merge_visual_results(result.get('visual_quality'), vq) + vq = visual_quality_metric( + sample_imgs, imaging_model, aesthetic_model, clip_model, process_batch_size, device + ) + result["visual_quality"] = merge_visual_results(result.get("visual_quality"), vq) - if 'dino' in requested_metrics: + if "dino" in requested_metrics: tqdm.write(f"{prefix}: [4/5] Computing dino metrics...") - dino = dino_mse_metric(sample_imgs, gt_imgs, dino_model, dino_processor, device, process_batch_size) - result['dino'] = merge_dino_results(result.get('dino'), dino) + dino = dino_mse_metric( + sample_imgs, gt_imgs, dino_model, dino_processor, device, process_batch_size + ) + result["dino"] = merge_dino_results(result.get("dino"), dino) del gt_imgs, sample_imgs del gt_reader, sample_reader torch.cuda.empty_cache() - if 'action' in requested_metrics: + if "action" in requested_metrics: tqdm.write(f"{prefix}: [4/5] Computing action accuracy (ViPE)...") # 计算action accuracy - actions = extract_actions_from_json(os.path.join(gt_dir, data_path, 'action.json'), mark_time, 97) + actions = extract_actions_from_json( + os.path.join(gt_dir, data_path, "action.json"), mark_time, 97 + ) action = action_accuracy_metric( - os.path.join(test_dir, data_path, 'video.mp4'), - os.path.join(gt_dir, data_path, 'video.mp4'), + os.path.join(test_dir, data_path, "video.mp4"), + os.path.join(gt_dir, data_path, "video.mp4"), mark_time, actions, - max_frames = 97, - gt_data_dir = os.path.join(gt_dir, data_path), + max_frames=97, + gt_data_dir=os.path.join(gt_dir, data_path), verbose_prefix=f" {prefix}", - gpu_id=gpu_id + gpu_id=gpu_id, ) if action is not None: - result['action'] = action + result["action"] = action tqdm.write(f"{prefix}: Finish Task!") @@ -192,12 +232,16 @@ def compute_metrics_single_gpu(task_queue, result_list, gt_root, test_root, dino stop_event.set() raise except Exception as e: - result['error'] = str(e) + result["error"] = str(e) if retry_count < 3: - tqdm.write(f"{prefix}: Error processing {data_path} on GPU{gpu_id}: {e}, retrying ({retry_count + 1}/3)") + tqdm.write( + f"{prefix}: Error processing {data_path} on GPU{gpu_id}: {e}, retrying ({retry_count + 1}/3)" + ) task_queue.put(Task(data, retry_count + 1)) else: - tqdm.write(f"{prefix}: Error processing {data_path} on GPU{gpu_id}: {e}, max retries exceeded, giving up") + tqdm.write( + f"{prefix}: Error processing {data_path} on GPU{gpu_id}: {e}, max retries exceeded, giving up" + ) result_list.append(result) continue @@ -206,42 +250,60 @@ def compute_metrics_single_gpu(task_queue, result_list, gt_root, test_root, dino finally: tqdm.write(f"GPU{gpu_id}: Cleaning up...") - if 'lcm' in requested_metrics or 'gsc' in requested_metrics: + if "lcm" in requested_metrics or "gsc" in requested_metrics: del lpips_metric, ssim_metric, psnr_metric - if 'visual' in requested_metrics: + if "visual" in requested_metrics: del imaging_model, aesthetic_model, clip_model torch.cuda.empty_cache() -def compute_metrics(gt_root, test_root, dino_path, output_path, requested_metrics=['lcm', 'visual', 'dino', 'action'], - video_max_time=100, process_batch_size=10, num_gpus=1): - result_dict = {'data':[], 'video_max_time':video_max_time} - mp.set_start_method('spawn', force=True) +def compute_metrics( + gt_root, + test_root, + dino_path, + output_path, + requested_metrics=["lcm", "visual", "dino", "action"], + video_max_time=100, + process_batch_size=10, + num_gpus=1, +): + result_dict = {"data": [], "video_max_time": video_max_time} + + mp.set_start_method("spawn", force=True) all_data = [] - for perspective in ['1st_data', '3rd_data']: - for test_type in ['mem_test', 'action_space_test', 'mirror_test']: - if test_type == 'mirror_test': - if 'gsc' in requested_metrics: + for perspective in ["1st_data", "3rd_data"]: + for test_type in ["mem_test", "action_space_test", "mirror_test"]: + if test_type == "mirror_test": + if "gsc" in requested_metrics: test_dir = os.path.join(test_root, perspective, test_type) - all_data += [{'path': d, 'perspective': perspective, 'test_type': test_type } - for d in os.listdir(test_dir)] + all_data += [ + {"path": d, "perspective": perspective, "test_type": test_type} for d in os.listdir(test_dir) + ] else: - if 'lcm' in requested_metrics or 'visual' in requested_metrics or 'dino' in requested_metrics or 'action' in requested_metrics: - gt_dir = os.path.join(gt_root, perspective, 'test', test_type) + if ( + "lcm" in requested_metrics + or "visual" in requested_metrics + or "dino" in requested_metrics + or "action" in requested_metrics + ): + gt_dir = os.path.join(gt_root, perspective, "test", test_type) test_dir = os.path.join(test_root, perspective, test_type) - all_data += [{'path': d, 'perspective': perspective, 'test_type': test_type} - for d in os.listdir(test_dir) if os.path.exists(os.path.join(gt_dir, d))] + all_data += [ + {"path": d, "perspective": perspective, "test_type": test_type} + for d in os.listdir(test_dir) + if os.path.exists(os.path.join(gt_dir, d)) + ] if len(all_data) == 0: - tqdm.write(f"No data found!") + tqdm.write("No data found!") return total_tasks = len(all_data) - tqdm.write(f"\n{'='*60}") + tqdm.write(f"\n{'=' * 60}") tqdm.write(f"Total data: {total_tasks}, Using {num_gpus} GPU(s) with task queue") - tqdm.write(f"{'='*60}\n") + tqdm.write(f"{'=' * 60}\n") ensure_all_models_downloaded() @@ -263,8 +325,8 @@ def monitor_progress(): if current_count > last_count: pbar.update(current_count - last_count) last_count = current_count - with open(output_path, 'w') as f: - result_dict['data'] = list(result_list) + with open(output_path, "w") as f: + result_dict["data"] = list(result_list) json.dump(result_dict, f, indent=2) if last_count >= total_tasks: break @@ -281,15 +343,26 @@ def monitor_progress(): try: with mp.get_context("spawn").Pool(processes=num_gpus) as pool: worker_args = [ - (task_queue, result_list, gt_root, test_root, dino_path, - requested_metrics, video_max_time, process_batch_size, f'cuda:{i}', i, stop_event) + ( + task_queue, + result_list, + gt_root, + test_root, + dino_path, + requested_metrics, + video_max_time, + process_batch_size, + f"cuda:{i}", + i, + stop_event, + ) for i in range(num_gpus) ] pool.starmap(compute_metrics_single_gpu, worker_args) monitor_thread.join() - result_dict['data'] = list(result_list) + result_dict["data"] = list(result_list) except KeyboardInterrupt: tqdm.write("\n\nInterrupted by user! Merging partial results...") @@ -297,35 +370,47 @@ def monitor_progress(): pool.terminate() pool.join() monitor_thread.join(timeout=2) - result_dict['data'] = list(result_list) + result_dict["data"] = list(result_list) tqdm.write(f" Merged {len(result_dict['data'])} results from shared list") return result_dict -if __name__ == '__main__': + +if __name__ == "__main__": import argparse + warnings.filterwarnings("ignore", message=".*pretrained.*") warnings.filterwarnings("ignore", message=".*Weights.*") warnings.filterwarnings("ignore", message=".*video.*deprecated.*") warnings.filterwarnings("ignore", message=".*timm.models.layers.*") warnings.filterwarnings("ignore", message=".*TRANSFORMERS_CACHE.*") - parser = argparse.ArgumentParser(description='Compute video metrics with multi-GPU support') - parser.add_argument('--gt_root', type=str, required=True, help='Ground truth data root directory') - parser.add_argument('--test_root', type=str, required=True, help='Test data root directory') - parser.add_argument('--dino_path', type=str, default='./dinov3_vitb16', - help='dinov3 weight directory, for example ./dinov3_vitb16') - parser.add_argument('--num_gpus', type=int, default=1, help='Number of GPUs to use (default: 1)') - parser.add_argument('--video_max_time', type=int, default=None, help='Maximum video frames (default: None = use all frames)') - parser.add_argument('--output', type=str, default=None, help='Output JSON file path') - parser.add_argument('--metrics', type=str, default='lcm,visual,dino,action,gsc', help='Requested metrics to compute, comma separated (e.g. dino,visual)') + parser = argparse.ArgumentParser(description="Compute video metrics with multi-GPU support") + parser.add_argument("--gt_root", type=str, required=True, help="Ground truth data root directory") + parser.add_argument("--test_root", type=str, required=True, help="Test data root directory") + parser.add_argument( + "--dino_path", type=str, default="./dinov3_vitb16", help="dinov3 weight directory, for example ./dinov3_vitb16" + ) + parser.add_argument("--num_gpus", type=int, default=1, help="Number of GPUs to use (default: 1)") + parser.add_argument( + "--video_max_time", type=int, default=None, help="Maximum video frames (default: None = use all frames)" + ) + parser.add_argument("--output", type=str, default=None, help="Output JSON file path") + parser.add_argument( + "--metrics", + type=str, + default="lcm,visual,dino,action,gsc", + help="Requested metrics to compute, comma separated (e.g. dino,visual)", + ) args = parser.parse_args() # 检查GPU数量 available_gpus = torch.cuda.device_count() if args.num_gpus > available_gpus: - tqdm.write(f"Warning: Requested {args.num_gpus} GPUs but only {available_gpus} available. Using {available_gpus} GPUs.") + tqdm.write( + f"Warning: Requested {args.num_gpus} GPUs but only {available_gpus} available. Using {available_gpus} GPUs." + ) args.num_gpus = available_gpus # 准备输出路径 @@ -333,7 +418,8 @@ def monitor_progress(): output_path = args.output else: from datetime import datetime - output_path = f'result_{os.path.basename(args.test_root)}_{datetime.now().strftime("%Y-%m-%d-%H:%M:%S")}.json' + + output_path = f"result_{os.path.basename(args.test_root)}_{datetime.now().strftime('%Y-%m-%d-%H:%M:%S')}.json" tqdm.write(f"Starting computation with {args.num_gpus} GPU(s)...") tqdm.write(f"Output will be saved to: {output_path}") @@ -345,14 +431,14 @@ def monitor_progress(): args.test_root, args.dino_path, output_path, - requested_metrics=[m.strip() for m in args.metrics.split(',')], + requested_metrics=[m.strip() for m in args.metrics.split(",")], video_max_time=args.video_max_time, - num_gpus=args.num_gpus + num_gpus=args.num_gpus, ) except KeyboardInterrupt: - tqdm.write("\n\n" + "="*60) + tqdm.write("\n\n" + "=" * 60) tqdm.write("INTERRUPTED BY USER (Ctrl+C)") - tqdm.write("="*60) + tqdm.write("=" * 60) if video_results is None: tqdm.write("No results to save (computation was interrupted too early)") @@ -362,9 +448,8 @@ def monitor_progress(): tqdm.write(f"{len(video_results['data'])} videos are processed") tqdm.write(f"checking output path {output_path}:") if os.path.exists(output_path): - with open(output_path, 'r') as f: + with open(output_path, "r") as f: saved_results = json.load(f) tqdm.write(f"{len(saved_results['data'])} videos are saved") else: tqdm.write("\nNo results to save.") - \ No newline at end of file diff --git a/worldfoundry/evaluation/tasks/execution/runners/wbench/runtime/wbench/src/metrics/video_quality/imaging_quality.py b/worldfoundry/evaluation/tasks/execution/runners/wbench/runtime/wbench/src/metrics/video_quality/imaging_quality.py index 69acd582e..643c1866f 100644 --- a/worldfoundry/evaluation/tasks/execution/runners/wbench/runtime/wbench/src/metrics/video_quality/imaging_quality.py +++ b/worldfoundry/evaluation/tasks/execution/runners/wbench/runtime/wbench/src/metrics/video_quality/imaging_quality.py @@ -1,10 +1,11 @@ -"""Imaging quality metric — MUSIQ (via pyiqa).""" +"""Imaging quality metric — the official MUSIQ-KonIQ variant via PyIQA.""" + import numpy as np import torch import torchvision.transforms as T from ..base import BaseMetric -from ..weight_utils import setup_torch_hub_dir, get_weights_dir +from ..weight_utils import wbench_asset_path def _vbench_transform(img_tensor, max_edge=512): @@ -20,13 +21,13 @@ def _vbench_transform(img_tensor, max_edge=512): class ImagingQualityMetric(BaseMetric): def __init__(self, device="cuda"): super().__init__(device) - setup_torch_hub_dir() - import os import pyiqa - import pyiqa.utils.download_util as _dl - # Keep pyiqa aligned with the WBench MUSIQ asset location. - _dl.DEFAULT_CACHE_DIR = get_weights_dir("pyiqa") - self.model = pyiqa.create_metric('musiq', device=self.device) + + self.model = pyiqa.create_metric( + "musiq", + device=self.device, + pretrained_model_path=str(wbench_asset_path("wbench_musiq_koniq_checkpoint")), + ) @property def name(self): diff --git a/worldfoundry/evaluation/tasks/execution/runners/wbench/runtime/wbench/src/metrics/weight_utils.py b/worldfoundry/evaluation/tasks/execution/runners/wbench/runtime/wbench/src/metrics/weight_utils.py index 97469dc0c..1535d2bbd 100644 --- a/worldfoundry/evaluation/tasks/execution/runners/wbench/runtime/wbench/src/metrics/weight_utils.py +++ b/worldfoundry/evaluation/tasks/execution/runners/wbench/runtime/wbench/src/metrics/weight_utils.py @@ -13,7 +13,7 @@ _DIR_ASSET_BY_SUBDIR = { "clip": "wbench_clip_vit_b32_checkpoint", - "pyiqa": "wbench_musiq_spaq_checkpoint", + "pyiqa": "wbench_musiq_koniq_checkpoint", } _LEGACY_ASSET_RELATIVE_PATHS = { @@ -23,9 +23,9 @@ "aesthetic/sa_0_4_vit_l_14_linear.pth", "aesthetic_model/emb_reader/sa_0_4_vit_l_14_linear.pth", ), - "wbench_musiq_spaq_checkpoint": ( - "pyiqa/musiq_spaq_ckpt-358bb6af.pth", - "pyiqa_model/musiq_spaq_ckpt-358bb6af.pth", + "wbench_musiq_koniq_checkpoint": ( + "pyiqa/musiq_koniq_ckpt-e95806b9.pth", + "pyiqa_model/musiq_koniq_ckpt-e95806b9.pth", ), "wbench_clip_vit_b16_model_dir": ( "clip-vit-base-patch16", diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py index 18475834d..26a6e863c 100644 --- a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/run_worldmembench_official_runner.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 WorldFoundry contributors # SPDX-License-Identifier: Apache-2.0 -"""Run ZJU-REAL/WorldMemBench or import its official per-model results.""" +"""Run the in-tree WorldMemBench evaluator or import official per-model results.""" from __future__ import annotations @@ -29,7 +29,7 @@ CONFIG = BenchRunnerConfig( benchmark_id="worldmembench", display_name="WorldMemBench", - root_env="WORLDFOUNDRY_WORLDMEMBENCH_ROOT", + root_env="", results_path_env="WORLDFOUNDRY_WORLDMEMBENCH_RESULTS_PATH", default_repo_subdir="", metric_order=tuple(METRIC_SPECS), @@ -43,7 +43,7 @@ def build_parser() -> argparse.ArgumentParser: - """Expose the standard benchmark runner flags and upstream runtime inputs.""" + """Expose the standard benchmark runner flags and metric asset configuration.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--benchmark-id", choices=["worldmembench"], default="worldmembench") parser.add_argument( @@ -59,12 +59,6 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--run-official", action="store_true") - parser.add_argument( - "--repo-root", - type=Path, - default=_env_path(CONFIG.root_env), - help="Optional official checkout; otherwise use the installed worldmembench package", - ) parser.add_argument("--python", default=os.environ.get("WORLDFOUNDRY_WORLDMEMBENCH_PYTHON", sys.executable)) parser.add_argument("--config", type=Path, default=_env_path("WORLDFOUNDRY_WORLDMEMBENCH_CONFIG")) parser.add_argument("--weights-dir", type=Path, default=_env_path("WORLDFOUNDRY_WORLDMEMBENCH_WEIGHTS_DIR")) @@ -87,13 +81,13 @@ def build_parser() -> argparse.ArgumentParser: def build_official_command(args: argparse.Namespace, results: Path, metrics: tuple[str, ...]) -> list[str]: - """Delegate metric computation and native multi-GPU execution to upstream.""" + """Run the in-tree official protocol, including its multi-GPU sharding.""" if args.generated_video_dir is None or not args.generated_video_dir.exists(): raise ValueError("--run-official requires an existing --generated-video-dir or JSONL --input-path") command = [ args.python, "-m", - "worldmembench", + "worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime", "eval", str(args.generated_video_dir.resolve()), "--output", @@ -137,21 +131,12 @@ def run(args: argparse.Namespace) -> tuple[dict[str, Any], int]: raise ValueError(f"Use a fresh output directory; official results already exist: {results}") command = build_official_command(args, results, metrics) environment = dict(os.environ) - if args.repo_root is not None: - args.repo_root = args.repo_root.expanduser().resolve() - source = args.repo_root / "src" - if not (source / "worldmembench" / "__main__.py").is_file(): - raise ValueError(f"Not a WorldMemBench checkout: {args.repo_root}") - environment["PYTHONPATH"] = os.pathsep.join( - filter(None, [str(source), environment.get("PYTHONPATH", "")]) - ) with ( - (output / "upstream_stdout.log").open("w", encoding="utf-8") as stdout, - (output / "upstream_stderr.log").open("w", encoding="utf-8") as stderr, + (output / "runtime_stdout.log").open("w", encoding="utf-8") as stdout, + (output / "runtime_stderr.log").open("w", encoding="utf-8") as stderr, ): completed = subprocess.run( command, - cwd=args.repo_root, env=environment, stdout=stdout, stderr=stderr, @@ -178,7 +163,6 @@ def run(args: argparse.Namespace) -> tuple[dict[str, Any], int]: command=command, duration_seconds=time.monotonic() - started, returncode=returncode, - repo_root=args.repo_root, generated_video_dir=args.generated_video_dir, run_status_override="failed" if returncode else None, ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/LICENSE b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/LICENSE new file mode 100644 index 000000000..dce4b5ccf --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 WorldMemBench + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/__init__.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/__init__.py new file mode 100644 index 000000000..e330cec41 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""WorldMemBench.""" diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/__main__.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/__main__.py new file mode 100644 index 000000000..9b3c4944c --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/__main__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.cli import main + +raise SystemExit(main()) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/cli.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/cli.py new file mode 100644 index 000000000..7ce38edb4 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/cli.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""WorldMemBench command-line entry point.""" + +import argparse +import sys +from pathlib import Path + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.data import load_samples, write_manifest +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.evaluate import evaluate +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics import registry +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.paths import DEFAULT_WEIGHTS_DIR + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="worldmembench", description="WorldMemBench evaluation") + commands = parser.add_subparsers(dest="command", required=True) + commands.add_parser("list-metrics", help="List available metrics") + run = commands.add_parser("eval", help="Evaluate a video, video directory, or JSONL manifest") + run.add_argument("input", help="Video file, directory, or JSONL manifest") + run.add_argument("--metrics", nargs="+", help="Metrics to run (default: all available)") + run.add_argument("--output", help="Result directory (default: eval_results/)") + devices = run.add_mutually_exclusive_group() + devices.add_argument("--device", help="Single device, e.g. cuda:0 or cpu (default: cuda)") + devices.add_argument("--gpus", help="Comma-separated visible GPU indices, e.g. 0,1; one process per GPU") + run.add_argument("--weights-dir", help=f"Weight directory (default: {DEFAULT_WEIGHTS_DIR})") + run.add_argument("--config", help="YAML settings for checkpoints, metric batch sizes, and the DA3 runtime") + vlm = run.add_argument_group("Dynamic Memory VLM") + vlm.add_argument( + "--vlm-backend", choices=("local", "api"), help="Self-hosted vLLM (default: local) or external API" + ) + vlm.add_argument("--vlm-base-url", help="VLM endpoint (default: http://127.0.0.1:8001/v1)") + vlm.add_argument("--vlm-model", help="Served model ID (default: Qwen3.8-27B)") + vlm.add_argument( + "--vlm-api-key-env", + help="API key environment variable name, never the key itself (default: WORLDMEMBENCH_API_KEY)", + ) + run.add_argument( + "--expected-revisits", + type=int, + choices=(1, 2, 3), + help="Fixed Human/Object revisit protocol, e.g. 2 for LRLR; omitted: preserve manifest K, otherwise detect returns automatically", + ) + manifest = commands.add_parser("manifest", help="Create an input manifest without evaluating videos") + manifest.add_argument("input", help="Video directory, file, or existing manifest") + manifest.add_argument("--output", required=True, help="New output JSONL manifest") + for command in (run, manifest): + command.add_argument("--models", nargs="+", help="Filter generation model IDs") + command.add_argument("--categories", nargs="+", help="Filter categories, e.g. human object environment") + command.add_argument("--trajectories", nargs="+", help="Filter trajectory IDs") + command.add_argument("--cases", nargs="+", help="Filter case IDs, e.g. case_002") + args = parser.parse_args(argv) + if args.command == "list-metrics": + print("\n".join(registry.names()) or "WorldMemBench has no metrics yet.") + return 0 + try: + filters = dict(models=args.models, categories=args.categories, trajectories=args.trajectories, cases=args.cases) + if args.command == "manifest": + samples = load_samples(args.input, **filters) + output = write_manifest(samples, args.output) + print(f"Saved {len(samples)} samples to {output}") + return 0 + config = {} + if args.config: + try: + import yaml + except ImportError as error: + raise RuntimeError("--config requires PyYAML") from error + try: + config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8")) or {} + except yaml.YAMLError as error: + raise ValueError(f"invalid YAML config: {error}") from error + if not isinstance(config, dict): + raise ValueError("config must be a YAML mapping") + vlm_options = { + key: value + for key, value in { + "backend": args.vlm_backend, + "base_url": args.vlm_base_url, + "model": args.vlm_model, + "api_key_env": args.vlm_api_key_env, + }.items() + if value is not None + } + if vlm_options: + config.setdefault("metrics", {}).setdefault("dynamic_memory", {}).update(vlm_options) + output, success = evaluate( + args.input, + metrics=args.metrics, + output=args.output, + expected_revisits=args.expected_revisits, + **filters, + device="cuda" if args.gpus is not None else args.device or config.get("device", "cuda"), + gpus=args.gpus, + weights_dir=args.weights_dir or config.get("cache_dir", DEFAULT_WEIGHTS_DIR), + config=config, + command=list(sys.orig_argv) if argv is None else ["worldmembench", *argv], + ) + except (ValueError, OSError, RuntimeError) as error: + print(f"error: {error}", file=sys.stderr) + return 1 + print(f"Results saved to {output}") + if not success: + print("Some cases failed; see error fields in metrics/*/cases.jsonl.", file=sys.stderr) + return 0 if success else 1 diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/context.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/context.py new file mode 100644 index 000000000..418d55ddf --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/context.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Per-run configuration, resource lifecycle, and generic value cache.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable, Hashable +from pathlib import Path +from typing import Any, TypeVar, cast + +T = TypeVar("T") + + +class EvaluationContext: + def __init__( + self, + *, + device: str, + cache_dir: str | Path, + config: dict[str, Any] | None = None, + logger: logging.Logger | None = None, + ) -> None: + self.device = device + self.cache_dir = Path(cache_dir).expanduser().resolve() + self.cache_dir.mkdir(parents=True, exist_ok=True) + self.config = dict(config or {}) + self.logger = logger or logging.getLogger("worldmembench") + self._resources: dict[Hashable, Any] = {} + self._cache: dict[Hashable, Any] = {} + + def metric_config(self, metric_name: str) -> dict[str, Any]: + metrics = self.config.get("metrics") or {} + value = metrics.get(metric_name) or {} + if not isinstance(value, dict): + raise TypeError(f"config.metrics.{metric_name} must be a mapping") + return dict(value) + + def get_resource(self, key: Hashable, factory: Callable[[], T]) -> T: + if key not in self._resources: + self._resources[key] = factory() + return cast(T, self._resources[key]) + + def get_cached(self, key: Hashable, factory: Callable[[], T]) -> T: + if key not in self._cache: + self._cache[key] = factory() + return cast(T, self._cache[key]) + + def close(self) -> None: + for resource in reversed(list(self._resources.values())): + close = getattr(resource, "close", None) + if callable(close): + close() + self._resources.clear() + self._cache.clear() + + def clear_frame_cache(self) -> None: + """Release decoded frames after all metrics have evaluated one sample.""" + self._cache.clear() diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/data.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/data.py new file mode 100644 index 000000000..cee29f92a --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/data.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Load videos directly, or a JSONL manifest with additional sample fields.""" + +import json +import os +from pathlib import Path + +VIDEO_EXTENSIONS = {".mp4", ".mov", ".avi", ".mkv", ".webm"} +ANNOTATION_PATH_FIELDS = ("condition_image", "target_object_mask_first_frame") +ARTIFACT_PATH_FIELDS = ("da3", "da3_scene", "target_poses", "vipe") + + +def _video_sample(path: Path, source: Path) -> dict: + # Both public layouts carry unambiguous labels at a named root. + if ( + len(path.parents) >= 5 + and path.parent.name.startswith("case_") + and path.name == "video.mp4" + and path.parents[4].name in {"samples", "inference_results", "WorldMemBench"} + ): + first, second, trajectory, case_id, _ = path.parts[-5:] + if path.parents[4].name == "samples": + model, category = first, second + else: + category, model = first, second + return { + "item_id": "/".join((model, category, trajectory, case_id)), + "model_id": model, + "category": category, + "condition_id": trajectory, + "sample_id": case_id, + "video": str(path), + } + return {"item_id": path.relative_to(source).as_posix(), "video": str(path)} + + +def load_samples( + source: str | Path, + *, + models: list[str] | None = None, + categories: list[str] | None = None, + trajectories: list[str] | None = None, + cases: list[str] | None = None, +) -> list[dict]: + source = Path(source).expanduser().resolve() + if source.is_dir(): + samples = [ + _video_sample(path, source) + for path in sorted(source.rglob("*")) + if path.is_file() and path.suffix.lower() in VIDEO_EXTENSIONS + ] + elif source.is_file() and source.suffix.lower() == ".jsonl": + samples = [] + for number, line in enumerate(source.read_text(encoding="utf-8").splitlines(), 1): + if not line.strip(): + continue + try: + sample = json.loads(line) + video = (source.parent / sample["video"]).resolve() + if not video.is_file(): + raise ValueError(f"video not found: {video}") + sample["video"] = str(video) + annotations = sample.get("annotations") + if isinstance(annotations, dict): + for field in ANNOTATION_PATH_FIELDS: + if isinstance(annotations.get(field), str) and annotations[field]: + annotations[field] = str((source.parent / annotations[field]).resolve()) + artifacts = sample.get("artifacts") + if isinstance(artifacts, dict): + for field in ARTIFACT_PATH_FIELDS: + if isinstance(artifacts.get(field), str) and artifacts[field]: + artifacts[field] = str((source.parent / artifacts[field]).resolve()) + sample.setdefault("item_id", str(number)) + samples.append(sample) + except (ValueError, KeyError, TypeError) as error: + raise ValueError(f"{source}:{number}: {error}") from error + elif source.is_file() and source.suffix.lower() in VIDEO_EXTENSIONS: + samples = [_video_sample(source, source.parent)] + else: + raise ValueError(f"expected a video, video directory, or JSONL manifest: {source}") + if not samples: + raise ValueError(f"no samples found: {source}") + ids = [sample["item_id"] for sample in samples] + if any(not isinstance(value, str) or not value.strip() for value in ids): + raise ValueError("item_id must be a non-empty string") + if len(ids) != len(set(ids)): + raise ValueError("duplicate item_id in manifest") + for field, values in ( + ("model_id", models), + ("category", categories), + ("condition_id", trajectories), + ("sample_id", cases), + ): + if values: + if any(not sample.get(field) for sample in samples): + raise ValueError(f"filter requires {field}; use the standard samples layout or a manifest") + samples = [sample for sample in samples if sample[field] in values] + if not samples: + raise ValueError("no samples match the requested filters") + return samples + + +def write_manifest(samples: list[dict], output: str | Path) -> Path: + """Save relative media and artifact paths alongside sample annotations.""" + output = Path(output).expanduser().resolve() + output.parent.mkdir(parents=True, exist_ok=True) + with output.open("x", encoding="utf-8") as handle: + for sample in samples: + row = {**sample, "video": os.path.relpath(sample["video"], output.parent)} + if isinstance(sample.get("annotations"), dict): + row["annotations"] = dict(sample["annotations"]) + for field in ANNOTATION_PATH_FIELDS: + value = row["annotations"].get(field) + if isinstance(value, str) and value: + row["annotations"][field] = os.path.relpath(value, output.parent) + if isinstance(sample.get("artifacts"), dict): + row["artifacts"] = dict(sample["artifacts"]) + for field in ARTIFACT_PATH_FIELDS: + value = row["artifacts"].get(field) + if isinstance(value, str) and value: + row["artifacts"][field] = os.path.relpath(value, output.parent) + handle.write(json.dumps(row, ensure_ascii=False, allow_nan=False) + "\n") + return output diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/errors.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/errors.py new file mode 100644 index 000000000..893bfdc94 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/errors.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Structured errors mapped to WorldMemBench result statuses.""" + + +class WorldMemBenchError(Exception): + """Base class for expected WorldMemBench errors.""" + + +class DataInvalidError(WorldMemBenchError): + """The benchmark item or its required annotations are invalid.""" + + +class EvaluationError(WorldMemBenchError): + """The evaluator failed independently of the generated model output.""" + + +class VideoDecodeError(EvaluationError): + """A required video frame could not be decoded.""" + + +class NotApplicableError(WorldMemBenchError): + """The requested metric is not applicable to the item.""" diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/evaluate.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/evaluate.py new file mode 100644 index 000000000..256d43cb6 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/evaluate.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Evaluation orchestration with per-case errors and complete result reporting.""" + +from datetime import datetime +from pathlib import Path + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.context import EvaluationContext +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.data import load_samples, write_manifest +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.execution import ( + evaluate_parallel, + evaluate_shard, + merge_shards, + parse_gpus, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics import MetricRegistry, registry +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.paths import DEFAULT_WEIGHTS_DIR +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.results import ResultStatus, ResultWriter + + +def evaluate( + source: str | Path, + *, + metrics: list[str] | None = None, + models: list[str] | None = None, + categories: list[str] | None = None, + trajectories: list[str] | None = None, + cases: list[str] | None = None, + expected_revisits: int | None = None, + output: str | Path | None = None, + device: str = "cuda", + gpus: str | None = None, + weights_dir: str | Path = DEFAULT_WEIGHTS_DIR, + config: dict | None = None, + metric_registry: MetricRegistry = registry, + command: list[str] | None = None, +) -> tuple[Path, bool]: + names = metric_registry.names() if metrics is None else metrics + if not names: + raise ValueError("No metrics selected. Choose an available WorldMemBench metric.") + if len(names) != len(set(names)): + raise ValueError("each metric may be selected only once") + selected = [metric_registry.get(name) for name in names] + gpu_ids = parse_gpus(gpus) if gpus is not None else None + if gpu_ids is not None and device != "cuda": + raise ValueError("use --device for one device or --gpus for GPU selection, not both") + if gpu_ids is not None: + import torch + + if not torch.cuda.is_available(): + raise RuntimeError("--gpus requires available CUDA devices") + if max(gpu_ids) >= torch.cuda.device_count(): + raise ValueError(f"GPU index out of range; {torch.cuda.device_count()} device(s) visible") + samples = load_samples(source, models=models, categories=categories, trajectories=trajectories, cases=cases) + output = Path(output) if output is not None else Path("eval_results") / datetime.now().strftime("%Y%m%d-%H%M%S") + if output.exists(): + raise FileExistsError(f"output already exists; choose another --output: {output}") + writer = ResultWriter(output) + writer.write_command(command) + if {"human_identity", "human_appearance"}.intersection(names): + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.human.preparation import ( + prepare_for_evaluation, + ) + + preparation_context = EvaluationContext( + device=f"cuda:{gpu_ids[0]}" if gpu_ids is not None else device, + cache_dir=weights_dir, + config=config, + ) + try: + prepare_for_evaluation( + samples, + preparation_context, + expected_revisits=expected_revisits, + weight_metric="human_identity" if "human_identity" in names else "human_appearance", + ) + finally: + preparation_context.close() + if {"object_appearance", "object_geometry"}.intersection(names): + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.object.preparation import ( + prepare_for_evaluation, + ) + + prepare_for_evaluation(samples, expected_revisits=expected_revisits) + if {"environment_visible_memory", "environment_depth_fidelity", "environment_scene_reconstruction"}.intersection( + names + ): + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.environment.preparation import ( + prepare_for_evaluation, + ) + + prepare_for_evaluation(samples) + if "dynamic_memory" in names: + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.dynamic.preparation import ( + prepare_for_evaluation, + ) + + prepare_for_evaluation(samples) + if {"camera_rotation", "camera_translation"}.intersection(names): + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.camera.preparation import ( + prepare_for_evaluation, + ) + + prepare_for_evaluation(samples, required=set(names) <= {"camera_rotation", "camera_translation"}) + write_manifest(samples, output / "manifest.jsonl") + if gpu_ids is None: + shards = [evaluate_shard(samples, selected, device, weights_dir, config)] + else: + devices = [f"cuda:{index}" for index in gpu_ids] + print(f"GPU devices: {', '.join(devices)}; one video shard per active GPU.", flush=True) + shards = evaluate_parallel(samples, selected, devices, weights_dir, config) + item_results = merge_shards(shards, samples, selected) + success = True + for metric in selected: + items = item_results[metric.name] + if any(item.status in {ResultStatus.DATA_INVALID, ResultStatus.EVAL_FAILED} for item in items): + success = False + summary = metric.aggregate_summary(items) + writer.write_metric(metric.name, items, summary) + return output, success diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/execution.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/execution.py new file mode 100644 index 000000000..88f926437 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/execution.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Evaluate video shards and merge item-level results in manifest order.""" + +import json +import multiprocessing +import sys +from dataclasses import asdict +from queue import Empty + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.context import EvaluationContext +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.results import ItemResult, ResultStatus + + +def parse_gpus(value: str) -> list[int]: + parts = value.split(",") + if any(not part.strip().isdigit() for part in parts): + raise ValueError("--gpus requires comma-separated GPU indices, e.g. 0,1") + ids = [int(part) for part in parts] + if len(ids) != len(set(ids)): + raise ValueError("--gpus must not contain duplicate GPU indices") + return ids + + +def evaluate_shard(samples, metrics, device, weights_dir, config): + context = EvaluationContext(device=device, cache_dir=weights_dir, config=config) + items_by_metric = {metric.name: [] for metric in metrics} + try: + names = [metric.name for metric in metrics] + if { + "environment_visible_memory", + "environment_depth_fidelity", + "environment_scene_reconstruction", + }.intersection(names): + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.environment.preparation import ( + prepare_artifacts, + ) + + prepare_artifacts(samples, context, names) + if {"camera_rotation", "camera_translation"}.intersection(names): + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.camera.preparation import ( + prepare_artifacts, + ) + + prepare_artifacts(samples, context, names) + for index, sample in enumerate(samples, 1): + identity = { + key: sample.get(key) for key in ("item_id", "model_id", "category", "sample_id", "condition_id") + } + for metric in metrics: + try: + units = metric.evaluate(sample, context) + if not units: + raise ValueError("metric returned no unit results") + item = metric.aggregate_item(sample, units) + item.status = ResultStatus(item.status) + for key, value in identity.items(): + setattr(item, key, value) + item.num_units = len(units) + json.dumps(asdict(item), allow_nan=False) + except Exception as error: + item = ItemResult( + **identity, + status=( + ResultStatus.DATA_INVALID + if isinstance(error, DataInvalidError) + else ResultStatus.EVAL_FAILED + ), + error=f"{type(error).__name__}: {error}", + ) + if item.status in {ResultStatus.DATA_INVALID, ResultStatus.EVAL_FAILED}: + item.scores = {} + print( + f"[{device}] {metric.name} / {sample['item_id']}: {item.status.value}: {item.error}", + file=sys.stderr, + flush=True, + ) + items_by_metric[metric.name].append(item) + context.clear_frame_cache() + print(f"[{device}] {index}/{len(samples)} {sample['item_id']}", flush=True) + return items_by_metric + finally: + context.close() + + +def _worker(samples, metrics, device, weights_dir, config): + if device.startswith("cuda"): + import torch + + torch.cuda.set_device(device) + torch.set_num_threads(4) + return evaluate_shard(samples, metrics, device, weights_dir, config) + + +def evaluate_parallel(samples, metrics, devices, weights_dir, config): + """Spawn one independent worker per active device; never fork CUDA state.""" + active = devices[: len(samples)] + context = multiprocessing.get_context("spawn") + queue = context.Queue() + processes = [] + received = {} + try: + for index, device in enumerate(active): + process = context.Process( + target=_report_worker, + args=(queue, index, samples[index :: len(active)], metrics, device, weights_dir, config), + ) + process.start() + processes.append(process) + while len(received) < len(processes): + try: + index, result, error = queue.get(timeout=0.2) + except Empty: + for index, process in enumerate(processes): + if index not in received and process.exitcode is not None and process.exitcode != 0: + raise RuntimeError(f"worker {active[index]} exited with code {process.exitcode}") + continue + if error: + raise RuntimeError(error) + received[index] = result + return [received[index] for index in range(len(processes))] + finally: + # Terminate only workers created by this run when any shard fails. + if len(received) < len(processes): + for process in processes: + if process.is_alive(): + process.terminate() + for process in processes: + process.join() + queue.close() + queue.join_thread() + + +def _report_worker(queue, index, samples, metrics, device, weights_dir, config): + try: + result = _worker(samples, metrics, device, weights_dir, config) + except BaseException as error: + queue.put((index, None, f"{type(error).__name__}: {error}")) + else: + queue.put((index, result, None)) + + +def merge_shards(shards, samples, metrics): + order = {sample["item_id"]: index for index, sample in enumerate(samples)} + items = {metric.name: [] for metric in metrics} + for shard_items in shards: + for metric in metrics: + items[metric.name].extend(shard_items[metric.name]) + for metric in metrics: + rows = items[metric.name] + if len(rows) != len(samples) or {row.item_id for row in rows} != set(order): + raise RuntimeError(f"incomplete or duplicate shard results for {metric.name}") + rows.sort(key=lambda row: order[row.item_id]) + return items diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/__init__.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/__init__.py new file mode 100644 index 000000000..f028df95c --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/__init__.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""WorldMemBench metric interfaces and built-in registry.""" + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.base import Metric +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.camera import ( + CameraRotationMetric, + CameraTranslationMetric, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.dynamic import DynamicMemoryMetric +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.environment import ( + EnvironmentDepthFidelityMetric, + EnvironmentSceneReconstructionMetric, + EnvironmentVisibleMemoryMetric, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.general_video_quality import ( + AestheticQualityMetric, + ImagingQualityMetric, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.human import ( + HumanAppearanceMetric, + HumanIdentityMetric, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.object import ( + ObjectAppearanceMetric, + ObjectGeometryMetric, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.registry import MetricRegistry + +registry = MetricRegistry() +registry.register(DynamicMemoryMetric()) +registry.register(CameraRotationMetric()) +registry.register(CameraTranslationMetric()) +registry.register(ImagingQualityMetric()) +registry.register(AestheticQualityMetric()) +registry.register(HumanIdentityMetric()) +registry.register(HumanAppearanceMetric()) +registry.register(ObjectAppearanceMetric()) +registry.register(ObjectGeometryMetric()) +registry.register(EnvironmentVisibleMemoryMetric()) +registry.register(EnvironmentDepthFidelityMetric()) +registry.register(EnvironmentSceneReconstructionMetric()) + +__all__ = ["Metric", "MetricRegistry", "registry"] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/aggregation.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/aggregation.py new file mode 100644 index 000000000..373b08d19 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/aggregation.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Equal-weight case aggregation shared by video metrics.""" + +from collections import Counter + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.results import ( + ItemResult, + ResultStatus, + SummaryResult, +) + + +def aggregate_video_summary(items: list[ItemResult]) -> SummaryResult: + # Weight videos equally within each model and retain status/coverage counts. + groups: dict[str, list[ItemResult]] = {} + detailed_groups: dict[tuple, list[ItemResult]] = {} + for item in items: + key = str(item.model_id) if item.model_id is not None else "all" + groups.setdefault(key, []).append(item) + detailed_groups.setdefault((item.model_id, item.category, item.condition_id), []).append(item) + + def aggregate_videos(group): + scored = [item for item in group if item.status.is_scored and bool(item.scores)] + score_values: dict[str, list[float]] = {} + for item in scored: + for score_key, value in item.scores.items(): + score_values.setdefault(score_key, []).append(float(value)) + return { + "scores": {name: sum(values) / len(values) for name, values in score_values.items() if values}, + "metadata": { + "n_total": len(group), + "n_scored": len(scored), + **{f"n_{status.value}": sum(item.status == status for item in group) for status in ResultStatus}, + }, + } + + return SummaryResult( + num_items=len(items), + status_counts=dict(Counter(item.status.value for item in items)), + results={key: aggregate_videos(group) for key, group in groups.items()}, + groups=[ + {"model_id": model, "category": category, "trajectory_id": trajectory, **aggregate_videos(group)} + for (model, category, trajectory), group in detailed_groups.items() + ], + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/base.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/base.py new file mode 100644 index 000000000..c460579e4 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/base.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Common interface for metric evaluation and aggregation.""" + +from abc import ABC, abstractmethod +from collections import Counter + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.context import EvaluationContext +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.results import ( + ItemResult, + SummaryResult, + UnitResult, +) + + +class Metric(ABC): + name: str + description: str = "" + + @abstractmethod + def evaluate(self, sample: dict, context: EvaluationContext) -> list[UnitResult]: + """Evaluate one sample and return its fine-grained units.""" + raise NotImplementedError + + @abstractmethod + def aggregate_item(self, sample: dict, units: list[UnitResult]) -> ItemResult: + """Apply this metric's unit-to-item scoring rule.""" + raise NotImplementedError + + def aggregate_summary(self, items: list[ItemResult]) -> SummaryResult: + """Count statuses; metrics add their explicitly defined score aggregation.""" + return SummaryResult( + num_items=len(items), + status_counts=dict(Counter(item.status.value for item in items)), + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/__init__.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/__init__.py new file mode 100644 index 000000000..f6be86d61 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Independent Camera Following rotation and translation RPE metrics.""" + +from .metric import CameraRotationMetric, CameraTranslationMetric + +__all__ = ["CameraRotationMetric", "CameraTranslationMetric"] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/metric.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/metric.py new file mode 100644 index 000000000..5428b2a3b --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/metric.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Separate rotation and translation RPE, preserving the reference estimator policy.""" + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import ( + DataInvalidError, + EvaluationError, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.aggregation import ( + aggregate_video_summary, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.base import Metric +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.results import ( + ItemResult, + ResultStatus, + UnitResult, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.video import VideoReader + + +def required_estimators(motion_type): + return ("da3",) if motion_type == "rotation" else ("da3", "vipe") + + +def load_inputs(sample, reader): + from .pose_io import load_trajectory + from .rpe import rotation_path_length_deg + from .trajectory import canonicalize_target + + motion = sample["annotations"]["camera_motion_type"] + try: + target = load_trajectory(sample["artifacts"]["target_poses"], "target") + target, path_length = canonicalize_target(target, motion) + rotation_path = rotation_path_length_deg(target) if motion in ("rotation", "mixed") else None + if rotation_path is not None and rotation_path <= 1e-12: + raise ValueError("rotation target must have a non-zero rotation path") + except (KeyError, ValueError) as error: + raise DataInvalidError(str(error)) from error + info = reader.probe(sample["video"]) + if info.total_frames != target.num_poses: + raise EvaluationError(f"video/target frame count mismatch: {info.total_frames} != {target.num_poses}") + return target, path_length, rotation_path + + +class _CameraMetric(Metric): + motion = "" + score_name = "" + + def __init__(self, *, reader=None): + self.reader = reader or VideoReader() + + def evaluate(self, sample, context): + from .pose_io import load_trajectory + from .rpe import compute_rpe + from .schemas import PoseTrajectory + from .trajectory import umeyama_align + + if error := sample.get("camera_preparation_error"): + return [UnitResult("video", status=ResultStatus(error["status"]), error=error["message"])] + motion = (sample.get("annotations") or {}).get("camera_motion_type") + if motion not in (self.motion, "mixed"): + return [UnitResult("video", status=ResultStatus.NOT_APPLICABLE)] + try: + target, path_length, rotation_path = load_inputs(sample, self.reader) + except DataInvalidError: + raise + except Exception as error: + return [UnitResult("video", status=ResultStatus.MODEL_FAILURE, error=str(error))] + + def analyze(): + results = {} + for name in required_estimators(motion): + preparation_error = (sample.get("camera_artifact_errors") or {}).get(name) + try: + if preparation_error: + raise EvaluationError(preparation_error) + trajectory = load_trajectory(sample["artifacts"][name], name) + if trajectory.num_poses != target.num_poses: + raise EvaluationError( + f"{name} pose/video count mismatch: {trajectory.num_poses} != {target.num_poses}" + ) + scale = None + if motion != "rotation": + alignment = umeyama_align(trajectory, target) + scale = alignment.scale + trajectory = PoseTrajectory(alignment.poses, source=name) + statistics = compute_rpe(target, trajectory, motion) + results[name] = { + "status": "ok", + "num_poses": trajectory.num_poses, + "alignment_scale": scale, + "rotation_rpe_mean_deg": statistics.rotation_mean_deg, + "rotation_rpe_path": statistics.rotation_path(rotation_path) + if rotation_path is not None + else None, + "translation_rpe_path": statistics.translation_path, + } + except Exception as error: + results[name] = {"status": "eval_failed", "error": f"{type(error).__name__}: {error}"} + return results + + results = context.get_cached(("camera:rpe", sample["item_id"]), analyze) + metadata = { + "camera_motion_type": motion, + "num_frames": target.num_poses, + "num_rpe_steps": target.num_poses - 1, + "target_path_length": path_length, + "target_rotation_path_deg": rotation_path, + "estimators": results, + "direction": "lower_is_better", + } + failures = [value["error"] for value in results.values() if value["status"] != "ok"] + if failures: + return [UnitResult("video", status=ResultStatus.EVAL_FAILED, error="; ".join(failures), metadata=metadata)] + value = ( + results["da3"]["rotation_rpe_path"] + if self.motion == "rotation" + else (results["da3"]["translation_rpe_path"] + results["vipe"]["translation_rpe_path"]) / 2.0 + ) + return [UnitResult("video", unit_type="video", scores={self.score_name: value}, metadata=metadata)] + + def aggregate_item(self, sample, units): + unit = units[0] + return ItemResult( + sample["item_id"], + scores=dict(unit.scores), + num_units=1, + status=unit.status, + error=unit.error, + metadata=dict(unit.metadata), + ) + + def aggregate_summary(self, items): + summary = aggregate_video_summary(items) + summary.metadata["direction"] = "lower_is_better" + return summary + + +class CameraRotationMetric(_CameraMetric): + name = "camera_rotation" + motion = "rotation" + score_name = "camera_rotation_rpe_path" + description = "DA3 adjacent rotation RPE normalized by target rotation path; lower is better." + + +class CameraTranslationMetric(_CameraMetric): + name = "camera_translation" + motion = "translation" + score_name = "camera_translation_rpe" + description = "Mean DA3/ViPE path-normalized translation RPE after Sim(3) alignment; lower is better." diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/pose_io.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/pose_io.py new file mode 100644 index 000000000..e6dd09ad2 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/pose_io.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Read full-frame target c2w, DA3 w2c, and ViPE c2w pose arrays.""" + +from pathlib import Path + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import ( + DataInvalidError, + EvaluationError, +) + +from .schemas import PoseTrajectory + + +def load_trajectory(path, kind: str) -> PoseTrajectory: + error_type = DataInvalidError if kind == "target" else EvaluationError + try: + path = Path(path).expanduser().resolve() + if kind == "target" and path.suffix == ".npy": + poses = np.load(path, allow_pickle=False) + else: + with np.load(path, allow_pickle=False) as data: + poses = np.asarray(data["extrinsics" if kind == "da3" else "c2w"], dtype=np.float64) + if "frame_indices" in data and not np.array_equal(data["frame_indices"], np.arange(len(poses))): + raise ValueError("frame_indices must cover every frame in exact 0..N-1 order") + if kind == "da3": + matrices = np.asarray(poses, dtype=np.float64) + if matrices.ndim == 3 and matrices.shape[1:] == (3, 4): + padded = np.repeat(np.eye(4)[None], len(matrices), axis=0) + padded[:, :3] = matrices + matrices = padded + if matrices.ndim != 3 or matrices.shape[1:] != (4, 4): + raise ValueError("DA3 extrinsics must have shape (N,3,4) or (N,4,4)") + poses = np.linalg.inv(matrices) + return PoseTrajectory(poses, source=kind) + except (OSError, ValueError, KeyError, TypeError, np.linalg.LinAlgError) as error: + raise error_type(f"invalid {kind} pose artifact {path}: {error}") from error diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/preparation.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/preparation.py new file mode 100644 index 000000000..5d03beea5 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/preparation.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Read target camera inputs and prepare required predicted poses before scoring.""" + +import json +from pathlib import Path + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.video import VideoReader + +from .metric import load_inputs, required_estimators + + +def prepare_for_evaluation(samples, *, required=False): + for sample in samples: + sample.pop("camera_preparation_error", None) + try: + directory = Path(sample["video"]).parent + annotations = dict(sample.get("annotations") or {}) + artifacts = dict(sample.get("artifacts") or {}) + if "camera_motion_type" not in annotations: + source = directory / "trajectory.json" + if source.is_file(): + value = json.loads(source.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise DataInvalidError("trajectory.json must be an object") + if "camera_motion_type" in value: + annotations["camera_motion_type"] = value["camera_motion_type"] + if "camera_motion_type" not in annotations: + if required or artifacts.get("target_poses") or (directory / "target_poses.npz").exists(): + raise DataInvalidError("Camera evaluation requires camera_motion_type in trajectory.json") + continue + if annotations["camera_motion_type"] not in ("rotation", "translation", "mixed"): + raise DataInvalidError("camera_motion_type must be rotation, translation or mixed") + artifacts.setdefault("target_poses", str(directory / "target_poses.npz")) + artifacts.setdefault("da3", str(directory / "da3/results.npz")) + if annotations["camera_motion_type"] != "rotation": + artifacts.setdefault("vipe", str(directory / "vipe/results.npz")) + sample.update(annotations=annotations, artifacts=artifacts) + except (DataInvalidError, OSError, ValueError, KeyError, TypeError) as error: + sample["camera_preparation_error"] = {"status": "data_invalid", "message": str(error)} + + +def prepare_artifacts(samples, context, metric_names): + eligible = [ + sample + for sample in samples + if not sample.get("camera_preparation_error") + and ( + (sample.get("annotations") or {}).get("camera_motion_type") == "mixed" + or "camera_" + str((sample.get("annotations") or {}).get("camera_motion_type")) in metric_names + ) + ] + if not eligible: + return + from ...preparation.da3 import ensure_artifacts as da3 + from ...preparation.vipe import ensure_artifacts as vipe + + reader = VideoReader() + print(f"[{context.device}][Camera poses] 0/{len(eligible)}", flush=True) + for index, sample in enumerate(eligible, 1): + try: + load_inputs(sample, reader) + except Exception as error: + sample["camera_preparation_error"] = { + "status": "data_invalid" if isinstance(error, DataInvalidError) else "model_failure", + "message": str(error), + } + else: + motion = sample["annotations"]["camera_motion_type"] + sample["camera_artifact_errors"] = {} + for name in required_estimators(motion): + if Path(sample["artifacts"][name]).is_file(): + continue + try: + if name == "da3": + da3( + sample, + context, + weight_metric="camera_rotation" if motion == "rotation" else "camera_translation", + require_scene=False, + ) + else: + vipe(sample, context) + except Exception as error: + sample["camera_artifact_errors"][name] = f"{type(error).__name__}: {error}" + failed = sample.get("camera_preparation_error") or sample.get("camera_artifact_errors") + print( + f"[{context.device}][Camera poses] {index}/{len(eligible)} " + f"{'failed' if failed else 'ready'}: {sample['item_id']}", + flush=True, + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/rpe.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/rpe.py new file mode 100644 index 000000000..4d881bffc --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/rpe.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Adjacent-step translation and rotation Relative Pose Error.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from .schemas import PoseTrajectory +from .trajectory import MOTION_TYPES + + +def _rotation_errors_deg(error_rotations: np.ndarray) -> np.ndarray: + traces = np.trace(error_rotations, axis1=1, axis2=2) + cosines = np.clip((traces - 1.0) / 2.0, -1.0, 1.0) + return np.degrees(np.arccos(cosines)) + + +@dataclass(frozen=True) +class RPEStatistics: + translation_errors: np.ndarray | None + rotation_errors_deg: np.ndarray + + @property + def translation_mean(self) -> float | None: + if self.translation_errors is None: + return None + return float(np.mean(self.translation_errors)) + + @property + def translation_median(self) -> float | None: + if self.translation_errors is None: + return None + return float(np.median(self.translation_errors)) + + @property + def translation_path(self) -> float | None: + """Sum of adjacent errors on the unit-path canonical target.""" + if self.translation_errors is None: + return None + return float(np.sum(self.translation_errors)) + + @property + def rotation_mean_deg(self) -> float: + return float(np.mean(self.rotation_errors_deg)) + + @property + def rotation_median_deg(self) -> float: + return float(np.median(self.rotation_errors_deg)) + + @property + def rotation_sum_deg(self) -> float: + return float(np.sum(self.rotation_errors_deg)) + + def rotation_path(self, target_rotation_path_deg: float) -> float: + """Normalize the accumulated rotation error by the target path.""" + + path = float(target_rotation_path_deg) + if not np.isfinite(path) or path <= 1e-12: + raise ValueError("target rotation path must be positive and finite") + return self.rotation_sum_deg / path + + +def rotation_path_length_deg(trajectory: PoseTrajectory) -> float: + """Return the exact sum of adjacent target rotation magnitudes.""" + + rotations = trajectory.poses[:, :3, :3] + relative = np.einsum( + "nij,njk->nik", + np.swapaxes(rotations[:-1], 1, 2), + rotations[1:], + ) + path = float(np.sum(_rotation_errors_deg(relative))) + if not np.isfinite(path): + raise ValueError("target rotation path must be finite") + return path + + +def compute_rpe( + target: PoseTrajectory, + estimated: PoseTrajectory, + motion_type: str, +) -> RPEStatistics: + if motion_type not in MOTION_TYPES: + raise ValueError(f"unsupported camera_motion_type: {motion_type!r}") + if target.num_poses != estimated.num_poses: + raise ValueError("target and estimated trajectories must have equal length") + + target_poses = target.poses + estimated_poses = estimated.poses + if motion_type == "rotation": + target_relative = np.einsum( + "nij,njk->nik", + np.swapaxes(target_poses[:-1, :3, :3], 1, 2), + target_poses[1:, :3, :3], + ) + estimated_relative = np.einsum( + "nij,njk->nik", + np.swapaxes(estimated_poses[:-1, :3, :3], 1, 2), + estimated_poses[1:, :3, :3], + ) + rotation_error = np.einsum( + "nij,njk->nik", + np.swapaxes(target_relative, 1, 2), + estimated_relative, + ) + translation_errors = None + else: + target_relative = np.linalg.inv(target_poses[:-1]) @ target_poses[1:] + estimated_relative = np.linalg.inv(estimated_poses[:-1]) @ estimated_poses[1:] + error_transform = np.linalg.inv(target_relative) @ estimated_relative + translation_errors = np.linalg.norm(error_transform[:, :3, 3], axis=1) + rotation_error = error_transform[:, :3, :3] + + rotation_errors = _rotation_errors_deg(rotation_error) + if translation_errors is not None and not np.isfinite(translation_errors).all(): + raise ValueError("translation RPE contains non-finite values") + if not np.isfinite(rotation_errors).all(): + raise ValueError("rotation RPE contains non-finite values") + return RPEStatistics( + translation_errors=translation_errors, + rotation_errors_deg=rotation_errors, + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/schemas.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/schemas.py new file mode 100644 index 000000000..65bfee72b --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/schemas.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Validated trajectory value objects for Camera Following.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +MIN_CAMERA_POSES = 9 + + +@dataclass(frozen=True) +class PoseTrajectory: + poses: np.ndarray + source: str = "" + + def __post_init__(self) -> None: + poses = np.asarray(self.poses, dtype=np.float64) + if poses.ndim != 3 or poses.shape[1:] not in {(3, 4), (4, 4)}: + raise ValueError(f"pose trajectory must have shape (N,3,4) or (N,4,4), got {poses.shape}") + if len(poses) < MIN_CAMERA_POSES: + raise ValueError(f"pose trajectory requires at least {MIN_CAMERA_POSES} poses; got {len(poses)}") + if poses.shape[1:] == (3, 4): + padded = np.repeat(np.eye(4, dtype=np.float64)[None], len(poses), axis=0) + padded[:, :3, :] = poses + poses = padded + if not np.isfinite(poses).all(): + raise ValueError("pose trajectory must contain only finite values") + expected_bottom = np.array([0.0, 0.0, 0.0, 1.0]) + if not np.allclose(poses[:, 3, :], expected_bottom, atol=1e-6): + raise ValueError("pose trajectory bottom rows must be [0,0,0,1]") + rotations = poses[:, :3, :3] + gram = np.swapaxes(rotations, 1, 2) @ rotations + determinants = np.linalg.det(rotations) + if not np.allclose(gram, np.eye(3), atol=5e-3) or not np.allclose( + determinants, + 1.0, + atol=5e-3, + ): + raise ValueError("pose trajectory contains an invalid rotation matrix") + + object.__setattr__(self, "poses", np.ascontiguousarray(poses)) + + @property + def num_poses(self) -> int: + return len(self.poses) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/trajectory.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/trajectory.py new file mode 100644 index 000000000..714f6fbf9 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/camera/trajectory.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Reference path normalization and Sim(3) alignment for full-frame trajectories.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from .schemas import PoseTrajectory + +MOTION_TYPES = frozenset({"translation", "rotation", "mixed"}) + + +def canonicalize_target( + trajectory: PoseTrajectory, + motion_type: str, + *, + epsilon: float = 1e-12, +) -> tuple[PoseTrajectory, float | None]: + if motion_type not in MOTION_TYPES: + raise ValueError(f"camera_motion_type must be one of {sorted(MOTION_TYPES)}; got {motion_type!r}") + if motion_type == "rotation": + return trajectory, None + + centers = trajectory.poses[:, :3, 3] + path_length = float(np.linalg.norm(np.diff(centers, axis=0), axis=1).sum()) + if not np.isfinite(path_length) or path_length <= epsilon: + raise ValueError("translation/mixed target must have non-zero finite total path length") + poses = trajectory.poses.copy() + poses[:, :3, 3] = (centers - centers[0]) / path_length + return ( + PoseTrajectory( + poses, + source=f"{trajectory.source}:canonical", + ), + path_length, + ) + + +@dataclass(frozen=True) +class UmeyamaAlignment: + scale: float + rotation: np.ndarray + translation: np.ndarray + poses: np.ndarray + + +def umeyama_align( + estimated: PoseTrajectory, + target: PoseTrajectory, + *, + epsilon: float = 1e-12, +) -> UmeyamaAlignment: + if estimated.num_poses != target.num_poses: + raise ValueError("estimated and target trajectories must have equal length") + source = estimated.poses[:, :3, 3] + destination = target.poses[:, :3, 3] + source_mean = source.mean(axis=0) + destination_mean = destination.mean(axis=0) + source_centered = source - source_mean + destination_centered = destination - destination_mean + source_variance = float(np.square(source_centered).sum() / len(source)) + + if source_variance <= epsilon: + alignment_rotation = np.eye(3, dtype=np.float64) + scale = 0.0 + else: + covariance = destination_centered.T @ source_centered / len(source) + left, singular_values, right_transpose = np.linalg.svd(covariance) + reflection = np.ones(3, dtype=np.float64) + if np.linalg.det(left @ right_transpose) < 0: + reflection[-1] = -1.0 + reflection_guard = np.diag(reflection) + alignment_rotation = left @ reflection_guard @ right_transpose + scale = float((singular_values * reflection).sum() / (source_variance + epsilon)) + + translation = destination_mean - scale * alignment_rotation @ source_mean + aligned = estimated.poses.copy() + aligned[:, :3, :3] = np.einsum( + "ij,njk->nik", + alignment_rotation, + estimated.poses[:, :3, :3], + ) + aligned[:, :3, 3] = scale * (alignment_rotation @ source.T).T + translation + return UmeyamaAlignment( + scale=scale, + rotation=alignment_rotation, + translation=translation, + poses=aligned, + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/__init__.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/__init__.py new file mode 100644 index 000000000..65874ceb5 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +from .metric import DynamicMemoryMetric + +__all__ = ["DynamicMemoryMetric"] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/checklist.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/checklist.py new file mode 100644 index 000000000..09eb95263 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/checklist.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Validate one frozen eight-question Dynamic Memory checklist.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Any, Literal + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError + +CATEGORIES = ( + "return_progress", + "persistence_reset", + "final_state", + "causal_physics", +) +EXPECTED_IDS = tuple(f"{category}_{index:02d}" for category in CATEGORIES for index in (1, 2)) +EVIDENCE_WINDOWS = { + "returned_view", + "returned_hold_sequence", + "returned_and_later_visible_sequence", + "final_stationary_view", + "ordered_visible_sequence", +} + + +@dataclass(frozen=True) +class ChecklistQuestion: + id: str + category: str + question: str + expected_answer: Literal["yes", "no"] + evidence_window: str + + +@dataclass(frozen=True) +class CaseChecklist: + case_id: str + questions: tuple[ChecklistQuestion, ...] + + +def parse_checklist(value: Any) -> CaseChecklist: + if not isinstance(value, dict): + raise DataInvalidError("checklist case must be an object") + case_id = str(value.get("case_id", "")) + if not re.fullmatch(r"case_[0-9]{3}", case_id): + raise DataInvalidError(f"invalid checklist case_id: {case_id!r}") + raw_questions = value.get("questions") + if not isinstance(raw_questions, list) or len(raw_questions) != 8: + raise DataInvalidError(f"{case_id}: checklist must contain 8 questions") + questions: list[ChecklistQuestion] = [] + for index, raw in enumerate(raw_questions): + if not isinstance(raw, dict): + raise DataInvalidError(f"{case_id}: question {index} is not an object") + question_id = str(raw.get("id", "")) + expected_id = EXPECTED_IDS[index] + category = str(raw.get("category", "")) + if question_id != expected_id or category != expected_id.rsplit("_", 1)[0]: + raise DataInvalidError(f"{case_id}: invalid question layout at {index}") + text = str(raw.get("question", "")).strip() + if not text.endswith("?"): + raise DataInvalidError(f"{case_id}/{question_id}: invalid question text") + expected_answer = str(raw.get("expected_answer", "")).lower() + if expected_answer not in {"yes", "no"}: + raise DataInvalidError(f"{case_id}/{question_id}: expected_answer must be yes or no") + evidence_window = str(raw.get("evidence_window", "")).strip() + if evidence_window not in EVIDENCE_WINDOWS: + raise DataInvalidError(f"{case_id}/{question_id}: invalid evidence_window {evidence_window!r}") + questions.append( + ChecklistQuestion( + id=question_id, + category=category, + question=text, + expected_answer=expected_answer, # type: ignore[arg-type] + evidence_window=evidence_window, + ) + ) + return CaseChecklist(case_id=case_id, questions=tuple(questions)) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/metric.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/metric.py new file mode 100644 index 000000000..5efa31934 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/metric.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Run VLM answer generation and score calculation through one evaluation interface.""" + +from dataclasses import asdict + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.aggregation import ( + aggregate_video_summary, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.base import Metric +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.results import ( + ItemResult, + ResultStatus, + UnitResult, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.video import VideoReader + +from .checklist import parse_checklist +from .scoring import aggregate_scores, score_answers +from .vlm import HTTPVLMBackend, InvalidVLMResponse, VLMConfig, generate_answers + + +class DynamicMemoryMetric(Metric): + name = "dynamic_memory" + description = "VLM answers to eight visual questions, scored against case-specific targets." + + def __init__(self, *, backend=None, reader=None): + self.backend = backend + self.reader = reader or VideoReader() + + def evaluate(self, sample, context): + if sample.get("category") not in (None, "dynamic_memory"): + return [UnitResult("checklist", status=ResultStatus.NOT_APPLICABLE)] + if error := sample.get("dynamic_preparation_error"): + raise DataInvalidError(error["message"]) + config = VLMConfig(**context.metric_config(self.name)) + annotations = sample["annotations"] + checklist = parse_checklist(annotations["dynamic_checklist"]) + backend = self.backend or context.get_resource("dynamic:vlm", HTTPVLMBackend) + try: + response, indices, protocol, anchors = generate_answers( + sample["video"], + checklist, + annotations["lr_hold_protocol"], + config, + reader=self.reader, + backend=backend, + ) + except InvalidVLMResponse as error: + return [ + UnitResult( + "checklist", + status=ResultStatus.EVAL_FAILED, + error=str(error), + metadata={"raw_vlm_response": error.raw_response}, + ) + ] + units = score_answers(checklist, response, indices) + event = response.event_started + units[0].metadata["case_metadata"] = { + "case_id": checklist.case_id, + "vlm_model_id": config.model, + "vlm_backend": config.backend, + "vlm_settings": {key: value for key, value in asdict(config).items() if key != "api_key_env"}, + "sampled_frame_indices": indices, + "event_started": event.answer, + "event_gate_applied": event.answer == "no", + "event_evidence_image_numbers": list(event.evidence_image_numbers), + "event_evidence_frame_indices": [indices[n - 1] for n in event.evidence_image_numbers], + "event_reason": event.reason, + "raw_vlm_response": response.raw_response, + "lr_hold_protocol": protocol.as_dict(), + "prompt_anchors": anchors.as_dict(), + "usage": response.usage, + } + return units + + def aggregate_item(self, sample, units): + if units[0].status != ResultStatus.OK: + unit = units[0] + return ItemResult(sample["item_id"], status=unit.status, error=unit.error, metadata=dict(unit.metadata)) + scores = aggregate_scores(units) + metadata = dict(units[0].metadata["case_metadata"]) + metadata["questions"] = [ + { + "id": unit.unit_id, + "score": unit.scores["score"], + **{key: value for key, value in unit.metadata.items() if key != "case_metadata"}, + } + for unit in units + ] + return ItemResult(sample["item_id"], scores=scores, num_units=8, metadata=metadata) + + def aggregate_summary(self, items): + return aggregate_video_summary(items) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/preparation.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/preparation.py new file mode 100644 index 000000000..da9af8caf --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/preparation.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Embed per-case checklist and planned LR-hold phases in the run manifest.""" + +import json +from dataclasses import asdict +from pathlib import Path + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import ( + DataInvalidError, + EvaluationError, +) + +from .checklist import parse_checklist +from .protocol import LR_HOLD_TRAJECTORY_ID, LRHoldProtocol + + +def prepare_for_evaluation(samples: list[dict]) -> None: + for sample in samples: + if sample.get("category") not in (None, "dynamic_memory"): + continue + sample.pop("dynamic_preparation_error", None) + try: + directory = Path(sample["video"]).parent + annotations = dict(sample.get("annotations") or {}) + if "dynamic_checklist" not in annotations: + annotations["dynamic_checklist"] = json.loads( + (directory / "checklist.json").read_text(encoding="utf-8") + ) + case = parse_checklist(annotations["dynamic_checklist"]) + if sample.get("sample_id") not in (None, case.case_id): + raise DataInvalidError("checklist case_id does not match the video case") + annotations["dynamic_checklist"] = { + "case_id": case.case_id, + "questions": [asdict(question) for question in case.questions], + } + if "lr_hold_protocol" not in annotations: + trajectory = json.loads((directory / "trajectory.json").read_text(encoding="utf-8")) + annotations["lr_hold_protocol"] = { + "trajectory_id": LR_HOLD_TRAJECTORY_ID, + "native_trajectory_id": sample.get("condition_id") or LR_HOLD_TRAJECTORY_ID, + **trajectory, + } + protocol = annotations["lr_hold_protocol"] + # The declared length belongs to the planned trajectory, not to the decoded video. + protocol = { + "trajectory_id": protocol["trajectory_id"], + "native_trajectory_id": protocol["native_trajectory_id"], + "actions": protocol["actions"], + "phase_transitions": protocol["phase_transitions"], + "tail_hold_transitions": protocol.get("tail_hold_transitions", 0), + "declared_output_frames": protocol.get( + "declared_output_frames", + sum(protocol["phase_transitions"]) + protocol.get("tail_hold_transitions", 0) + 1, + ), + } + parsed = LRHoldProtocol.from_metadata(protocol, actual_output_frames=protocol["declared_output_frames"]) + if sample.get("condition_id") not in (None, parsed.native_trajectory_id): + raise DataInvalidError("condition_id does not match LR-hold native trajectory") + annotations["lr_hold_protocol"] = protocol + sample["annotations"] = annotations + except (DataInvalidError, EvaluationError, OSError, ValueError, KeyError, TypeError, AttributeError) as error: + sample["dynamic_preparation_error"] = {"status": "data_invalid", "message": str(error)} diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/protocol.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/protocol.py new file mode 100644 index 000000000..71467fd60 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/protocol.py @@ -0,0 +1,227 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Planned five-phase protocol for the single-revisit LR-hold trajectory.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Iterable + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import ( + DataInvalidError, + EvaluationError, +) + +LR_HOLD_TRAJECTORY_ID = "lr_hold_sym_0p5dpf" +LR_HOLD_ACTIONS = ("hold", "yaw_left", "hold", "yaw_right", "hold") + + +@dataclass(frozen=True) +class FrameRange: + """Inclusive output-frame range.""" + + start: int + end: int + + def __post_init__(self) -> None: + if self.start < 0 or self.end < self.start: + raise ValueError(f"invalid inclusive frame range: {self.start}..{self.end}") + + def contains(self, frame_index: int) -> bool: + return self.start <= frame_index <= self.end + + def as_dict(self) -> dict[str, int]: + return {"start": self.start, "end": self.end} + + +@dataclass(frozen=True) +class PromptAnchors: + """Sampled-image anchors supplied to the VLM rather than inferred by it.""" + + frame_map: tuple[dict[str, Any], ...] + departure_baseline_images: tuple[int, ...] + away_hold_images: tuple[int, ...] + returned_hold_images: tuple[int, ...] + returned_before_final_images: tuple[int, ...] + final_stationary_images: tuple[int, ...] + + def as_dict(self) -> dict[str, list[int]]: + return { + "departure_baseline_images": list(self.departure_baseline_images), + "away_hold_images": list(self.away_hold_images), + "returned_hold_images": list(self.returned_hold_images), + "returned_before_final_images": list(self.returned_before_final_images), + "final_stationary_images": list(self.final_stationary_images), + } + + def phase_ranges(self) -> dict[str, list[int]]: + ranges: dict[str, list[int]] = {} + for phase in ( + "initial_hold", + "yaw_left", + "away_hold", + "yaw_right", + "returned_hold", + ): + numbers = [int(value["image_number"]) for value in self.frame_map if value["phase"] == phase] + if numbers: + ranges[phase] = [numbers[0], numbers[-1]] + return ranges + + def compact_ranges(self) -> dict[str, list[int]]: + def bounds(numbers: tuple[int, ...]) -> list[int]: + return [numbers[0], numbers[-1]] + + return { + "departure_baseline": bounds(self.departure_baseline_images), + "away_hold": bounds(self.away_hold_images), + "returned_hold": bounds(self.returned_hold_images), + "returned_before_final": bounds(self.returned_before_final_images), + "final_stationary": bounds(self.final_stationary_images), + } + + +@dataclass(frozen=True) +class LRHoldProtocol: + """Model-native planned LR-hold phases mapped onto one output video.""" + + trajectory_id: str + native_trajectory_id: str + phase_transitions: tuple[int, int, int, int, int] + tail_hold_transitions: int + declared_output_frames: int + actual_output_frames: int + initial_hold: FrameRange + yaw_left: FrameRange + away_hold: FrameRange + yaw_right: FrameRange + returned_hold: FrameRange + + @classmethod + def from_metadata( + cls, + value: Any, + *, + actual_output_frames: int, + ) -> "LRHoldProtocol": + if not isinstance(value, dict): + raise DataInvalidError("metadata.lr_hold_protocol must be an object") + trajectory_id = str(value.get("trajectory_id", "")) + if trajectory_id != LR_HOLD_TRAJECTORY_ID: + raise DataInvalidError(f"Dynamic Memory requires {LR_HOLD_TRAJECTORY_ID}, got {trajectory_id!r}") + native_trajectory_id = str(value.get("native_trajectory_id", "")) + if not ( + native_trajectory_id == LR_HOLD_TRAJECTORY_ID + or native_trajectory_id.startswith(LR_HOLD_TRAJECTORY_ID + "_action") + ): + raise DataInvalidError(f"invalid native LR-hold trajectory: {native_trajectory_id!r}") + raw_actions = value.get("actions") + if tuple(raw_actions or ()) != LR_HOLD_ACTIONS: + raise DataInvalidError(f"LR-hold actions must be {LR_HOLD_ACTIONS}, got {raw_actions!r}") + raw_counts = value.get("phase_transitions") + if ( + not isinstance(raw_counts, list) + or len(raw_counts) != 5 + or any(isinstance(count, bool) or not isinstance(count, int) for count in raw_counts) + or any(count <= 0 for count in raw_counts) + ): + raise DataInvalidError("lr_hold_protocol.phase_transitions must contain five positive integers") + counts = tuple(raw_counts) + tail_hold = value.get("tail_hold_transitions", 0) + declared_frames = value.get("declared_output_frames") + if isinstance(tail_hold, bool) or not isinstance(tail_hold, int) or tail_hold < 0: + raise DataInvalidError("tail_hold_transitions must be a non-negative integer") + if isinstance(declared_frames, bool) or not isinstance(declared_frames, int) or declared_frames < 2: + raise DataInvalidError("declared_output_frames must be an integer >= 2") + if actual_output_frames < 2: + raise EvaluationError("LR-hold video must contain at least two frames") + if sum(counts) + tail_hold != declared_frames - 1: + raise DataInvalidError("phase transitions and tail hold do not cover declared output frames") + + native_boundaries = [] + cumulative = 0 + for count in counts[:4]: + cumulative += count + native_boundaries.append(cumulative) + + def map_frame(index: int) -> int: + return round(index * (actual_output_frames - 1) / (declared_frames - 1)) + + b1, b2, b3, b4 = (map_frame(value) for value in native_boundaries) + last = actual_output_frames - 1 + if not (0 < b1 < b2 < b3 < b4 < last): + raise EvaluationError( + f"mapped LR-hold phase boundaries are not strictly increasing: {(b1, b2, b3, b4, last)}" + ) + return cls( + trajectory_id=trajectory_id, + native_trajectory_id=native_trajectory_id, + phase_transitions=counts, # type: ignore[arg-type] + tail_hold_transitions=tail_hold, + declared_output_frames=declared_frames, + actual_output_frames=actual_output_frames, + initial_hold=FrameRange(0, b1), + yaw_left=FrameRange(b1 + 1, b2 - 1), + away_hold=FrameRange(b2, b3), + yaw_right=FrameRange(b3 + 1, b4 - 1), + returned_hold=FrameRange(b4, last), + ) + + def phase_for_frame(self, frame_index: int) -> str: + if not 0 <= frame_index < self.actual_output_frames: + raise ValueError(f"frame {frame_index} outside [0, {self.actual_output_frames - 1}]") + if self.initial_hold.contains(frame_index): + return "initial_hold" + if self.yaw_left.contains(frame_index): + return "yaw_left" + if self.away_hold.contains(frame_index): + return "away_hold" + if self.yaw_right.contains(frame_index): + return "yaw_right" + return "returned_hold" + + def prompt_anchors(self, sampled_frame_indices: Iterable[int]) -> PromptAnchors: + indices = tuple(int(index) for index in sampled_frame_indices) + if not indices or tuple(sorted(set(indices))) != indices: + raise ValueError("sampled frame indices must be sorted and unique") + frame_map = tuple( + { + "image_number": image_number, + "source_frame_index": frame_index, + "phase": self.phase_for_frame(frame_index), + } + for image_number, frame_index in enumerate(indices, start=1) + ) + + def image_numbers(phase: str) -> tuple[int, ...]: + return tuple(int(value["image_number"]) for value in frame_map if value["phase"] == phase) + + initial = image_numbers("initial_hold") + away = image_numbers("away_hold") + returned = image_numbers("returned_hold") + if len(initial) < 3 or len(away) < 1 or len(returned) < 4: + raise EvaluationError("stride sampling does not preserve enough LR-hold anchor images") + return PromptAnchors( + frame_map=frame_map, + departure_baseline_images=initial[-3:], + away_hold_images=away, + returned_hold_images=returned, + returned_before_final_images=returned[:-3], + final_stationary_images=returned[-3:], + ) + + def as_dict(self) -> dict[str, Any]: + return { + "trajectory_id": self.trajectory_id, + "native_trajectory_id": self.native_trajectory_id, + "phase_transitions": list(self.phase_transitions), + "tail_hold_transitions": self.tail_hold_transitions, + "declared_output_frames": self.declared_output_frames, + "actual_output_frames": self.actual_output_frames, + "initial_hold": self.initial_hold.as_dict(), + "yaw_left": self.yaw_left.as_dict(), + "away_hold": self.away_hold.as_dict(), + "yaw_right": self.yaw_right.as_dict(), + "returned_hold": self.returned_hold.as_dict(), + } diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/scoring.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/scoring.py new file mode 100644 index 000000000..5985adcda --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/scoring.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Calculate Dynamic Memory scores from structured answers, without model inference.""" + +from statistics import mean + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.results import UnitResult + +from .checklist import CATEGORIES + + +def score_answers(checklist, response, frame_indices) -> list[UnitResult]: + """Stage two: compare each answer with its target and apply the event prerequisite.""" + event_started = response.event_started.answer == "yes" + units = [] + for question, answer in zip(checklist.questions, response.answers, strict=True): + ungated = float(answer.answer == question.expected_answer) + units.append( + UnitResult( + question.id, + unit_type="question", + scores={"score": ungated if event_started else 0.0}, + metadata={ + "category": question.category, + "question": question.question, + "expected_answer": question.expected_answer, + "predicted_answer": answer.answer, + "evidence_window": question.evidence_window, + "reason": answer.reason, + "evidence_image_numbers": list(answer.evidence_image_numbers), + "evidence_frame_indices": [frame_indices[n - 1] for n in answer.evidence_image_numbers], + "ungated_score": ungated, + }, + ) + ) + return units + + +def aggregate_scores(units) -> dict[str, float]: + if len(units) != 8: + raise EvaluationError("Dynamic Memory requires 8 question units") + scores = {"score": mean(unit.scores["score"] for unit in units)} + for category in CATEGORIES: + scores[f"{category}_score"] = mean( + unit.scores["score"] for unit in units if unit.metadata["category"] == category + ) + return scores diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/vlm.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/vlm.py new file mode 100644 index 000000000..55e6ff4ff --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/dynamic/vlm.py @@ -0,0 +1,396 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Generate structured VLM answers from chronologically sampled video frames.""" + +from __future__ import annotations + +import base64 +import json +import os +import urllib.error +import urllib.parse +import urllib.request +from dataclasses import dataclass, field, replace +from typing import Literal + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + +from .checklist import CaseChecklist +from .protocol import LRHoldProtocol, PromptAnchors + +DEFAULT_STRIDE = 8 +DEFAULT_JPEG_QUALITY = 90 + + +@dataclass(frozen=True) +class VLMConfig: + base_url: str = "http://127.0.0.1:8001/v1" + model: str = "Qwen3.8-27B" + max_tokens: int = 4096 + timeout_seconds: float = 600.0 + jpeg_quality: int = DEFAULT_JPEG_QUALITY + sampling_stride_frames: int = DEFAULT_STRIDE + backend: str = "local" + api_key_env: str = "WORLDMEMBENCH_API_KEY" + + +@dataclass(frozen=True) +class QuestionAnswer: + id: str + answer: Literal["yes", "no"] + evidence_image_numbers: tuple[int, ...] + reason: str + + +@dataclass(frozen=True) +class EventStartedAnswer: + answer: Literal["yes", "no"] + evidence_image_numbers: tuple[int, ...] + reason: str + + +@dataclass(frozen=True) +class VLMResponse: + raw_response: str + event_started: EventStartedAnswer + answers: tuple[QuestionAnswer, ...] + usage: dict = field(default_factory=dict) + + +class InvalidVLMResponse(EvaluationError): + def __init__(self, message: str, *, raw_response: str | None = None) -> None: + super().__init__(message) + self.raw_response = raw_response + + +def sample_indices(total_frames: int, *, stride: int = DEFAULT_STRIDE) -> list[int]: + if total_frames <= 0: + return [] + if stride <= 0: + raise ValueError("sampling stride must be positive") + indices = list(range(0, total_frames, stride)) + if indices[-1] != total_frames - 1: + indices.append(total_frames - 1) + return indices + + +def encode_jpegs( + frames_rgb: list, + *, + quality: int = DEFAULT_JPEG_QUALITY, +) -> list[bytes]: + import cv2 + + encoded: list[bytes] = [] + for frame in frames_rgb: + ok, buffer = cv2.imencode( + ".jpg", + cv2.cvtColor(frame, cv2.COLOR_RGB2BGR), + [int(cv2.IMWRITE_JPEG_QUALITY), quality], + ) + if not ok: + raise EvaluationError("failed to JPEG-encode a sampled frame") + encoded.append(bytes(buffer)) + return encoded + + +def generate_answers(video_path, checklist, trajectory, config, *, reader, backend): + """Stage one: sample the video and obtain answers in one VLM request.""" + if config.backend not in ("local", "api") or not config.base_url or not config.model: + raise EvaluationError("VLM requires backend local/api, base_url and model") + if config.max_tokens <= 0 or config.timeout_seconds <= 0 or not 1 <= config.jpeg_quality <= 100: + raise EvaluationError("invalid VLM token, timeout or JPEG quality setting") + info = reader.probe(video_path) + indices = sample_indices(info.total_frames, stride=config.sampling_stride_frames) + protocol = LRHoldProtocol.from_metadata(trajectory, actual_output_frames=info.total_frames) + anchors = protocol.prompt_anchors(indices) + images = encode_jpegs( + [frame.rgb for frame in reader.read_indices(video_path, indices)], quality=config.jpeg_quality + ) + initial = anchors.phase_ranges()["initial_hold"] + response = backend.generate( + prompt=build_prompt(checklist, anchors), + jpeg_images=images, + config=config, + expected_ids=tuple(question.id for question in checklist.questions), + event_image_numbers=tuple(range(initial[0], initial[1] + 1)), + ) + return response, indices, protocol, anchors + + +def build_prompt(checklist: CaseChecklist, anchors: PromptAnchors) -> str: + """Ask the VLM for visual evidence without disclosing expected answers.""" + + questions = [ + { + "id": question.id, + "category": question.category, + "evidence_window": question.evidence_window, + "question": question.question, + } + for question in checklist.questions + ] + return "\n".join( + [ + "You are judging one generated Dynamic Memory LR-hold video using eight case-specific questions.", + "The attached images are sampled in chronological order.", + "The target generation prompt and expected answers are deliberately hidden.", + "Use only visible evidence; do not complete a plausible story from the question wording.", + "Question wording describes the intended evaluation target, but it is not visual evidence.", + "For a change claim, cite visible before-and-after images that show a concrete difference; otherwise answer no.", + "If the required subject or object disappears, is replaced, or changes into a different scene, do not treat its absence as successful completion.", + "In the reason, describe only concrete evidence visible in the cited images and do not invent an unseen state or event.", + "The evaluator has assigned the images to the model's planned camera phases.", + "Each range below is an inclusive pair of 1-based image numbers.", + "Planned phase image ranges:", + json.dumps(anchors.phase_ranges(), ensure_ascii=False, indent=2), + "", + "Useful comparison ranges:", + json.dumps(anchors.compact_ranges(), ensure_ascii=False, indent=2), + "", + "Phase guidance:", + "- returned_view means the returned_hold range, compared with the departure_baseline range when the question asks about change.", + "- returned_hold_sequence means the progression inside the returned_hold range.", + "- returned_and_later_visible_sequence means the returned_hold range through its final images.", + "- final_stationary_view means the final_stationary range.", + "- ordered_visible_sequence may use the full chronological image sequence.", + "- For a comparison or temporal claim, use visible before/after evidence when available.", + "- If an event never occurs, say that directly. Never invent a reset, reversal, worsening, disappearance, or change that is not visible.", + "- Interpret extent words relative to the stated baseline, not by absolute impression.", + "- Camera motion, framing, lighting, an action pose, or a replaced subject alone is not evidence of state progress.", + "- If the exact state is not visibly established, answer no.", + "- Cite only decisive 1-based image numbers.", + "", + "Event-start prerequisite:", + "- First decide whether the central event described by the questions visibly starts during initial_hold, before the camera leaves the original view.", + "- event_started is yes only when at least two initial_hold images show concrete temporal progress by the principal subject or object.", + "- Mere pose, contact, camera motion, background motion, or a plausible intention is not an event start.", + "- If no relevant action or state change is visibly underway by the end of initial_hold, event_started is no.", + "- Even when event_started is no, answer all eight questions from visible evidence; the evaluator applies the zero-score gate.", + "", + "Questions (expected answers are intentionally hidden):", + json.dumps(questions, ensure_ascii=False, indent=2), + "", + "Return one JSON object with exactly the keys event_started and answers.", + "event_started must contain:", + '- answer: "yes" or "no";', + "- evidence_image_numbers: at least two decisive 1-based image numbers from initial_hold;", + "- reason: one concise sentence describing the visible change, or the visible lack of change.", + "answers must be a list of exactly eight objects in the same ID order.", + "Each object must contain:", + "- id: the question ID;", + '- answer: "yes" or "no";', + "- evidence_image_numbers: a short list of decisive 1-based image numbers;", + "- reason: one concise sentence describing only the cited visible evidence.", + "Return JSON only.", + ] + ) + + +def parse_response( + raw_text: str, + expected_ids: tuple[str, ...], + *, + num_images: int, + event_image_numbers: tuple[int, ...] | None = None, +) -> VLMResponse: + text = raw_text.strip() + if text.startswith("```") and text.endswith("```"): + text = "\n".join(text.splitlines()[1:-1]).strip() + try: + value = json.loads(text) + except json.JSONDecodeError as error: + # Accept only lossless repair of missing outer braces around the two + # required top-level fields. Do not tolerate truncation or extra fields. + try: + wrapped = json.loads("{" + text + "}") + except json.JSONDecodeError: + raise InvalidVLMResponse("VLM response is not valid JSON") from error + if not isinstance(wrapped, dict) or set(wrapped) != { + "event_started", + "answers", + }: + raise InvalidVLMResponse("VLM response is not valid JSON") from error + value = wrapped + if not isinstance(value, dict) or set(value) != {"event_started", "answers"}: + raise InvalidVLMResponse("VLM response must contain exactly event_started and answers") + event = _parse_event_started( + value.get("event_started"), + num_images=num_images, + allowed_image_numbers=event_image_numbers, + ) + rows = _answer_rows(value) + if len(rows) != len(expected_ids): + raise InvalidVLMResponse(f"expected {len(expected_ids)} answers, got {len(rows)}") + answers: list[QuestionAnswer] = [] + for index, (answer, expected_id) in enumerate(zip(rows, expected_ids, strict=True)): + if not isinstance(answer, dict): + raise InvalidVLMResponse(f"answer {index} must be an object") + if answer.get("id") != expected_id: + raise InvalidVLMResponse(f"answer {index} ID must be {expected_id!r}") + predicted = str(answer.get("answer", "")).strip().lower() + if predicted not in {"yes", "no"}: + raise InvalidVLMResponse(f"{expected_id}: answer must be yes or no") + evidence = answer.get("evidence_image_numbers") + if not isinstance(evidence, list) or not evidence: + raise InvalidVLMResponse(f"{expected_id}: evidence_image_numbers must be a non-empty list") + try: + evidence_numbers = tuple(int(number) for number in evidence) + except (TypeError, ValueError) as error: + raise InvalidVLMResponse(f"{expected_id}: evidence image number must be an integer") from error + if any(number < 1 or number > num_images for number in evidence_numbers): + raise InvalidVLMResponse(f"{expected_id}: evidence image number is outside [1, {num_images}]") + reason = str(answer.get("reason", "")).strip() + if not reason: + raise InvalidVLMResponse(f"{expected_id}: reason is required") + answers.append( + QuestionAnswer( + id=expected_id, + answer=predicted, # type: ignore[arg-type] + evidence_image_numbers=evidence_numbers, + reason=reason, + ) + ) + return VLMResponse( + raw_response=raw_text, + event_started=event, + answers=tuple(answers), + ) + + +def _parse_event_started( + value, + *, + num_images: int, + allowed_image_numbers: tuple[int, ...] | None, +) -> EventStartedAnswer: + if not isinstance(value, dict): + raise InvalidVLMResponse("event_started must be an object") + answer = str(value.get("answer", "")).strip().lower() + if answer not in {"yes", "no"}: + raise InvalidVLMResponse("event_started.answer must be yes or no") + evidence = value.get("evidence_image_numbers") + if not isinstance(evidence, list) or len(evidence) < 2: + raise InvalidVLMResponse("event_started.evidence_image_numbers must contain at least two images") + try: + evidence_numbers = tuple(int(number) for number in evidence) + except (TypeError, ValueError) as error: + raise InvalidVLMResponse("event_started evidence image number must be an integer") from error + if len(set(evidence_numbers)) < 2: + raise InvalidVLMResponse("event_started evidence images must be distinct") + if any(number < 1 or number > num_images for number in evidence_numbers): + raise InvalidVLMResponse(f"event_started evidence image number is outside [1, {num_images}]") + if allowed_image_numbers is not None: + allowed = set(allowed_image_numbers) + if any(number not in allowed for number in evidence_numbers): + raise InvalidVLMResponse("event_started evidence images must come from initial_hold") + reason = str(value.get("reason", "")).strip() + if not reason: + raise InvalidVLMResponse("event_started.reason is required") + return EventStartedAnswer( + answer=answer, # type: ignore[arg-type] + evidence_image_numbers=evidence_numbers, + reason=reason, + ) + + +def _answer_rows(value) -> list: + """Read the question answers from the structured VLM response.""" + + rows = value.get("answers") + if isinstance(rows, list): + return rows + raise InvalidVLMResponse("VLM response must contain answers") + + +def chat_completions_url(base_url: str) -> str: + base = base_url.rstrip("/") + path = urllib.parse.urlparse(base).path.rstrip("/") + if path.endswith("/chat/completions"): + return base + if path.endswith("/v1"): + return base + "/chat/completions" + return base + "/v1/chat/completions" + + +def build_payload( + config: VLMConfig, + *, + prompt: str, + jpeg_images: list[bytes], +) -> dict: + content = [ + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64," + base64.b64encode(image).decode("ascii")}, + } + for image in jpeg_images + ] + content.append({"type": "text", "text": prompt}) + payload = { + "model": config.model, + "messages": [{"role": "user", "content": content}], + "temperature": 0, + "max_tokens": config.max_tokens, + "chat_template_kwargs": { + "enable_thinking": False, + "preserve_thinking": False, + }, + } + if config.backend == "api": + payload.pop("chat_template_kwargs") + payload["enable_thinking"] = False + return payload + + +class HTTPVLMBackend: + def generate( + self, + *, + prompt: str, + jpeg_images: list[bytes], + config: VLMConfig, + expected_ids: tuple[str, ...], + event_image_numbers: tuple[int, ...], + ) -> VLMResponse: + headers = {"Content-Type": "application/json"} + key = "" + if config.backend == "api": + key = os.environ.get(config.api_key_env, "") + if not key: + raise EvaluationError(f"set the API key environment variable {config.api_key_env}") + headers["Authorization"] = "Bearer " + key + request = urllib.request.Request( + chat_completions_url(config.base_url), + data=json.dumps(build_payload(config, prompt=prompt, jpeg_images=jpeg_images)).encode("utf-8"), + headers=headers, + ) + try: + with urllib.request.urlopen(request, timeout=config.timeout_seconds) as response: + body = json.loads(response.read()) + except urllib.error.HTTPError as error: + detail = error.read().decode("utf-8", errors="replace")[:2000] + if key: + detail = detail.replace(key, "[REDACTED]") + raise EvaluationError(f"VLM HTTP {error.code}: {detail}") from error + except urllib.error.URLError as error: + raise EvaluationError(f"VLM request failed: {error.reason}") from error + try: + raw_text = body["choices"][0]["message"]["content"] + except (KeyError, IndexError, TypeError) as error: + raise EvaluationError("VLM response has no chat content") from error + if not isinstance(raw_text, str) or not raw_text.strip(): + raise EvaluationError("VLM returned empty chat content") + try: + if body["choices"][0].get("finish_reason") == "length": + raise InvalidVLMResponse("VLM output reached max_tokens") + parsed = parse_response( + raw_text, + expected_ids, + num_images=len(jpeg_images), + event_image_numbers=event_image_numbers, + ) + return replace(parsed, usage=body.get("usage") or {}) + except InvalidVLMResponse as error: + raise InvalidVLMResponse(str(error), raw_response=raw_text) from error diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/__init__.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/__init__.py new file mode 100644 index 000000000..90e9af0a3 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Formal Visible Memory, Depth Spearman and Scene Reconstruction metrics.""" + +from .metric import EnvironmentDepthFidelityMetric, EnvironmentSceneReconstructionMetric, EnvironmentVisibleMemoryMetric + +__all__ = ["EnvironmentVisibleMemoryMetric", "EnvironmentDepthFidelityMetric", "EnvironmentSceneReconstructionMetric"] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/artifacts.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/artifacts.py new file mode 100644 index 000000000..3f5773edb --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/artifacts.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Read the two DA3 artifacts used by Environment scoring.""" + +from dataclasses import dataclass + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError + + +def load_poses(path, total_frames: int) -> np.ndarray: + with np.load(path, allow_pickle=False) as data: + matrices = np.asarray(data["extrinsics"], dtype=np.float64) + indices = np.asarray(data["frame_indices"]) + if total_frames < 9 or len(matrices) != total_frames: + raise DataInvalidError("DA3 must provide one pose per video frame, with at least nine frames") + if not np.array_equal(indices, np.arange(total_frames)): + raise DataInvalidError("DA3 frame_indices must equal 0..video_frame_count-1") + if matrices.shape == (total_frames, 3, 4): + padded = np.repeat(np.eye(4)[None], total_frames, axis=0) + padded[:, :3] = matrices + matrices = padded + if matrices.shape != (total_frames, 4, 4) or not np.isfinite(matrices).all(): + raise DataInvalidError("DA3 extrinsics must be finite [T,3,4] or [T,4,4] matrices") + poses = np.linalg.inv(matrices) + rotations = poses[:, :3, :3] + if ( + not np.allclose(poses[:, 3], [0, 0, 0, 1], atol=1e-6) + or not np.allclose(np.swapaxes(rotations, 1, 2) @ rotations, np.eye(3), atol=5e-3) + or not np.allclose(np.linalg.det(rotations), 1.0, atol=5e-3) + ): + raise DataInvalidError("DA3 extrinsics contain invalid camera transformations") + return np.ascontiguousarray(poses) + + +@dataclass(frozen=True) +class SceneGeometry: + frame_indices: np.ndarray + depths: np.ndarray + confidence: np.ndarray + intrinsics: np.ndarray + c2w: np.ndarray + + @property + def index_to_position(self) -> dict[int, int]: + return {int(index): position for position, index in enumerate(self.frame_indices)} + + +def load_scene_geometry(path, total_frames: int) -> SceneGeometry: + with np.load(path, allow_pickle=False) as data: + if int(data["schema_version"].item()) != 1: + raise DataInvalidError("unsupported DA3 scene geometry schema") + if int(data["source_frame_count"].item()) != total_frames: + raise DataInvalidError("DA3 scene geometry source_frame_count does not match the video") + geometry = SceneGeometry( + np.asarray(data["frame_indices"], dtype=np.int64), + np.asarray(data["depths"], dtype=np.float32), + np.asarray(data["confidence"], dtype=np.float32), + np.asarray(data["intrinsics"], dtype=np.float64), + np.asarray(data["c2w"], dtype=np.float64), + ) + indices, depths = geometry.frame_indices, geometry.depths + count = len(indices) + if ( + indices.ndim != 1 + or count < 2 + or indices[0] != 0 + or indices[-1] >= total_frames + or not np.all(np.diff(indices) > 0) + ): + raise DataInvalidError("scene frame_indices must start at zero and increase within the video") + if ( + depths.ndim != 3 + or len(depths) != count + or geometry.confidence.shape != depths.shape + or geometry.intrinsics.shape != (count, 3, 3) + or geometry.c2w.shape != (count, 4, 4) + ): + raise DataInvalidError("inconsistent DA3 scene geometry array shapes") + if any(not np.isfinite(array).all() for array in (depths, geometry.confidence, geometry.intrinsics, geometry.c2w)): + raise DataInvalidError("DA3 scene geometry contains non-finite values") + if np.any(depths < 0) or np.any(np.sum(depths > 0, axis=(1, 2)) == 0) or np.any(geometry.confidence < 0): + raise DataInvalidError("DA3 depth/confidence must be non-negative with positive depth per frame") + return geometry diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/backends.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/backends.py new file mode 100644 index 000000000..a18b213f9 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/backends.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Local-only full-frame DINOv2 and SegFormer inference.""" + +from pathlib import Path +from types import SimpleNamespace + +import numpy as np + +DEFAULT_REJECT_LABELS = ( + "person", + "car", + "boat", + "bus", + "truck", + "van", + "ship", + "minibike", + "animal", + "bicycle", + "sky", + "water", + "sea", + "river", + "lake", + "mirror", +) + + +class DinoVisibleFeatures: + def __init__(self, checkpoint, device, batch_size=8, *, model): + import torch + from transformers import AutoImageProcessor + + if batch_size < 1: + raise ValueError("DINO batch size must be positive") + checkpoint = Path(checkpoint).expanduser().resolve() + self.torch, self.device, self.batch_size = torch, torch.device(device), batch_size + self.processor = AutoImageProcessor.from_pretrained(checkpoint, local_files_only=True) + self.model = model + + def embed(self, images_rgb): + from PIL import Image + + output = [] + for start in range(0, len(images_rgb), self.batch_size): + images = [ + Image.fromarray(np.asarray(value, dtype=np.uint8)) + for value in images_rgb[start : start + self.batch_size] + ] + inputs = self.processor(images=images, return_tensors="pt").to(self.device) + with self.torch.inference_mode(): + features = self.model(**inputs).float() + features = self.torch.nn.functional.normalize(features, dim=-1) + output.append(features.cpu().numpy()) + return np.concatenate(output, axis=0) + + +class SegformerMasks: + def __init__(self, checkpoint, device, batch_size=4, reject_labels=DEFAULT_REJECT_LABELS): + import torch + from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation + + if batch_size < 1: + raise ValueError("semantic batch size must be positive") + checkpoint = Path(checkpoint).expanduser().resolve() + self.torch, self.device, self.batch_size = torch, torch.device(device), batch_size + self.processor = AutoImageProcessor.from_pretrained(checkpoint, local_files_only=True) + self.model = ( + AutoModelForSemanticSegmentation.from_pretrained(checkpoint, local_files_only=True).to(self.device).eval() + ) + rejected = {str(label).strip().casefold() for label in reject_labels} + # Exact label matching: rejecting car must not reject carpet. + self.rejected_ids = [ + int(index) + for index, label in self.model.config.id2label.items() + if str(label).strip().casefold() in rejected + ] + + def keep_masks(self, images_rgb): + from PIL import Image + + output = [] + for start in range(0, len(images_rgb), self.batch_size): + chunk = images_rgb[start : start + self.batch_size] + images = [Image.fromarray(np.asarray(value, dtype=np.uint8)) for value in chunk] + inputs = self.processor(images=images, return_tensors="pt").to(self.device) + with self.torch.inference_mode(): + logits = self.model(**inputs).logits + segmentations = self.processor.post_process_semantic_segmentation( + SimpleNamespace(logits=logits), target_sizes=[tuple(image.shape[:2]) for image in chunk] + ) + output.extend(~np.isin(labels.cpu().numpy(), self.rejected_ids) for labels in segmentations) + return output diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/geometry.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/geometry.py new file mode 100644 index 000000000..0e238da15 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/geometry.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Reference DA3 return selection and spherical nearest-surface reconstruction.""" + +from dataclasses import dataclass + +import numpy as np + + +def rotation_distance_deg(first, second) -> float: + relative = first.T @ second + cosine = float(np.clip((np.trace(relative) - 1.0) / 2.0, -1.0, 1.0)) + return float(np.degrees(np.arccos(cosine))) + + +@dataclass(frozen=True) +class ReturnCandidate: + slot_id: int + axis: str + return_index: int + rotation_error_deg: float + returned: bool + + +def select_returns( + poses, protocol, *, rotation_threshold_deg=30.0, translation_threshold=0.15, return_buffer_frames=30 +): + if rotation_threshold_deg <= 0 or return_buffer_frames < 0: + raise ValueError("rotation threshold must be positive and return buffer non-negative") + results = [] + reference = poses[0] + for slot in protocol.slots: + window = protocol.phases[slot.return_phase] + # Extend only to the right; a left extension can steal the previous return. + last = min(window.end + return_buffer_frames, len(poses) - 1) + candidates = np.arange(window.start, last + 1, dtype=np.int64) + rotations = np.asarray( + [rotation_distance_deg(reference[:3, :3], poses[i, :3, :3]) for i in candidates], dtype=np.float64 + ) + translation_gated = slot.axis in {"lateral", "longitudinal", "translation"} + if translation_gated: + positions = poses[:, :3, 3] + span = float(2.0 * np.linalg.norm(positions - positions.mean(axis=0), axis=1).max()) + translations = np.asarray( + [np.linalg.norm(reference[:3, 3] - poses[i, :3, 3]) / max(span, 1e-8) for i in candidates], + dtype=np.float64, + ) + best = int(np.argmin(rotations / rotation_threshold_deg + translations / max(translation_threshold, 1e-8))) + else: + best = int(np.argmin(rotations)) + error = float(rotations[best]) + returned = error <= rotation_threshold_deg and ( + not translation_gated or float(translations[best]) <= translation_threshold + ) + results.append(ReturnCandidate(slot.slot_id, slot.axis, int(candidates[best]), error, returned)) + return tuple(results) + + +def departure_evidence(poses, window, threshold_deg): + errors = [rotation_distance_deg(poses[0, :3, :3], poses[i, :3, :3]) for i in window.indices] + best = int(np.argmax(errors)) + maximum = float(errors[best]) + return { + "departed": maximum >= threshold_deg, + "departure_peak_index": window.start + best, + "departure_rotation_deg": maximum, + } + + +def window_indices(available, start, end, max_count=8): + values = np.asarray(available, dtype=np.int64) + values = values[(values >= start) & (values <= end)] + if len(values) > max_count: + values = values[np.asarray(np.round(np.linspace(0, len(values) - 1, max_count)), dtype=np.int64)] + return tuple(int(value) for value in values) + + +def spherical_range_field( + geometry, + indices, + anchor_rotation, + *, + semantic_masks=None, + angular_width=128, + confidence_quantile=0.5, + near_depth_percentile=5.0, + far_depth_percentile=85.0, + max_points_per_frame=5000, +): + if angular_width < 8 or angular_width % 2: + raise ValueError("angular_width must be an even integer >= 8") + if not 0 <= confidence_quantile < 1 or not 0 <= near_depth_percentile < far_depth_percentile <= 100: + raise ValueError("invalid confidence or depth quantiles") + if max_points_per_frame < 1: + raise ValueError("max_points_per_frame must be positive") + height = angular_width // 2 + flat = np.full(angular_width * height, np.inf, dtype=np.float64) + by_index = geometry.index_to_position + anchor_inverse = np.asarray(anchor_rotation, dtype=np.float64).T + for index in indices: + position = by_index[index] + depth = np.asarray(geometry.depths[position], dtype=np.float64) + confidence = np.asarray(geometry.confidence[position], dtype=np.float64) + valid = np.isfinite(depth) & np.isfinite(confidence) & (depth > 0) + if semantic_masks is not None: + valid &= semantic_masks[index] + if not valid.any(): + continue + threshold = float(np.quantile(confidence[valid], confidence_quantile)) + valid &= confidence >= threshold + low, high = np.percentile(depth[valid], [near_depth_percentile, far_depth_percentile]) + valid &= (depth >= low) & (depth <= high) + ys, xs = np.where(valid) + if not len(xs): + continue + if len(xs) > max_points_per_frame: + selected = np.asarray(np.round(np.linspace(0, len(xs) - 1, max_points_per_frame)), dtype=np.int64) + ys, xs = ys[selected], xs[selected] + pixels = np.stack([xs.astype(np.float64), ys.astype(np.float64), np.ones(len(xs))], axis=0) + rays = np.linalg.inv(geometry.intrinsics[position]) @ pixels + ranges = np.linalg.norm(rays, axis=0) * depth[ys, xs] + directions = anchor_inverse @ geometry.c2w[position, :3, :3] @ rays + directions /= np.maximum(np.linalg.norm(directions, axis=0), 1e-12) + azimuth = np.arctan2(directions[0], directions[2]) + elevation = np.arcsin(np.clip(directions[1], -1.0, 1.0)) + x = np.floor((azimuth + np.pi) * angular_width / (2.0 * np.pi)).astype(np.int64) + y = np.floor((elevation + np.pi / 2.0) * height / np.pi).astype(np.int64) + bins = np.clip(y, 0, height - 1) * angular_width + np.clip(x, 0, angular_width - 1) + np.minimum.at(flat, bins, np.log(np.maximum(ranges, 1e-12))) + values = flat.reshape(height, angular_width) + values[~np.isfinite(values)] = np.nan + return values + + +def compare_range_fields(reference, revisit): + ref_valid, rev_valid = np.isfinite(reference), np.isfinite(revisit) + common = ref_valid & rev_valid + union_bins, common_bins = int((ref_valid | rev_valid).sum()), int(common.sum()) + if not union_bins or not common_bins: + return {"reconstruction_loss": 1.0, "common_bins": common_bins, "union_bins": union_bins} + offset = float(np.median(reference[common] - revisit[common])) + residual = np.abs(reference[common] - (revisit[common] + offset)) + similarity = float(np.exp(-residual).sum() / union_bins) + similarity = float(np.clip(similarity, 0.0, 1.0)) + return {"reconstruction_loss": 1.0 - similarity, "common_bins": common_bins, "union_bins": union_bins} diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/metric.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/metric.py new file mode 100644 index 000000000..8e32f5741 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/metric.py @@ -0,0 +1,293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""The three formal Environment scores, sharing DA3 inputs and return selection.""" + +from dataclasses import asdict + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import ( + DataInvalidError, + EvaluationError, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.aggregation import ( + aggregate_video_summary, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.base import Metric +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.results import ( + ItemResult, + ResultStatus, + UnitResult, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.video import VideoReader + +from .protocol import protocol_for_sample + + +class _EnvironmentMetric(Metric): + def __init__(self, *, reader=None): + self.reader = reader or VideoReader() + + def evaluate(self, sample, context): + if sample.get("category") not in (None, "environment"): + return [UnitResult("video", status=ResultStatus.NOT_APPLICABLE)] + if error := sample.get("environment_preparation_error"): + exception = DataInvalidError if error["status"] == "data_invalid" else EvaluationError + raise exception(error["message"]) + + def load_inputs(): + from .artifacts import load_poses + from .geometry import select_returns + + info = self.reader.probe(sample["video"]) + protocol = protocol_for_sample(sample, total_frames=info.total_frames) + artifacts = sample.get("artifacts") or {} + if not artifacts.get("da3"): + raise DataInvalidError("missing artifacts.da3 (da3/results.npz)") + poses = load_poses(artifacts["da3"], info.total_frames) + config = context.metric_config("environment_return_validity") + returns = select_returns( + poses, + protocol, + rotation_threshold_deg=float(config.get("rotation_threshold_deg", 30.0)), + translation_threshold=float(config.get("translation_threshold", 0.15)), + return_buffer_frames=int(config.get("return_buffer_frames", 30)), + ) + return protocol, poses, returns + + protocol, poses, returns = context.get_cached(("environment:inputs", sample["video"]), load_inputs) + scores, metadata = self.score(sample, context, protocol, poses, returns) + return [ + UnitResult( + "video", + unit_type="video", + scores=scores, + metadata={ + "video_path": sample["video"], + "num_frames": protocol.total_frames, + "actions": list(protocol.actions), + "phase_transitions": list(protocol.phase_transitions), + "tail_hold_transitions": protocol.tail_hold_transitions, + "phases": [asdict(phase) for phase in protocol.phases], + **metadata, + }, + ) + ] + + @staticmethod + def geometry(sample, context, total_frames): + from .artifacts import load_scene_geometry + + path = (sample.get("artifacts") or {}).get("da3_scene") + if not path: + raise DataInvalidError("missing artifacts.da3_scene (da3/scene_geometry.npz)") + return context.get_cached(("environment:geometry", path), lambda: load_scene_geometry(path, total_frames)) + + def aggregate_item(self, sample, units): + unit = units[0] + return ItemResult( + sample["item_id"], + scores=dict(unit.scores), + status=unit.status, + error=unit.error, + metadata=dict(unit.metadata), + num_units=1, + ) + + def aggregate_summary(self, items): + return aggregate_video_summary(items) + + +class EnvironmentVisibleMemoryMetric(_EnvironmentMetric): + name = "environment_visible_memory" + description = "Mean full-frame DINOv2 similarity at DA3 return frames." + + def __init__(self, *, reader=None, backend=None): + super().__init__(reader=reader) + self.backend = backend + + def score(self, sample, context, protocol, poses, returns): + import numpy as np + + from ..features import get_dino_model + from .backends import DinoVisibleFeatures + + config = context.metric_config(self.name) + checkpoint = config.get("dino_checkpoint") or context.cache_dir / "metrics" / self.name / "dinov2-large" + required = sorted({0, *(value.return_index for value in returns)}) + frames = self.reader.read_indices(sample["video"], required) + backend = self.backend or context.get_resource( + "environment:visible_dino", + lambda: DinoVisibleFeatures( + checkpoint, + context.device, + int(config.get("batch_size", 8)), + model=get_dino_model(context, checkpoint), + ), + ) + features = backend.embed([frame.rgb for frame in frames]) + if len(features) != len(required): + raise EvaluationError("DINO returned the wrong number of feature rows") + by_index = dict(zip(required, features)) + reference = np.asarray(by_index[0], dtype=np.float64) + + def similarity(index): + return float(np.clip(reference @ np.asarray(by_index[index], dtype=np.float64), 0.0, 1.0)) + + slots = [ + {"slot_id": value.slot_id, "return_index": value.return_index, "similarity": similarity(value.return_index)} + for value in returns + ] + visible = float(np.mean([slot["similarity"] for slot in slots])) + return {"score": visible}, {"visible_similarity": visible, "slots": slots} + + +class EnvironmentDepthFidelityMetric(_EnvironmentMetric): + name = "environment_depth_fidelity" + description = "Pixel-aligned DA3 depth-order Spearman correlation at return frames." + + def score(self, sample, context, protocol, poses, returns): + import numpy as np + from scipy.stats import spearmanr + + geometry = self.geometry(sample, context, protocol.total_frames) + by_index = geometry.index_to_position + reference = np.asarray(geometry.depths[by_index[0]], dtype=np.float64).ravel() + slots = [] + for value in returns: + sampled = min(by_index, key=lambda index: abs(index - value.return_index)) + candidate = np.asarray(geometry.depths[by_index[sampled]], dtype=np.float64).ravel() + valid = np.isfinite(reference) & np.isfinite(candidate) & (reference > 0) & (candidate > 0) + if int(valid.sum()) < 100: + raise EvaluationError("depth fidelity requires at least 100 valid pixels") + score = float(spearmanr(reference[valid], candidate[valid]).statistic) + if not np.isfinite(score): + raise EvaluationError("depth fidelity produced a non-finite Spearman value") + slots.append( + { + "slot_id": value.slot_id, + "axis": value.axis, + "return_status": "returned" if value.returned else "not_revisited", + "return_index": value.return_index, + "sampled_return_index": sampled, + "depth_spearman": score, + } + ) + return {"depth_spearman": float(np.mean([slot["depth_spearman"] for slot in slots]))}, {"slots": slots} + + +class EnvironmentSceneReconstructionMetric(_EnvironmentMetric): + name = "environment_scene_reconstruction" + description = "Pure-rotation spherical reconstruction score (1 - reference loss), higher is better." + + def __init__(self, *, reader=None, semantic_filter=None): + super().__init__(reader=reader) + self.semantic_filter = semantic_filter + + def score(self, sample, context, protocol, poses, returns): + import cv2 + import numpy as np + + from .backends import DEFAULT_REJECT_LABELS, SegformerMasks + from .geometry import compare_range_fields, departure_evidence, spherical_range_field, window_indices + + if any(slot.axis not in {"yaw", "pitch"} for slot in protocol.slots): + raise EvaluationError("Scene Reconstruction supports pure-rotation yaw/pitch trajectories only") + config = context.metric_config(self.name) + geometry = self.geometry(sample, context, protocol.total_frames) + available = geometry.frame_indices + home_fraction, home_frames = ( + float(config.get("home_window_fraction", 0.2)), + int(config.get("home_window_frames", 8)), + ) + if not 0 < home_fraction <= 1 or home_frames < 1: + raise ValueError("invalid reconstruction home window") + phase0 = protocol.phases[0] + home = window_indices( + available, phase0.start, int(round(phase0.start + home_fraction * (phase0.end - phase0.start))), home_frames + ) + radius = int(config.get("revisit_radius_frames", 14)) + threshold = float(context.metric_config("environment_return_validity").get("rotation_threshold_deg", 30.0)) + prepared, required = [], set(home) + for slot, value in zip(protocol.slots, returns, strict=True): + indices = window_indices( + available, + max(int(available[0]), value.return_index - radius), + min(int(available[-1]), value.return_index + radius), + home_frames, + ) + departure = departure_evidence(poses, protocol.phases[slot.outbound_phase], threshold) + prepared.append((value, indices, departure)) + required.update(indices) + masks = None + if config.get("use_semantic_filter", True) and required: + backend = self.semantic_filter or context.get_resource( + "environment:segformer", + lambda: SegformerMasks( + config.get("semantic_checkpoint") + or context.cache_dir / "metrics" / self.name / "segformer-b0-ade20k", + context.device, + int(config.get("semantic_batch_size", 4)), + config.get("semantic_reject_labels") or DEFAULT_REJECT_LABELS, + ), + ) + indices = sorted(required) + images = [frame.rgb for frame in self.reader.read_indices(sample["video"], indices)] + predicted = backend.keep_masks(images) + if len(predicted) != len(indices): + raise EvaluationError("SegFormer returned the wrong number of masks") + height, width = geometry.depths.shape[1:] + masks = { + index: cv2.resize( + np.asarray(mask, dtype=np.uint8), (width, height), interpolation=cv2.INTER_NEAREST + ).astype(bool) + for index, mask in zip(indices, predicted) + } + settings = dict( + semantic_masks=masks, + angular_width=int(config.get("angular_width", 128)), + confidence_quantile=float(config.get("confidence_quantile", 0.5)), + near_depth_percentile=float(config.get("near_depth_percentile", 5.0)), + far_depth_percentile=float(config.get("far_depth_percentile", 85.0)), + max_points_per_frame=int(config.get("max_points_per_frame", 5000)), + ) + reference = spherical_range_field(geometry, home, poses[0, :3, :3], **settings) + minimum_bins, require_departure = ( + int(config.get("minimum_common_bins", 64)), + bool(config.get("require_pose_departure", True)), + ) + losses, slots = [], [] + for value, indices, departure in prepared: + revisit = spherical_range_field(geometry, indices, poses[value.return_index, :3, :3], **settings) + comparison = compare_range_fields(reference, revisit) + usable = ( + value.returned + and (departure["departed"] or not require_departure) + and comparison["common_bins"] >= minimum_bins + ) + loss = comparison["reconstruction_loss"] if usable else 1.0 + losses.append(float(loss)) + if usable: + status = "ok" + elif require_departure and not departure["departed"]: + status = "not_departed" + elif not value.returned: + status = "not_revisited" + else: + status = "insufficient_common_bins" + slots.append( + { + "slot_id": value.slot_id, + "axis": value.axis, + "status": status, + "return_status": "returned" if value.returned else "not_revisited", + "return_index": value.return_index, + "return_rotation_error_deg": value.rotation_error_deg, + "reference_indices": list(home), + "revisit_indices": list(indices), + **departure, + "reconstruction_score": 1.0 - loss, + "common_bins": comparison["common_bins"], + "union_bins": comparison["union_bins"], + } + ) + return {"reconstruction_score": 1.0 - float(np.mean(losses))}, {"slots": slots} diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/preparation.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/preparation.py new file mode 100644 index 000000000..79e1c1685 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/preparation.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Resolve the Environment case layout without inferring missing trajectory data.""" + +import json +from pathlib import Path + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError + +from .protocol import protocol_fields + + +def prepare_for_evaluation(samples: list[dict]) -> None: + """Read trajectory.json and register DA3 paths; inference is a separate step.""" + for sample in samples: + if sample.get("category") not in (None, "environment"): + continue + sample.pop("environment_preparation_error", None) + try: + case_dir = Path(sample["video"]).parent + annotations = dict(sample.get("annotations") or {}) + if "environment_protocol" not in annotations: + annotations["environment_protocol"] = json.loads( + (case_dir / "trajectory.json").read_text(encoding="utf-8") + ) + annotations["environment_protocol"] = protocol_fields(annotations["environment_protocol"]) + artifacts = dict(sample.get("artifacts") or {}) + for key, filename in (("da3", "results.npz"), ("da3_scene", "scene_geometry.npz")): + artifacts.setdefault(key, str((case_dir / "da3" / filename).resolve())) + sample["annotations"] = annotations + sample["artifacts"] = artifacts + except (DataInvalidError, OSError, ValueError, KeyError, TypeError) as error: + sample["environment_preparation_error"] = { + "status": "data_invalid", + "message": f"invalid Environment trajectory: {error}", + } + + +def prepare_artifacts(samples: list[dict], context, metric_names) -> None: + """Run before a shard loads scoring models, then release the DA3 subprocess.""" + from ...preparation.da3 import ensure_artifacts + + require_scene = bool({"environment_depth_fidelity", "environment_scene_reconstruction"}.intersection(metric_names)) + weight_metric = next(name for name in metric_names if name.startswith("environment_")) + eligible = [ + sample + for sample in samples + if sample.get("category") in (None, "environment") and not sample.get("environment_preparation_error") + ] + if eligible: + print(f"[{context.device}][DA3] 0/{len(eligible)}", flush=True) + for index, sample in enumerate(eligible, 1): + try: + ensure_artifacts(sample, context, weight_metric=weight_metric, require_scene=require_scene) + except Exception as error: + sample["environment_preparation_error"] = { + "status": "eval_failed", + "message": f"DA3 preparation failed: {type(error).__name__}: {error}", + } + print( + f"[{context.device}][DA3] {index}/{len(eligible)} failed: {sample['item_id']}: " + f"{sample['environment_preparation_error']['message']}", + flush=True, + ) + else: + print(f"[{context.device}][DA3] {index}/{len(eligible)} ready: {sample['item_id']}", flush=True) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/protocol.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/protocol.py new file mode 100644 index 000000000..eba0c6217 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/environment/protocol.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Map explicit trajectory phases to video frames using the reference protocol.""" + +from dataclasses import dataclass + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError + +INVERSE_ACTION = { + "yaw_left": "yaw_right", + "yaw_right": "yaw_left", + "pitch_up": "pitch_down", + "pitch_down": "pitch_up", + "move_left": "move_right", + "move_right": "move_left", + "move_forward": "move_backward", + "move_backward": "move_forward", +} + + +@dataclass(frozen=True) +class PhaseWindow: + phase_id: int + action: str + axis: str + start: int + end: int + + @property + def indices(self) -> range: + return range(self.start, self.end + 1) + + +@dataclass(frozen=True) +class RevisitSlot: + slot_id: int + outbound_phase: int + return_phase: int + axis: str + + +@dataclass(frozen=True) +class EnvironmentProtocol: + condition_id: str + total_frames: int + actions: tuple[str, ...] + phase_transitions: tuple[int, ...] + tail_hold_transitions: int + phases: tuple[PhaseWindow, ...] + slots: tuple[RevisitSlot, ...] + + +def protocol_fields(value: dict) -> dict: + """Read the three scoring fields; do not infer timing from a directory name.""" + if not isinstance(value, dict): + raise DataInvalidError("Environment trajectory must be a JSON object") + actions = value.get("actions") + transitions = value.get("phase_transitions") + tail = value.get("tail_hold_transitions", 0) + if ( + not isinstance(actions, list) + or not actions + or len(actions) % 2 + or any(not isinstance(action, str) or action not in INVERSE_ACTION for action in actions) + ): + raise DataInvalidError("trajectory.actions must contain supported outbound/return action pairs") + if any(INVERSE_ACTION[actions[i]] != actions[i + 1] for i in range(0, len(actions), 2)): + raise DataInvalidError("each Environment action pair must be inverse") + if ( + not isinstance(transitions, list) + or len(transitions) != len(actions) + or any(type(count) is not int or count <= 0 for count in transitions) + ): + raise DataInvalidError("trajectory.phase_transitions must contain one positive integer per action") + if type(tail) is not int or tail < 0: + raise DataInvalidError("trajectory.tail_hold_transitions must be a non-negative integer") + return {"actions": list(actions), "phase_transitions": list(transitions), "tail_hold_transitions": tail} + + +def protocol_for_sample(sample: dict, *, total_frames: int) -> EnvironmentProtocol: + fields = protocol_fields((sample.get("annotations") or {}).get("environment_protocol")) + if total_frames < 3: + raise DataInvalidError("Environment video requires at least three frames") + actions = tuple(fields["actions"]) + transitions = tuple(fields["phase_transitions"]) + tail = fields["tail_hold_transitions"] + source_last = sum(transitions) + tail + phases = [] + cursor = 0 + for index, (action, count) in enumerate(zip(actions, transitions)): + start = int(round(cursor * (total_frames - 1) / source_last)) + cursor += count + end = int(round(cursor * (total_frames - 1) / source_last)) + if end <= start: + raise DataInvalidError(f"phase {index} collapses in decoded frame space") + if action in ("move_left", "move_right"): + axis = "lateral" + elif action in ("move_forward", "move_backward"): + axis = "longitudinal" + else: + axis = action.split("_", 1)[0] + phases.append(PhaseWindow(index, action, axis, start, end)) + slots = tuple(RevisitSlot(i // 2 + 1, i, i + 1, phases[i].axis) for i in range(0, len(phases), 2)) + return EnvironmentProtocol( + str(sample.get("condition_id") or ""), total_frames, actions, transitions, tail, tuple(phases), slots + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/features.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/features.py new file mode 100644 index 000000000..24e3a544a --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/features.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Masked DINOv2 appearance features shared by human and object protocols.""" + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.context import EvaluationContext +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + +from .human.target_association import normalized_embedding + + +def get_dino_model(context: EvaluationContext, checkpoint: str | Path) -> Any: + """Share a local CLS backbone while keeping each protocol's preprocessing separate.""" + from worldfoundry.base_models.perception_core.general_perception.dino_embeddings import load_dinov2_feature_model + + checkpoint = Path(checkpoint).expanduser().resolve() + return context.get_resource( + ("dinov2-cls", checkpoint, context.device), + lambda: load_dinov2_feature_model(checkpoint, device=context.device), + ) + + +class MaskedDinoV2Features: + """Batched masked-crop DINOv2-Large CLS embeddings.""" + + def __init__( + self, + *, + model: Any, + device: str, + input_size: int = 518, + batch_size: int = 16, + ) -> None: + if batch_size < 1: + raise ValueError("DINO batch_size must be positive") + try: + import torch + import torchvision.transforms as transforms + except ImportError as error: + raise EvaluationError( + "Masked DINOv2 features require torch and torchvision; install the WorldMemBench metric requirements" + ) from error + self.model = model + self.transform = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Resize((input_size, input_size), antialias=True), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225], + ), + ] + ) + self.device = device + self.torch = torch + self.effective_batch_size = batch_size + + @staticmethod + def _masked_crop( + frame_rgb: np.ndarray, + mask: np.ndarray, + ) -> np.ndarray | None: + if mask.shape != frame_rgb.shape[:2] or not mask.any(): + return None + y_coordinates, x_coordinates = np.where(mask) + x1, x2 = int(x_coordinates.min()), int(x_coordinates.max()) + 1 + y1, y2 = int(y_coordinates.min()), int(y_coordinates.max()) + 1 + crop = frame_rgb[y1:y2, x1:x2].copy() + crop_mask = mask[y1:y2, x1:x2] + crop[~crop_mask] = 0 + return crop + + @staticmethod + def _is_out_of_memory(error: RuntimeError) -> bool: + return "out of memory" in str(error).lower() + + def embed_batch( + self, + samples: Sequence[tuple[np.ndarray, np.ndarray]], + ) -> list[np.ndarray | None]: + results: list[np.ndarray | None] = [None] * len(samples) + position = 0 + while position < len(samples): + chunk_size = min( + self.effective_batch_size, + len(samples) - position, + ) + chunk = samples[position : position + chunk_size] + tensors = [] + result_indices = [] + for offset, (frame_rgb, mask) in enumerate(chunk): + crop = self._masked_crop(frame_rgb, mask) + if crop is None: + continue + tensors.append(self.transform(crop)) + result_indices.append(position + offset) + if not tensors: + position += chunk_size + continue + try: + tensor_batch = self.torch.stack(tensors).to(self.device) + with self.torch.inference_mode(): + vectors = self.model(tensor_batch) + vectors = vectors.detach().float().cpu().numpy() + except RuntimeError as error: + if not self._is_out_of_memory(error) or self.effective_batch_size == 1: + raise + self.effective_batch_size = max( + 1, + self.effective_batch_size // 2, + ) + if self.torch.cuda.is_available(): + self.torch.cuda.empty_cache() + continue + for result_index, vector in zip(result_indices, vectors): + results[result_index] = normalized_embedding(vector) + position += chunk_size + return results + + def embed(self, frame_rgb: np.ndarray, mask: np.ndarray) -> np.ndarray | None: + """Extract one masked crop with the same preprocessing as a batch.""" + return self.embed_batch([(frame_rgb, mask)])[0] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/__init__.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/__init__.py new file mode 100644 index 000000000..c4a035fcd --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""General Video Quality metrics; model backends are imported lazily.""" + +from .aesthetic_quality import AestheticQualityMetric +from .imaging_quality import ImagingQualityMetric + +__all__ = ["AestheticQualityMetric", "ImagingQualityMetric"] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/aesthetic.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/aesthetic.py new file mode 100644 index 000000000..d2f6ba4e7 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/aesthetic.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""OpenAI CLIP + LAION aesthetic predictor backend.""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + +CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073) +CLIP_STD = (0.26862954, 0.26130258, 0.27577711) + + +def preprocess_clip_rgb( + frames: list[np.ndarray], + image_size: int = 224, +) -> Any: + """Apply VBench's tensor CLIP transform to RGB uint8 frames.""" + + if image_size <= 0: + raise ValueError("image_size must be positive") + if not frames: + import torch + + return torch.empty((0, 3, image_size, image_size), dtype=torch.float32) + + import torch + from torchvision.transforms import CenterCrop, InterpolationMode, Resize + + resize = Resize( + image_size, + interpolation=InterpolationMode.BICUBIC, + antialias=False, + ) + center_crop = CenterCrop(image_size) + tensors = [] + for frame in frames: + if frame.ndim != 3 or frame.shape[2] != 3: + raise ValueError(f"expected HxWx3 RGB image, got shape {frame.shape}") + tensor = torch.from_numpy(np.ascontiguousarray(frame)).permute(2, 0, 1).to(dtype=torch.float32) + tensors.append(center_crop(resize(tensor))) + + batch = torch.stack(tensors).to(dtype=torch.float32).div_(255.0) + mean = torch.tensor(CLIP_MEAN, dtype=batch.dtype).view(1, 3, 1, 1) + std = torch.tensor(CLIP_STD, dtype=batch.dtype).view(1, 3, 1, 1) + return (batch - mean) / std + + +class LaionAestheticBackend: + """Fixed VBench-compatible ViT-L/14 aesthetic scoring backend.""" + + def __init__( + self, + *, + clip_checkpoint: str | Path, + aesthetic_checkpoint: str | Path, + device: str, + batch_size: int = 32, + ) -> None: + self.clip_checkpoint = Path(clip_checkpoint).expanduser().resolve() + self.aesthetic_checkpoint = Path(aesthetic_checkpoint).expanduser().resolve() + if not self.clip_checkpoint.is_file(): + raise EvaluationError(f"clip checkpoint not found: {self.clip_checkpoint}") + if not self.aesthetic_checkpoint.is_file(): + raise EvaluationError(f"aesthetic checkpoint not found: {self.aesthetic_checkpoint}") + if batch_size <= 0: + raise ValueError("batch_size must be positive") + + self.device = device + self.batch_size = batch_size + try: + from worldfoundry.base_models.perception_core.general_perception import openai_clip as clip + from worldfoundry.evaluation.tasks.metrics._shared.aesthetic import load_laion_aesthetic_linear_head + except ImportError as error: + raise EvaluationError( + "aesthetic_quality requires torch and torchvision; install the WorldMemBench metric requirements" + ) from error + + self.clip_model, _ = clip.load( + str(self.clip_checkpoint), + device=device, + jit=False, + ) + self.clip_model.eval() + + self.aesthetic_model = load_laion_aesthetic_linear_head(self.aesthetic_checkpoint).to(device) + + def score(self, frames: list[np.ndarray]) -> list[float]: + if not frames: + return [] + + import torch + import torch.nn.functional as functional + + outputs: list[float] = [] + for start in range(0, len(frames), self.batch_size): + batch = preprocess_clip_rgb(frames[start : start + self.batch_size]).to(self.device) + with torch.inference_mode(): + features: Any = self.clip_model.encode_image(batch) + features = functional.normalize( + features.to(torch.float32), + dim=-1, + p=2, + ) + raw: Any = self.aesthetic_model(features).squeeze(dim=-1) + outputs.extend(float(value) for value in raw.detach().float().cpu().reshape(-1).tolist()) + + if len(outputs) != len(frames): + raise EvaluationError(f"LAION aesthetic backend returned {len(outputs)} scores for {len(frames)} frames") + if not all(math.isfinite(value) for value in outputs): + raise EvaluationError("LAION aesthetic backend returned a non-finite score") + return outputs diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/aesthetic_quality.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/aesthetic_quality.py new file mode 100644 index 000000000..5b0c6a9f2 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/aesthetic_quality.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""OpenAI CLIP ViT-L/14 + LAION aesthetic quality.""" + +from pathlib import Path + +from .common import FrameQualityMetric + + +class AestheticQualityMetric(FrameQualityMetric): + name = "aesthetic_quality" + description = "CLIP ViT-L/14 + LAION aesthetic quality (stride 4, raw / 10)." + raw_key = "raw_score" + normalizer = 10.0 + + def _backend(self, context): + from .aesthetic import LaionAestheticBackend + + config = context.metric_config(self.name) + clip_path = ( + Path(config.get("clip_checkpoint") or (context.cache_dir / "metrics" / self.name / "ViT-L-14.pt")) + .expanduser() + .resolve() + ) + aesthetic_path = ( + Path( + config.get("aesthetic_checkpoint") + or (context.cache_dir / "metrics" / self.name / "sa_0_4_vit_l_14_linear.pth") + ) + .expanduser() + .resolve() + ) + batch_size = int(config.get("batch_size", 32)) + return context.get_resource( + f"laion_aesthetic:{clip_path}:{aesthetic_path}:{context.device}:{batch_size}", + lambda: LaionAestheticBackend( + clip_checkpoint=clip_path, + aesthetic_checkpoint=aesthetic_path, + device=context.device, + batch_size=batch_size, + ), + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/common.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/common.py new file mode 100644 index 000000000..faf3320a6 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/common.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Reference stride-4 frame scoring and equal-weight video aggregation.""" + +import math + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.context import EvaluationContext +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.aggregation import ( + aggregate_video_summary, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.base import Metric +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.results import ItemResult, UnitResult +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.video import ( + VideoReader, + fixed_stride_indices, +) + + +class FrameQualityMetric(Metric): + raw_key: str + normalizer: float + unit_type = "frame" + + def __init__(self, *, backend=None, reader=None) -> None: + self.backend = backend + self.reader = reader or VideoReader() + + def _backend(self, context: EvaluationContext): + raise NotImplementedError + + def evaluate(self, sample: dict, context: EvaluationContext) -> list[UnitResult]: + path = sample["video"] + info = self.reader.probe(path) + indices = fixed_stride_indices(info.total_frames, 4) + decoded = context.get_cached( + ("rgb_frames", path, tuple(indices)), + lambda: self.reader.read_indices(path, indices), + ) + backend = self.backend if self.backend is not None else self._backend(context) + raw_scores = [float(value) for value in backend.score([frame.rgb for frame in decoded])] + if len(raw_scores) != len(decoded): + raise EvaluationError(f"{self.name} backend returned {len(raw_scores)} scores for {len(decoded)} frames") + if not all(math.isfinite(value) for value in raw_scores): + raise EvaluationError(f"{self.name} backend returned a non-finite score") + return [ + UnitResult( + unit_id=f"frame_{frame.frame_idx:08d}", + unit_type="frame", + scores={"score": max(0.0, min(1.0, raw / self.normalizer)), self.raw_key: raw}, + metadata={"frame_idx": frame.frame_idx, "video_path": path}, + ) + for frame, raw in zip(decoded, raw_scores) + ] + + def aggregate_item(self, sample: dict, units: list[UnitResult]) -> ItemResult: + if not units: + raise EvaluationError(f"{self.name} produced no frame units") + identity = {key: sample.get(key) for key in ("item_id", "model_id", "category", "sample_id", "condition_id")} + non_scored = [unit for unit in units if not unit.status.is_scored] + if non_scored: + first = non_scored[0] + return ItemResult(**identity, num_units=len(units), status=first.status, error=first.error) + return ItemResult( + **identity, + scores={key: sum(unit.scores[key] for unit in units) / len(units) for key in ("score", self.raw_key)}, + num_units=len(units), + metadata={"sampled_frame_indices": [unit.metadata["frame_idx"] for unit in units]}, + ) + + aggregate_summary = staticmethod(aggregate_video_summary) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/imaging_quality.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/imaging_quality.py new file mode 100644 index 000000000..e6468c64d --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/imaging_quality.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""MUSIQ-SPAQ imaging quality, unchanged from the reference calculation.""" + +from pathlib import Path + +from .common import FrameQualityMetric + + +class ImagingQualityMetric(FrameQualityMetric): + name = "imaging_quality" + description = "MUSIQ-SPAQ imaging quality (stride 4, MOS / 100)." + raw_key = "raw_mos" + normalizer = 100.0 + + def _backend(self, context): + from .musiq import MusiqSpaqBackend + + config = context.metric_config(self.name) + checkpoint = ( + Path( + config.get("checkpoint") or (context.cache_dir / "metrics" / self.name / "musiq_spaq_ckpt-358bb6af.pth") + ) + .expanduser() + .resolve() + ) + batch_size = int(config.get("batch_size", 8)) + return context.get_resource( + f"musiq_spaq:{checkpoint}:{context.device}:{batch_size}", + lambda: MusiqSpaqBackend(checkpoint=checkpoint, device=context.device, batch_size=batch_size), + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/musiq.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/musiq.py new file mode 100644 index 000000000..bdbc3cb39 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/general_video_quality/musiq.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Fixed MUSIQ-SPAQ backend with VBench-compatible longer preprocessing.""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + + +def resize_longer_rgb(rgb: np.ndarray, max_side: int = 512) -> np.ndarray: + if rgb.ndim != 3 or rgb.shape[2] != 3: + raise ValueError(f"expected HxWx3 RGB image, got shape {rgb.shape}") + height, width = rgb.shape[:2] + longest = max(height, width) + if longest <= max_side: + return rgb + import cv2 + + scale = max_side / float(longest) + resized_height = max(1, int(scale * height)) + resized_width = max(1, int(scale * width)) + return cv2.resize( + rgb, + (resized_width, resized_height), + interpolation=cv2.INTER_LINEAR, + ) + + +class MusiqSpaqBackend: + def __init__( + self, + *, + checkpoint: str | Path, + device: str, + batch_size: int = 8, + ) -> None: + self.checkpoint = Path(checkpoint).expanduser().resolve() + if not self.checkpoint.is_file(): + raise EvaluationError(f"MUSIQ-SPAQ checkpoint not found: {self.checkpoint}") + if batch_size <= 0: + raise ValueError("batch_size must be positive") + self.device = device + self.batch_size = batch_size + try: + from pyiqa.archs.musiq_arch import MUSIQ + except ImportError as error: + raise EvaluationError( + "imaging_quality requires pyiqa; install the WorldMemBench metric requirements" + ) from error + self.model = MUSIQ(pretrained_model_path=str(self.checkpoint)) + self.model.to(device) + self.model.eval() + + def score(self, frames: list[np.ndarray]) -> list[float]: + if not frames: + return [] + import torch + + outputs: list[float] = [] + for start in range(0, len(frames), self.batch_size): + resized = [ + np.ascontiguousarray(resize_longer_rgb(frame)) for frame in frames[start : start + self.batch_size] + ] + tensor = ( + torch.from_numpy(np.stack(resized)) + .permute(0, 3, 1, 2) + .contiguous() + .to(device=self.device, dtype=torch.float32) + / 255.0 + ) + with torch.inference_mode(): + raw: Any = self.model(tensor) + values = raw.detach().float().cpu().reshape(-1).tolist() + outputs.extend(float(value) for value in values) + if len(outputs) != len(frames): + raise EvaluationError(f"MUSIQ returned {len(outputs)} scores for {len(frames)} frames") + if not all(math.isfinite(value) for value in outputs): + raise EvaluationError("MUSIQ returned a non-finite score") + return outputs diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/__init__.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/__init__.py new file mode 100644 index 000000000..31d888ccc --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Human Entity Memory metrics.""" + +from .metric import HumanAppearanceMetric, HumanIdentityMetric + +__all__ = ["HumanAppearanceMetric", "HumanIdentityMetric"] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/appearance_features.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/appearance_features.py new file mode 100644 index 000000000..357397399 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/appearance_features.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Phase-local body tracking and masked DINOv2 appearance scoring.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Protocol + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + +from .person_mask_tracker import MaskPropagator +from .schemas import ( + AppearanceAnalysis, + AppearanceTargetResult, + BoundingBox, + IdentityAnalysis, +) +from .target_association import ( + cosine_matrix, + normalized_embedding, +) + + +class AppearanceBackend(Protocol): + def extract( + self, + identity: IdentityAnalysis, + ) -> dict[int, dict[int, np.ndarray]]: ... + + +class AppearanceFeatureExtractor(Protocol): + def embed_batch( + self, + samples: Sequence[tuple[np.ndarray, np.ndarray]], + ) -> list[np.ndarray | None]: ... + + +class YoloSamDinoAppearanceBackend: + def __init__( + self, + *, + mask_propagator: MaskPropagator, + feature_extractor: AppearanceFeatureExtractor, + ) -> None: + self.mask_propagator = mask_propagator + self.feature_extractor = feature_extractor + + def _target_phase_features( + self, + identity: IdentityAnalysis, + target_id: int, + phase_indices: Sequence[int], + person_assignments: dict[int, dict[int, BoundingBox]], + ) -> dict[int, dict[int, np.ndarray]]: + if not phase_indices: + return {} + frames = [identity.decoded_frames[index] for index in phase_indices] + target_prompts: list[tuple[int, BoundingBox]] = [] + for local_index, frame_index in enumerate(phase_indices): + person_box = person_assignments.get(frame_index, {}).get(target_id) + if person_box is not None: + target_prompts.append((local_index, person_box)) + prompts = {target_id: target_prompts} if target_prompts else {} + # Appearance keeps only frames with a direct body-box prompt below. + # Propagated masks cannot replace those direct masks and are discarded. + local_masks = self.mask_propagator.track( + frames, + prompts, + propagate=False, + ) + features: dict[int, dict[int, np.ndarray]] = {} + sample_keys: list[tuple[int, int]] = [] + samples: list[tuple[np.ndarray, np.ndarray]] = [] + for tracked_target_id, target_masks in sorted(local_masks.items()): + for local_index, mask in sorted(target_masks.items()): + if local_index < 0 or local_index >= len(phase_indices): + raise EvaluationError("SAM2 returned an out-of-range phase frame") + frame_index = phase_indices[local_index] + # A propagated mask is valid only while the body association + # confirms that the tracked mask still overlaps a person. + if tracked_target_id not in person_assignments.get(frame_index, {}): + continue + sample_keys.append((tracked_target_id, frame_index)) + samples.append( + ( + identity.decoded_frames[frame_index], + np.asarray(mask, dtype=bool), + ) + ) + embeddings = self.feature_extractor.embed_batch(samples) + if len(embeddings) != len(sample_keys): + raise EvaluationError("DINO batch output count does not match its input count") + for (tracked_target_id, frame_index), embedding in zip( + sample_keys, + embeddings, + ): + if embedding is not None: + features.setdefault(tracked_target_id, {})[frame_index] = normalized_embedding(embedding) + return features + + def extract(self, identity: IdentityAnalysis) -> dict[int, dict[int, np.ndarray]]: + person_assignments = identity.body_assignments + merged: dict[int, dict[int, np.ndarray]] = {} + for phase in identity.appearance_target_phases: + if not phase.is_scorable: + continue + visible = self._target_phase_features( + identity, + phase.target_id, + phase.visible_indices, + person_assignments, + ) + returned = self._target_phase_features( + identity, + phase.target_id, + phase.return_indices, + person_assignments, + ) + for source in (visible, returned): + for target_id, frame_features in source.items(): + merged.setdefault(target_id, {}).update(frame_features) + return merged + + +def analyze_appearance( + identity: IdentityAnalysis, + backend: AppearanceBackend, +) -> AppearanceAnalysis: + features = backend.extract(identity) + results: list[AppearanceTargetResult] = [] + for phase in identity.appearance_target_phases: + target_id = phase.target_id + target_features = features.get(target_id, {}) + references = [target_features[index] for index in phase.visible_indices if index in target_features] + frame_scores: dict[int, float] = {} + n_return_features = 0 + for frame_index in phase.return_indices: + candidate = target_features.get(frame_index) + if candidate is None: + frame_scores[frame_index] = 0.0 + continue + n_return_features += 1 + if not references or not phase.is_scorable: + frame_scores[frame_index] = 0.0 + continue + similarities = cosine_matrix(references, [candidate])[:, 0] + frame_scores[frame_index] = float(np.clip(similarities.max(), 0.0, 1.0)) + return_frame_count = len(phase.return_indices) + score = sum(frame_scores.values()) / n_return_features if n_return_features else 0.0 + coverage = n_return_features / return_frame_count if return_frame_count else 0.0 + if not phase.is_scorable: + score = 0.0 + revisit_scores: list[float] = [] + revisit_coverages: list[float] = [] + n_revisit_features: list[int] = [] + for revisit_indices in phase.revisit_indices: + feature_indices = [index for index in revisit_indices if index in target_features] + n_features = len(feature_indices) + n_revisit_features.append(n_features) + revisit_coverages.append(n_features / len(revisit_indices) if revisit_indices else 0.0) + if n_features and phase.is_scorable: + revisit_scores.append(sum(frame_scores[index] for index in feature_indices) / n_features) + else: + revisit_scores.append(0.0) + results.append( + AppearanceTargetResult( + target_id=target_id, + phase=phase, + frame_scores=frame_scores, + score=score, + reappearance_coverage=coverage, + n_reference_features=len(references), + n_return_features=n_return_features, + revisit_scores=tuple(revisit_scores), + revisit_coverages=tuple(revisit_coverages), + n_revisit_features=tuple(n_revisit_features), + ) + ) + return AppearanceAnalysis(target_results=tuple(results)) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/body_tracking.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/body_tracking.py new file mode 100644 index 000000000..0ff2515e3 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/body_tracking.py @@ -0,0 +1,486 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Target-specific body phases with frozen-face reacquisition and local SAM2.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Protocol + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.revisit_protocol import ( + MAX_SHORT_MISSED_FRAMES, + MIN_REAPPEAR_MATCH_RUN, + REVISIT_BOUNDARY_BUFFER_FRAMES, + _close_short_internal_gaps, + _merge_nearby_runs, + _true_runs, + infer_target_phase, + phase_from_revisit_runs, +) + +from .face_identity import ( + _frame_matches, + build_frozen_reference_libraries, +) +from .person_mask_tracker import MaskPropagator, PersonDetector +from .schemas import ( + BodyTrackAnalysis, + BoundingBox, + FaceObservation, + FaceTrackAnalysis, + HumanEntityProtocol, + TargetPhase, +) +from .target_association import assign_faces_to_person_boxes, hungarian_matches + +BODY_MASK_BOX_IOU_THRESHOLD = 0.1 +BODY_REACQUIRE_PROMPT_STRIDE = 8 +BODY_REACQUISITION_FACE_THRESHOLD = 0.35 + + +class BodyTrackerBackend(Protocol): + def analyze( + self, + protocol: HumanEntityProtocol, + decoded_frames: dict[int, np.ndarray], + *, + max_revisit_gap_frames: int, + revisit_boundary_buffer_frames: int = REVISIT_BOUNDARY_BUFFER_FRAMES, + face_tracks: FaceTrackAnalysis | None = None, + face_reacquisition_threshold: float = BODY_REACQUISITION_FACE_THRESHOLD, + ) -> BodyTrackAnalysis: ... + + +def _mask_bbox(mask: np.ndarray) -> BoundingBox | None: + y_coordinates, x_coordinates = np.where(mask) + if not len(x_coordinates): + return None + return BoundingBox( + x1=float(x_coordinates.min()), + y1=float(y_coordinates.min()), + x2=float(x_coordinates.max() + 1), + y2=float(y_coordinates.max() + 1), + ) + + +def _box_iou(first: BoundingBox, second: BoundingBox) -> float: + x1 = max(first.x1, second.x1) + y1 = max(first.y1, second.y1) + x2 = min(first.x2, second.x2) + y2 = min(first.y2, second.y2) + intersection = max(0.0, x2 - x1) * max(0.0, y2 - y1) + union = first.area + second.area - intersection + return intersection / union if union > 0.0 else 0.0 + + +def _initial_phase( + *, + target_id: int, + total_frames: int, + assignments: dict[int, dict[int, BoundingBox]], +) -> TargetPhase: + presence = [ + frame_index == 0 or target_id in assignments.get(frame_index, {}) for frame_index in range(total_frames) + ] + smoothed = _close_short_internal_gaps( + presence, + max_gap=MAX_SHORT_MISSED_FRAMES, + ) + initial_runs = _true_runs(smoothed) + visible_end = initial_runs[0][1] + visible = tuple(range(visible_end + 1)) + if visible_end == total_frames - 1: + return TargetPhase( + target_id=target_id, + visible_indices=visible, + departure_indices=(), + return_indices=(), + status="not_departed", + ) + return TargetPhase( + target_id=target_id, + visible_indices=visible, + departure_indices=tuple(range(visible_end + 1, total_frames)), + return_indices=(), + status="not_returned", + ) + + +def _stable_face_runs( + *, + target_id: int, + visible_end: int, + total_frames: int, + matches_by_frame: dict[int, dict[int, FaceObservation]], + max_revisit_gap_frames: int, +) -> list[tuple[int, int]]: + presence = [target_id in matches_by_frame.get(frame_index, {}) for frame_index in range(total_frames)] + runs = [ + (start, end) + for start, end in _true_runs(presence, start=visible_end + 1) + if end - start + 1 >= MIN_REAPPEAR_MATCH_RUN + ] + runs = _merge_nearby_runs( + runs, + max_gap=max_revisit_gap_frames, + ) + return runs + + +def _segment_bounds( + runs: Sequence[tuple[int, int]], + *, + post_start: int, + final_frame: int, +) -> list[tuple[int, int]]: + bounds = [] + for index, (start, end) in enumerate(runs): + segment_start = post_start if index == 0 else (runs[index - 1][1] + start) // 2 + 1 + segment_end = final_frame if index == len(runs) - 1 else (end + runs[index + 1][0]) // 2 + bounds.append((segment_start, segment_end)) + return bounds + + +def _sparse_prompts( + assignments: dict[int, BoundingBox], + *, + start: int, + end: int, +) -> list[tuple[int, BoundingBox]]: + frames = sorted(frame_index for frame_index in assignments if start <= frame_index <= end) + if not frames: + return [] + selected = [frames[0]] + for frame_index in frames[1:-1]: + if frame_index - selected[-1] >= BODY_REACQUIRE_PROMPT_STRIDE: + selected.append(frame_index) + if frames[-1] != selected[-1]: + selected.append(frames[-1]) + return [(frame_index, assignments[frame_index]) for frame_index in selected] + + +def _anchored_component( + candidate_frames: set[int], + *, + segment_start: int, + segment_end: int, + anchor_start: int, + anchor_end: int, + max_gap: int, +) -> tuple[int, int]: + presence = [frame_index in candidate_frames for frame_index in range(segment_start, segment_end + 1)] + smoothed = _close_short_internal_gaps(presence, max_gap=max_gap) + components = [(segment_start + start, segment_start + end) for start, end in _true_runs(smoothed)] + anchored = [(start, end) for start, end in components if not (end < anchor_start or start > anchor_end)] + if not anchored: + return anchor_start, anchor_end + return ( + min(start for start, _ in anchored), + max(end for _, end in anchored), + ) + + +class YoloSamBodyTracker: + """Use initial SAM2 presence, then face-anchored independent local tracks.""" + + def __init__( + self, + *, + person_detector: PersonDetector, + mask_propagator: MaskPropagator, + mask_box_iou_threshold: float = BODY_MASK_BOX_IOU_THRESHOLD, + ) -> None: + self.person_detector = person_detector + self.mask_propagator = mask_propagator + self.mask_box_iou_threshold = mask_box_iou_threshold + + def _initial_tracks( + self, + protocol: HumanEntityProtocol, + decoded_frames: dict[int, np.ndarray], + ) -> tuple[ + dict[int, dict[int, BoundingBox]], + dict[int, list[tuple[BoundingBox, float]]], + ]: + frame_indices = list(protocol.required_frame_indices) + frames = [decoded_frames[index] for index in frame_indices] + first_people = self.person_detector.detect(frames[0]) + people_by_frame = {0: first_people} + first_assignments = assign_faces_to_person_boxes( + {target_id: box for target_id, box in enumerate(protocol.target_face_boxes)}, + first_people, + ) + masks = self.mask_propagator.track( + frames, + {target_id: [(0, box)] for target_id, box in first_assignments.items()}, + ) + assignments = {0: first_assignments} + for local_index, frame_index in enumerate(frame_indices[1:], start=1): + people = self.person_detector.detect(decoded_frames[frame_index]) + people_by_frame[frame_index] = people + target_ids = [] + mask_boxes = [] + for target_id in range(protocol.n_targets): + mask = masks.get(target_id, {}).get(local_index) + box = _mask_bbox(mask) if mask is not None else None + if box is not None: + target_ids.append(target_id) + mask_boxes.append(box) + if not mask_boxes or not people: + assignments[frame_index] = {} + continue + scores = np.asarray( + [[_box_iou(mask_box, person_box) for person_box, _ in people] for mask_box in mask_boxes], + dtype=np.float32, + ) + matches = hungarian_matches( + scores, + threshold=self.mask_box_iou_threshold, + ) + assignments[frame_index] = {target_ids[row]: people[column][0] for row, column in matches.items()} + return assignments, people_by_frame + + def _face_reacquisition_matches( + self, + protocol: HumanEntityProtocol, + face_tracks: FaceTrackAnalysis, + initial_phases: Sequence[TargetPhase], + *, + face_reacquisition_threshold: float, + ) -> dict[int, dict[int, FaceObservation]]: + references = build_frozen_reference_libraries( + face_tracks, + initial_phases, + ) + matches_by_frame: dict[int, dict[int, FaceObservation]] = {0: {}} + for frame_index in range(1, protocol.total_frames): + active_target_ids = [ + target_id + for target_id, phase in enumerate(initial_phases) + if (phase.status != "not_departed" and frame_index > phase.visible_indices[-1]) + ] + reference_sets = [references[target_id] for target_id in active_target_ids] + matches_by_frame[frame_index] = _frame_matches( + reference_sets, + face_tracks.candidates_by_frame[frame_index], + threshold=face_reacquisition_threshold, + target_ids=active_target_ids, + ) + return matches_by_frame + + @staticmethod + def _direct_person_assignments( + *, + stable_runs: Sequence[list[tuple[int, int]]], + matches_by_frame: dict[int, dict[int, FaceObservation]], + people_by_frame: dict[int, list[tuple[BoundingBox, float]]], + total_frames: int, + ) -> list[dict[int, BoundingBox]]: + direct = [dict() for _ in stable_runs] + for frame_index in range(total_frames): + eligible_faces = {} + for target_id, runs in enumerate(stable_runs): + if any(start <= frame_index <= end for start, end in runs): + observation = matches_by_frame.get(frame_index, {}).get(target_id) + if observation is not None: + eligible_faces[target_id] = observation.bbox + if not eligible_faces: + continue + mapped = assign_faces_to_person_boxes( + eligible_faces, + people_by_frame.get(frame_index, []), + ) + for target_id, person_box in mapped.items(): + direct[target_id][frame_index] = person_box + return direct + + def _local_revisit_tracks( + self, + *, + target_id: int, + visible_end: int, + runs: Sequence[tuple[int, int]], + direct_assignments: dict[int, BoundingBox], + decoded_frames: dict[int, np.ndarray], + people_by_frame: dict[int, list[tuple[BoundingBox, float]]], + max_revisit_gap_frames: int, + ) -> tuple[dict[int, BoundingBox], list[tuple[int, int]]]: + viable_runs = [ + (start, end) + for start, end in runs + if any(start <= frame_index <= end for frame_index in direct_assignments) + ] + assignments = {} + kept_runs = [] + bounds = _segment_bounds( + viable_runs, + post_start=visible_end + 1, + final_frame=len(decoded_frames) - 1, + ) + for (anchor_start, anchor_end), (segment_start, segment_end) in zip( + viable_runs, + bounds, + ): + global_prompts = _sparse_prompts( + direct_assignments, + start=anchor_start, + end=anchor_end, + ) + segment_frames = [decoded_frames[frame_index] for frame_index in range(segment_start, segment_end + 1)] + masks = self.mask_propagator.track( + segment_frames, + {target_id: [(frame_index - segment_start, box) for frame_index, box in global_prompts]}, + ).get(target_id, {}) + candidates = { + frame_index: box + for frame_index, box in direct_assignments.items() + if anchor_start <= frame_index <= anchor_end + } + for local_index, mask in masks.items(): + frame_index = segment_start + int(local_index) + if frame_index < segment_start or frame_index > segment_end or frame_index in candidates: + continue + mask_box = _mask_bbox(np.asarray(mask, dtype=bool)) + people = people_by_frame.get(frame_index, []) + if mask_box is None or not people: + continue + scores = [_box_iou(mask_box, person_box) for person_box, _ in people] + best_index = int(np.argmax(scores)) + if scores[best_index] >= self.mask_box_iou_threshold: + candidates[frame_index] = people[best_index][0] + run_start, run_end = _anchored_component( + set(candidates), + segment_start=segment_start, + segment_end=segment_end, + anchor_start=anchor_start, + anchor_end=anchor_end, + max_gap=max_revisit_gap_frames, + ) + assignments.update( + {frame_index: box for frame_index, box in candidates.items() if run_start <= frame_index <= run_end} + ) + kept_runs.append((run_start, run_end)) + return assignments, kept_runs + + def analyze( + self, + protocol: HumanEntityProtocol, + decoded_frames: dict[int, np.ndarray], + *, + max_revisit_gap_frames: int, + revisit_boundary_buffer_frames: int = REVISIT_BOUNDARY_BUFFER_FRAMES, + face_tracks: FaceTrackAnalysis | None = None, + face_reacquisition_threshold: float = BODY_REACQUISITION_FACE_THRESHOLD, + ) -> BodyTrackAnalysis: + initial_assignments, people_by_frame = self._initial_tracks( + protocol, + decoded_frames, + ) + if face_tracks is None: + inferred = tuple( + infer_target_phase( + target_id=target_id, + total_frames=protocol.total_frames, + matched_frame_indices={ + frame_index for frame_index, matched in initial_assignments.items() if target_id in matched + }, + max_revisit_gap_frames=max_revisit_gap_frames, + expected_revisits=protocol.expected_revisits, + revisit_boundary_buffer_frames=(revisit_boundary_buffer_frames), + ) + for target_id in range(protocol.n_targets) + ) + return BodyTrackAnalysis( + target_phases=inferred, + assignments=initial_assignments, + ) + + initial_phases = tuple( + _initial_phase( + target_id=target_id, + total_frames=protocol.total_frames, + assignments=initial_assignments, + ) + for target_id in range(protocol.n_targets) + ) + matches_by_frame = self._face_reacquisition_matches( + protocol, + face_tracks, + initial_phases, + face_reacquisition_threshold=face_reacquisition_threshold, + ) + stable_runs = [ + ( + [] + if phase.status == "not_departed" + else _stable_face_runs( + target_id=target_id, + visible_end=phase.visible_indices[-1], + total_frames=protocol.total_frames, + matches_by_frame=matches_by_frame, + max_revisit_gap_frames=max_revisit_gap_frames, + ) + ) + for target_id, phase in enumerate(initial_phases) + ] + direct_assignments = self._direct_person_assignments( + stable_runs=stable_runs, + matches_by_frame=matches_by_frame, + people_by_frame=people_by_frame, + total_frames=protocol.total_frames, + ) + + final_assignments = {frame_index: {} for frame_index in range(protocol.total_frames)} + target_phases = [] + for target_id, initial_phase in enumerate(initial_phases): + visible_end = initial_phase.visible_indices[-1] + if initial_phase.status == "not_departed": + for frame_index, matched in initial_assignments.items(): + person_box = matched.get(target_id) + if person_box is not None: + final_assignments[frame_index][target_id] = person_box + target_phases.append( + phase_from_revisit_runs( + target_id=target_id, + visible_indices=initial_phase.visible_indices, + total_frames=protocol.total_frames, + runs=[], + expected_revisits=protocol.expected_revisits, + revisit_boundary_buffer_frames=(revisit_boundary_buffer_frames), + ) + ) + continue + + for frame_index in initial_phase.visible_indices: + person_box = initial_assignments.get(frame_index, {}).get(target_id) + if person_box is not None: + final_assignments[frame_index][target_id] = person_box + local_assignments, revisit_runs = self._local_revisit_tracks( + target_id=target_id, + visible_end=visible_end, + runs=stable_runs[target_id], + direct_assignments=direct_assignments[target_id], + decoded_frames=decoded_frames, + people_by_frame=people_by_frame, + max_revisit_gap_frames=max_revisit_gap_frames, + ) + for frame_index, person_box in local_assignments.items(): + final_assignments[frame_index][target_id] = person_box + target_phases.append( + phase_from_revisit_runs( + target_id=target_id, + visible_indices=initial_phase.visible_indices, + total_frames=protocol.total_frames, + runs=revisit_runs, + expected_revisits=protocol.expected_revisits, + revisit_boundary_buffer_frames=(revisit_boundary_buffer_frames), + ) + ) + + return BodyTrackAnalysis( + target_phases=tuple(target_phases), + assignments=final_assignments, + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/evaluator.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/evaluator.py new file mode 100644 index 000000000..3be65ab87 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/evaluator.py @@ -0,0 +1,314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Shared resource construction and cached Human Entity analyses.""" + +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.context import EvaluationContext +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.video import VideoReader + +from ..features import MaskedDinoV2Features, get_dino_model +from ..revisit_protocol import MAX_SHORT_MISSED_FRAMES, REVISIT_BOUNDARY_BUFFER_FRAMES +from .appearance_features import ( + AppearanceBackend, + YoloSamDinoAppearanceBackend, + analyze_appearance, +) +from .body_tracking import ( + BODY_REACQUISITION_FACE_THRESHOLD, + BodyTrackerBackend, + YoloSamBodyTracker, +) +from .face_identity import ( + TRACK_ASSOCIATION_THRESHOLD, + FaceIdentityBackend, + InsightFaceBackend, + analyze_identity, + prepare_face_tracks, +) +from .person_mask_tracker import ( + TransformersSam2MaskPropagator, + YoloWorldPersonDetector, +) +from .phase_protocol import load_human_entity_protocol +from .schemas import AppearanceAnalysis, IdentityAnalysis + +YOLO_PERSON_CONFIDENCE = 0.05 +MIN_MASK_AREA_RATIO = 0.0 +DINO_BATCH_SIZE = 16 + + +def _face_backend(context: EvaluationContext, weight_metric="human_identity") -> FaceIdentityBackend: + config = context.metric_config("human_identity") + backend_name = str(config.get("backend", "insightface")).lower() + if backend_name != "insightface": + raise EvaluationError("human_identity.backend is fixed to 'insightface' for comparable scores") + model_name = str(config.get("insightface_model", "buffalo_l")) + model_root = config.get("insightface_root") + if model_root is None: + # InsightFace accepts an absolute model name. This loads the metric's + # buffalo_l link directly without creating another models/ directory. + model_root = context.cache_dir + model_name = str((context.cache_dir / "metrics" / weight_metric / model_name).resolve()) + key = ( + "human_entity:face_backend:insightface", + str((Path(model_root).expanduser() / "models" / model_name).resolve()), + context.device, + ) + return context.get_resource( + key, + lambda: InsightFaceBackend( + model_name=model_name, + model_root=model_root, + device=context.device, + ), + ) + + +def _person_detector(context: EvaluationContext, weight_metric="human_appearance") -> YoloWorldPersonDetector: + config = context.metric_config("human_appearance") + checkpoint = str( + Path(config.get("yolo_checkpoint") or context.cache_dir / "metrics" / weight_metric / "yolov8s-worldv2.pt") + .expanduser() + .resolve() + ) + key = ( + "human_entity:person_detector", + checkpoint, + YOLO_PERSON_CONFIDENCE, + context.device, + ) + return context.get_resource( + key, + lambda: YoloWorldPersonDetector( + checkpoint=checkpoint, + clip_checkpoint=str(context.cache_dir / "metrics" / weight_metric / "ViT-B-32.pt"), + confidence=YOLO_PERSON_CONFIDENCE, + device=context.device, + ), + ) + + +def _mask_propagator( + context: EvaluationContext, + weight_metric="human_appearance", +) -> TransformersSam2MaskPropagator: + config = context.metric_config("human_appearance") + checkpoint = str( + Path(config.get("sam2_checkpoint") or context.cache_dir / "metrics" / weight_metric / "sam2-hiera-large") + .expanduser() + .resolve() + ) + key = ( + "human_entity:mask_propagator", + checkpoint, + MIN_MASK_AREA_RATIO, + context.device, + ) + return context.get_resource( + key, + lambda: TransformersSam2MaskPropagator( + checkpoint=checkpoint, + min_mask_area_ratio=MIN_MASK_AREA_RATIO, + device=context.device, + ), + ) + + +def _dino_extractor( + context: EvaluationContext, +) -> MaskedDinoV2Features: + config = context.metric_config("human_appearance") + checkpoint = str( + config.get("dino_checkpoint") or context.cache_dir / "metrics" / "human_appearance" / "dinov2-large" + ) + raw_batch_size = config.get("dino_batch_size", DINO_BATCH_SIZE) + if isinstance(raw_batch_size, bool) or not isinstance(raw_batch_size, int) or raw_batch_size < 1: + raise EvaluationError("human_appearance.dino_batch_size must be positive") + key = ( + "human_entity:dino_extractor", + checkpoint, + raw_batch_size, + context.device, + ) + return context.get_resource( + key, + lambda: MaskedDinoV2Features( + model=get_dino_model(context, checkpoint), + device=context.device, + batch_size=raw_batch_size, + ), + ) + + +def _appearance_backend(context: EvaluationContext) -> AppearanceBackend: + mask_propagator = _mask_propagator(context) + feature_extractor = _dino_extractor(context) + key = ( + "human_entity:appearance_backend", + id(mask_propagator), + id(feature_extractor), + ) + return context.get_resource( + key, + lambda: YoloSamDinoAppearanceBackend( + mask_propagator=mask_propagator, + feature_extractor=feature_extractor, + ), + ) + + +def _body_backend(context: EvaluationContext, weight_metric="human_identity") -> BodyTrackerBackend: + person_detector = _person_detector(context, weight_metric) + mask_propagator = _mask_propagator(context, weight_metric) + key = ( + "human_entity:body_backend", + id(person_detector), + id(mask_propagator), + ) + return context.get_resource( + key, + lambda: YoloSamBodyTracker( + person_detector=person_detector, + mask_propagator=mask_propagator, + ), + ) + + +def _non_negative_int_config(name: str, value: object) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise EvaluationError(f"{name} must be a non-negative integer") + return value + + +def identity_analysis_for_item( + item: dict, + context: EvaluationContext, + *, + backend: FaceIdentityBackend | None = None, + body_backend: BodyTrackerBackend | None = None, + reader: VideoReader | None = None, + weight_metric: str = "human_identity", +) -> IdentityAnalysis: + local_path = item["video"] + video_reader = reader or VideoReader() + info = video_reader.probe(local_path) + protocol = load_human_entity_protocol( + item, + total_frames=info.total_frames, + frame_width=info.width, + frame_height=info.height, + ) + selected_backend = backend or _face_backend(context, weight_metric) + selected_body_backend = body_backend or _body_backend(context, weight_metric) + identity_config = context.metric_config("human_identity") + max_revisit_gap_frames = _non_negative_int_config( + "human_identity.max_revisit_gap_frames", + identity_config.get("max_revisit_gap_frames", MAX_SHORT_MISSED_FRAMES), + ) + revisit_boundary_buffer_frames = _non_negative_int_config( + "human_identity.revisit_boundary_buffer_frames", + identity_config.get( + "revisit_boundary_buffer_frames", + REVISIT_BOUNDARY_BUFFER_FRAMES, + ), + ) + track_threshold = TRACK_ASSOCIATION_THRESHOLD + body_reacquisition_threshold = BODY_REACQUISITION_FACE_THRESHOLD + cache_key = ( + "human_entity:identity_analysis", + item["item_id"], + local_path, + id(selected_backend), + id(selected_body_backend), + track_threshold, + body_reacquisition_threshold, + max_revisit_gap_frames, + revisit_boundary_buffer_frames, + ) + + def create() -> IdentityAnalysis: + decoded = video_reader.read_indices( + local_path, + list(protocol.required_frame_indices), + ) + decoded_by_index = {record.frame_idx: record.rgb for record in decoded} + face_tracks = prepare_face_tracks( + protocol=protocol, + decoded_frames=decoded_by_index, + backend=selected_backend, + track_association_threshold=track_threshold, + ) + body = selected_body_backend.analyze( + protocol, + decoded_by_index, + max_revisit_gap_frames=max_revisit_gap_frames, + revisit_boundary_buffer_frames=revisit_boundary_buffer_frames, + face_tracks=face_tracks, + face_reacquisition_threshold=body_reacquisition_threshold, + ) + return analyze_identity( + protocol=protocol, + decoded_frames=decoded_by_index, + backend=selected_backend, + video_path=local_path, + face_tracks=face_tracks, + track_association_threshold=track_threshold, + target_phases=body.target_phases, + body_assignments=body.assignments, + max_revisit_gap_frames=max_revisit_gap_frames, + revisit_boundary_buffer_frames=revisit_boundary_buffer_frames, + ) + + return context.get_cached( + cache_key, + create, + ) + + +def appearance_analysis_for_item( + item: dict, + context: EvaluationContext, + *, + face_backend: FaceIdentityBackend | None = None, + appearance_backend: AppearanceBackend | None = None, + reader: VideoReader | None = None, + body_backend: BodyTrackerBackend | None = None, +) -> tuple[IdentityAnalysis, AppearanceAnalysis]: + identity = identity_analysis_for_item( + item, + context, + backend=face_backend, + body_backend=body_backend, + reader=reader, + weight_metric="human_appearance", + ) + video_reader = reader or VideoReader() + local_path = identity.video_path + decoded = video_reader.read_indices( + local_path, + list(identity.protocol.required_frame_indices), + ) + # Appearance currently re-decodes the full item to keep the shared cache + # free of hundreds of decoded RGB frames between metric evaluations. + identity_with_frames = replace( + identity, + decoded_frames={record.frame_idx: record.rgb for record in decoded}, + ) + selected_backend = appearance_backend or _appearance_backend(context) + cache_key = ( + "human_entity:appearance_analysis", + item["item_id"], + identity.video_path, + id(selected_backend), + ) + appearance = context.get_cached( + cache_key, + lambda: analyze_appearance(identity_with_frames, selected_backend), + ) + return identity, appearance diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/face_identity.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/face_identity.py new file mode 100644 index 000000000..66e54e191 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/face_identity.py @@ -0,0 +1,400 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Face inference, seeded phase tracking, and frozen-target identity scoring.""" + +from __future__ import annotations + +import os +from collections.abc import Sequence +from pathlib import Path +from typing import Protocol + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import ( + DataInvalidError, + EvaluationError, +) + +from ..revisit_protocol import ( + MAX_SHORT_MISSED_FRAMES, + REVISIT_BOUNDARY_BUFFER_FRAMES, + infer_target_phase, +) +from .schemas import ( + BoundingBox, + FaceObservation, + FaceTrackAnalysis, + HumanEntityProtocol, + IdentityAnalysis, + IdentityTargetResult, + TargetPhase, +) +from .target_association import cosine_matrix, hungarian_matches, normalized_embedding + +TRACK_ASSOCIATION_THRESHOLD = 0.4 + + +class FaceIdentityBackend(Protocol): + def embed_box(self, frame_rgb: np.ndarray, bbox: BoundingBox) -> np.ndarray: ... + + def detect(self, frame_rgb: np.ndarray) -> list[FaceObservation]: ... + + +class InsightFaceBackend: + """SCRFD detection and ArcFace embeddings from InsightFace buffalo_l.""" + + def __init__( + self, + *, + device: str, + model_name: str = "buffalo_l", + model_root: str | Path, + detection_size: int = 640, + ) -> None: + try: + import onnxruntime as ort + from insightface.app import FaceAnalysis + except ImportError as error: + raise EvaluationError( + "InsightFace backend requires insightface and onnxruntime; " + "install the WorldMemBench metric requirements" + ) from error + + root = Path(model_root).expanduser().resolve() + model_dir = root / "models" / model_name + if not all((model_dir / name).is_file() for name in ("det_10g.onnx", "w600k_r50.onnx")): + raise EvaluationError( + f"InsightFace model directory not found: {model_dir}; " + "WorldMemBench never downloads evaluator weights implicitly" + ) + if device.startswith("cuda"): + # Changing LD_LIBRARY_PATH after Python starts does not reliably + # affect dlopen. ORT's loader explicitly preloads CUDA/cuDNN from + # the active environment's nvidia packages. If the caller already + # supplied a cuDNN path, avoid loading the same libraries twice. + preload = getattr(ort, "preload_dlls", None) + if preload is not None and "cudnn" not in os.environ.get("LD_LIBRARY_PATH", ""): + preload(msvc=False, directory="") + raw_device = device.partition(":")[2] + device_id = int(raw_device) if raw_device else 0 + providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] + else: + device_id = -1 + providers = ["CPUExecutionProvider"] + self.app = FaceAnalysis( + name=model_name, + root=str(root), + allowed_modules=["detection", "recognition"], + providers=providers, + # Enable TF32 for CUDA face inference. + provider_options=([{"device_id": device_id, "use_tf32": 1}, {}] if device_id >= 0 else [{}]), + ) + self.app.prepare( + ctx_id=device_id, + det_size=(detection_size, detection_size), + ) + if device_id >= 0: + for model in self.app.models.values(): + if "CUDAExecutionProvider" not in model.session.get_providers(): + raise EvaluationError( + "InsightFace could not activate CUDAExecutionProvider; check ONNX Runtime CUDA/cuDNN libraries" + ) + self._cached_frame: np.ndarray | None = None + self._cached_observations: list[FaceObservation] = [] + + def _observations(self, frame_rgb: np.ndarray) -> list[FaceObservation]: + if frame_rgb is self._cached_frame: + return list(self._cached_observations) + import cv2 + + faces = self.app.get(cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)) + observations: list[FaceObservation] = [] + for face in faces: + embedding = getattr(face, "normed_embedding", None) + raw_box = getattr(face, "bbox", None) + if embedding is None or raw_box is None: + continue + try: + bbox = BoundingBox.from_value( + np.asarray(raw_box).tolist(), + field_name="InsightFace bbox", + ) + normalized = normalized_embedding(embedding) + except (ValueError, EvaluationError): + continue + observations.append(FaceObservation(bbox=bbox, embedding=normalized)) + self._cached_frame = frame_rgb + self._cached_observations = observations + return list(observations) + + @staticmethod + def _box_iou(first: BoundingBox, second: BoundingBox) -> float: + x1 = max(first.x1, second.x1) + y1 = max(first.y1, second.y1) + x2 = min(first.x2, second.x2) + y2 = min(first.y2, second.y2) + intersection = max(0.0, x2 - x1) * max(0.0, y2 - y1) + union = first.area + second.area - intersection + return intersection / union if union > 0.0 else 0.0 + + def embed_box(self, frame_rgb: np.ndarray, bbox: BoundingBox) -> np.ndarray: + observations = self._observations(frame_rgb) + if not observations: + raise DataInvalidError("InsightFace found no face in the first frame") + match = max( + observations, + key=lambda observation: self._box_iou(observation.bbox, bbox), + ) + if self._box_iou(match.bbox, bbox) < 0.1: + raise DataInvalidError("target face bbox does not overlap an InsightFace detection") + return match.embedding + + def detect(self, frame_rgb: np.ndarray) -> list[FaceObservation]: + return self._observations(frame_rgb) + + +def _frame_matches( + reference_sets: Sequence[Sequence[np.ndarray]], + candidates: Sequence[FaceObservation], + *, + threshold: float, + target_ids: Sequence[int] | None = None, +) -> dict[int, FaceObservation]: + if target_ids is None: + target_ids = tuple(range(len(reference_sets))) + if len(target_ids) != len(reference_sets): + raise ValueError("target_ids and reference_sets must have the same length") + if not candidates or not reference_sets: + return {} + candidate_embeddings = [candidate.embedding for candidate in candidates] + score_rows: list[np.ndarray] = [] + for references in reference_sets: + similarities = cosine_matrix(references, candidate_embeddings) + score_rows.append(similarities.max(axis=0)) + scores = np.stack(score_rows, axis=0) + matches = hungarian_matches(scores, threshold=threshold) + return {target_ids[row_index]: candidates[candidate_index] for row_index, candidate_index in matches.items()} + + +def _seeded_track_associations( + *, + anchors: Sequence[np.ndarray], + seeded_association: dict[int, FaceObservation], + candidates_by_frame: dict[int, list[FaceObservation]], + total_frames: int, + threshold: float, +) -> dict[int, dict[int, FaceObservation]]: + """MBench-style rolling tracks seeded only by the frozen frame-0 faces.""" + + track_embeddings: list[list[np.ndarray]] = [[anchor] for anchor in anchors] + missed_frames = [0 for _ in anchors] + references_frozen = [False for _ in anchors] + associations: dict[int, dict[int, FaceObservation]] = {0: seeded_association} + for frame_index in range(1, total_frames): + reference_sets = [[embeddings[-1], np.mean(embeddings, axis=0)] for embeddings in track_embeddings] + matches = _frame_matches( + reference_sets, + candidates_by_frame[frame_index], + threshold=threshold, + ) + associations[frame_index] = matches + for target_id in range(len(anchors)): + observation = matches.get(target_id) + if observation is None: + if not references_frozen[target_id]: + missed_frames[target_id] += 1 + if missed_frames[target_id] > MAX_SHORT_MISSED_FRAMES: + references_frozen[target_id] = True + continue + missed_frames[target_id] = 0 + if references_frozen[target_id]: + continue + track_embeddings[target_id].append(normalized_embedding(observation.embedding)) + return associations + + +def prepare_face_tracks( + *, + protocol: HumanEntityProtocol, + decoded_frames: dict[int, np.ndarray], + backend: FaceIdentityBackend, + track_association_threshold: float = TRACK_ASSOCIATION_THRESHOLD, +) -> FaceTrackAnalysis: + if not 0.0 <= track_association_threshold <= 1.0: + raise ValueError("track_association_threshold must be in [0, 1]") + first_frame = decoded_frames.get(0) + if first_frame is None: + raise EvaluationError("the decoded frame set does not contain frame 0") + anchors = tuple(normalized_embedding(backend.embed_box(first_frame, bbox)) for bbox in protocol.target_face_boxes) + if len({anchor.shape for anchor in anchors}) != 1: + raise EvaluationError("ArcFace anchors have inconsistent dimensions") + seeded_association = { + target_id: FaceObservation(box, anchor) + for target_id, (box, anchor) in enumerate(zip(protocol.target_face_boxes, anchors)) + } + candidates_by_frame: dict[int, list[FaceObservation]] = {} + for frame_index in range(1, protocol.total_frames): + try: + frame = decoded_frames[frame_index] + except KeyError as error: + raise EvaluationError(f"the decoded frame set does not contain frame {frame_index}") from error + candidates_by_frame[frame_index] = backend.detect(frame) + track_associations = _seeded_track_associations( + anchors=anchors, + seeded_association=seeded_association, + candidates_by_frame=candidates_by_frame, + total_frames=protocol.total_frames, + threshold=track_association_threshold, + ) + return FaceTrackAnalysis( + anchors=anchors, + seeded_association=seeded_association, + candidates_by_frame=candidates_by_frame, + track_associations=track_associations, + ) + + +def build_frozen_reference_libraries( + face_tracks: FaceTrackAnalysis, + target_phases: Sequence[TargetPhase], +) -> tuple[tuple[np.ndarray, ...], ...]: + references: list[list[np.ndarray]] = [[anchor] for anchor in face_tracks.anchors] + for phase in target_phases: + for frame_index in phase.visible_indices: + if frame_index == 0: + continue + observation = face_tracks.track_associations[frame_index].get(phase.target_id) + if observation is not None: + references[phase.target_id].append(normalized_embedding(observation.embedding)) + return tuple(tuple(values) for values in references) + + +def _face_phases_from_associations( + *, + protocol: HumanEntityProtocol, + associations: dict[int, dict[int, FaceObservation]], + max_revisit_gap_frames: int, + revisit_boundary_buffer_frames: int, +) -> tuple[TargetPhase, ...]: + return tuple( + infer_target_phase( + target_id=target_id, + total_frames=protocol.total_frames, + matched_frame_indices={ + frame_index for frame_index, matches in associations.items() if target_id in matches + }, + max_revisit_gap_frames=max_revisit_gap_frames, + expected_revisits=protocol.expected_revisits, + revisit_boundary_buffer_frames=revisit_boundary_buffer_frames, + ) + for target_id in range(protocol.n_targets) + ) + + +def analyze_identity( + *, + protocol: HumanEntityProtocol, + face_tracks: FaceTrackAnalysis | None = None, + decoded_frames: dict[int, np.ndarray], + backend: FaceIdentityBackend, + video_path: str, + target_phases: tuple[TargetPhase, ...], + body_assignments: dict[int, dict[int, BoundingBox]], + track_association_threshold: float = TRACK_ASSOCIATION_THRESHOLD, + max_revisit_gap_frames: int = MAX_SHORT_MISSED_FRAMES, + revisit_boundary_buffer_frames: int = REVISIT_BOUNDARY_BUFFER_FRAMES, +) -> IdentityAnalysis: + if not 0.0 <= track_association_threshold <= 1.0: + raise ValueError("track_association_threshold must be in [0, 1]") + if face_tracks is None: + face_tracks = prepare_face_tracks( + protocol=protocol, + decoded_frames=decoded_frames, + backend=backend, + track_association_threshold=track_association_threshold, + ) + appearance_target_phases = target_phases + target_phases = _face_phases_from_associations( + protocol=protocol, + associations=face_tracks.track_associations, + max_revisit_gap_frames=max_revisit_gap_frames, + revisit_boundary_buffer_frames=revisit_boundary_buffer_frames, + ) + # Identity and appearance infer D independently. Only the face-derived V + # may extend the frozen identity library; return observations never do. + frozen_references = build_frozen_reference_libraries( + face_tracks, + target_phases, + ) + associations = face_tracks.track_associations + + results: list[IdentityTargetResult] = [] + for phase in target_phases: + target_id = phase.target_id + + return_indices = phase.return_indices + frame_scores = {} + returned_embeddings: list[np.ndarray] = [] + returned_by_frame: dict[int, np.ndarray] = {} + for frame_index in return_indices: + observation = associations.get(frame_index, {}).get(target_id) + if observation is None: + frame_scores[frame_index] = 0.0 + continue + embedding = normalized_embedding(observation.embedding) + similarities = cosine_matrix( + frozen_references[target_id], + [embedding], + )[:, 0] + frame_scores[frame_index] = float(np.clip(similarities, 0.0, 1.0).max()) + returned_embeddings.append(embedding) + returned_by_frame[frame_index] = embedding + coverage = len(returned_embeddings) / len(return_indices) if return_indices else 0.0 + if returned_embeddings and phase.is_scorable: + cross = cosine_matrix(frozen_references[target_id], returned_embeddings) + best_by_embedding = np.clip(cross, 0.0, 1.0).max(axis=0) + similarity = float(best_by_embedding.mean()) + else: + similarity = 0.0 + revisit_scores: list[float] = [] + revisit_coverages: list[float] = [] + n_revisit_embeddings: list[int] = [] + for revisit_indices in phase.revisit_indices: + embeddings = [returned_by_frame[index] for index in revisit_indices if index in returned_by_frame] + n_revisit_embeddings.append(len(embeddings)) + revisit_coverages.append(len(embeddings) / len(revisit_indices) if revisit_indices else 0.0) + if embeddings and phase.is_scorable: + cross = cosine_matrix(frozen_references[target_id], embeddings) + best_by_embedding = np.clip(cross, 0.0, 1.0).max(axis=0) + revisit_scores.append(float(best_by_embedding.mean())) + else: + revisit_scores.append(0.0) + results.append( + IdentityTargetResult( + target_id=target_id, + phase=phase, + frame_scores=frame_scores, + similarity=similarity, + coverage=coverage, + n_reference_embeddings=len(frozen_references[target_id]), + n_return_embeddings=len(returned_embeddings), + revisit_scores=tuple(revisit_scores), + revisit_coverages=tuple(revisit_coverages), + n_revisit_embeddings=tuple(n_revisit_embeddings), + ) + ) + return IdentityAnalysis( + protocol=protocol, + target_phases=target_phases, + appearance_target_phases=appearance_target_phases, + body_assignments=body_assignments, + target_results=tuple(results), + # Keep the cross-metric cache lightweight. Appearance re-decodes the + # full item and attaches frames to a transient analysis copy. + decoded_frames={}, + video_path=video_path, + max_revisit_gap_frames=max_revisit_gap_frames, + revisit_boundary_buffer_frames=revisit_boundary_buffer_frames, + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/metric.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/metric.py new file mode 100644 index 000000000..f2daff208 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/metric.py @@ -0,0 +1,565 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Human Identity and Human Appearance metric plugins.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.context import EvaluationContext +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import ( + DataInvalidError, + EvaluationError, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.aggregation import ( + aggregate_video_summary, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.base import Metric +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.results import ( + ItemResult, + ResultStatus, + UnitResult, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.video import VideoReader + +if TYPE_CHECKING: + from .appearance_features import AppearanceBackend + from .body_tracking import BodyTrackerBackend + from .face_identity import FaceIdentityBackend +from ..revisit_protocol import ( + MAX_EXPECTED_REVISITS, + MIN_REAPPEAR_MATCH_RUN, + validate_expected_revisits, +) +from .schemas import IdentityAnalysis, TargetPhase + +_SHARED_ITEM_METADATA = ( + "video_path", + "n_targets", + "target_face_bboxes_first_frame", + "target_descriptions", + "expected_revisits", + "revisit_mode", + "first_frame_match_status", + "first_frame_match_error", + "phase_source", + "phase_inference", +) +_SHARED_TARGET_METADATA = ( + "target_id", + "phase_status", + "V", + "D", + "R", + "revisits", +) + + +def _identity_metadata( + identity: IdentityAnalysis, + *, + phase_source: str, +) -> dict: + from .body_tracking import BODY_MASK_BOX_IOU_THRESHOLD, BODY_REACQUISITION_FACE_THRESHOLD + from .face_identity import TRACK_ASSOCIATION_THRESHOLD + + protocol = identity.protocol + return { + "video_path": identity.video_path, + "n_targets": protocol.n_targets, + "target_face_bboxes_first_frame": [box.as_list() for box in protocol.target_face_boxes], + "target_descriptions": list(protocol.target_descriptions), + "expected_revisits": protocol.expected_revisits, + "revisit_mode": "auto" if protocol.expected_revisits is None else "fixed", + "phase_source": phase_source, + "phase_inference": { + "body_mask_box_iou_threshold": BODY_MASK_BOX_IOU_THRESHOLD, + "body_presence_model": "first_frame_face_to_yolo_person_then_sam2_track", + "track_association_threshold": TRACK_ASSOCIATION_THRESHOLD, + "body_reacquisition_face_threshold": (BODY_REACQUISITION_FACE_THRESHOLD), + "max_revisit_gap_frames": identity.max_revisit_gap_frames, + "revisit_boundary_buffer_frames": (identity.revisit_boundary_buffer_frames), + "identity_revisit_model": "rolling_face_presence_frozen_v_cosine", + "revisit_partition": ( + "observed_return_components" + if protocol.expected_revisits is None + else "boundary_buffer_d_insertion_components" + ), + "min_reappear_match_run": MIN_REAPPEAR_MATCH_RUN, + "continuous_component_split": "balanced_full_component", + }, + } + + +def _phase_metadata(phase: TargetPhase) -> dict: + return { + "phase_status": phase.status, + "V": list(phase.visible_indices), + "D": list(phase.departure_indices), + "R": list(phase.return_indices), + } + + +def _revisit_score_fields( + scores: tuple[float, ...], + coverages: tuple[float, ...], + *, + expected_revisits: int | None, +) -> dict[str, float]: + # Automatic mode has a target-dependent number of segments. Keep those + # scores in per_target.revisits, not misaligned video-level scalar fields. + if expected_revisits is None: + return {} + fields: dict[str, float] = {} + for index in range(MAX_EXPECTED_REVISITS): + revisit_id = index + 1 + fields[f"revisit_{revisit_id}_score"] = scores[index] if index < len(scores) else 0.0 + fields[f"revisit_{revisit_id}_coverage"] = coverages[index] if index < len(coverages) else 0.0 + return fields + + +def _official_revisit_score( + scores: tuple[float, ...], + *, + expected_revisits: int | None, +) -> float: + if expected_revisits is None: + return sum(scores) / len(scores) if scores else 0.0 + return sum(scores[:expected_revisits]) / expected_revisits + + +def _revisit_metadata( + phase: TargetPhase, + scores: tuple[float, ...], + coverages: tuple[float, ...], + counts: tuple[int, ...], + *, + count_key: str, + expected_revisits: int | None, +) -> list[dict]: + output: list[dict] = [] + for index in range(len(scores) if expected_revisits is None else MAX_EXPECTED_REVISITS): + indices = phase.revisit_indices[index] if index < len(phase.revisit_indices) else () + slot = phase.revisit_slots[index] if index < len(phase.revisit_slots) else () + boundary = phase.revisit_nominal_boundaries[index] if index < len(phase.revisit_nominal_boundaries) else None + search_window = phase.revisit_search_windows[index] if index < len(phase.revisit_search_windows) else () + split_source = phase.revisit_sources[index] if index < len(phase.revisit_sources) else "missing" + departure_evidence = ( + phase.revisit_departure_evidence[index] if index < len(phase.revisit_departure_evidence) else False + ) + output.append( + { + "revisit_id": index + 1, + "indices": list(indices), + "slot": list(slot), + "nominal_boundary": boundary, + "search_window": list(search_window), + "actual_window": ([indices[0], indices[-1] + 1] if indices else []), + "window_source": split_source, + "departure_evidence": departure_evidence, + "score": scores[index] if index < len(scores) else 0.0, + "coverage": coverages[index] if index < len(coverages) else 0.0, + count_key: counts[index] if index < len(counts) else 0, + } + ) + return output + + +def _first_frame_match( + item: dict, +) -> tuple[str, str | None]: + preparation_error = item.get("human_preparation_error") + if preparation_error: + error_type = DataInvalidError if preparation_error["status"] == "data_invalid" else EvaluationError + raise error_type(preparation_error["message"]) + raw_status = item.get("annotations", {}).get("first_frame_match_status") + if raw_status not in {"ok", "failed"}: + raise DataInvalidError( + "annotations.first_frame_match_status must be 'ok' or 'failed'; " + "provide the source case files and run eval to prepare annotations" + ) + raw_error = item.get("annotations", {}).get("first_frame_match_error") + if raw_error is not None and not isinstance(raw_error, str): + raise DataInvalidError("annotations.first_frame_match_error must be a string or null") + error = raw_error.strip() if isinstance(raw_error, str) else None + if raw_status == "failed" and not error: + raise DataInvalidError( + "failed first_frame_match_status requires a non-blank annotations.first_frame_match_error" + ) + return raw_status, error + + +def _zero_revisits(*, count_key: str) -> list[dict]: + return [ + { + "revisit_id": index + 1, + "indices": [], + "slot": [], + "nominal_boundary": None, + "search_window": [], + "actual_window": [], + "window_source": "first_frame_match_failed", + "departure_evidence": False, + "score": 0.0, + "coverage": 0.0, + count_key: 0, + } + for index in range(MAX_EXPECTED_REVISITS) + ] + + +def _zero_units_for_first_frame_failure( + item: dict, + *, + metric_name: str, + unit_type: str, + prefix: str, + diagnostic_scores: tuple[str, str], + count_metadata: tuple[str, str], + error: str, +) -> list[UnitResult]: + boxes = item.get("annotations", {}).get("target_face_bboxes_first_frame") + descriptions = item.get("annotations", {}).get("target_descriptions") + expected_revisits = validate_expected_revisits(item.get("annotations", {}).get("expected_revisits")) + if not isinstance(boxes, (list, tuple)) or not boxes: + raise DataInvalidError("annotations.target_face_bboxes_first_frame must be non-empty") + if not isinstance(descriptions, (list, tuple)) or len(descriptions) != len(boxes): + raise DataInvalidError("annotations.target_descriptions must align with target face boxes") + video_path = item["video"] + shared = { + "video_path": video_path, + "n_targets": len(boxes), + "target_face_bboxes_first_frame": list(boxes), + "target_descriptions": [str(value) for value in descriptions], + "expected_revisits": expected_revisits, + "revisit_mode": "auto" if expected_revisits is None else "fixed", + "first_frame_match_status": "failed", + "first_frame_match_error": error, + "phase_source": "first_frame_match_failure", + "phase_inference": { + "zero_score_policy": "whole_model_case_trajectory", + "reason": error, + }, + } + zero_scores = { + "score": 0.0, + diagnostic_scores[0]: 0.0, + diagnostic_scores[1]: 0.0, + **_revisit_score_fields((), (), expected_revisits=expected_revisits), + } + return [ + UnitResult( + unit_type=unit_type, + unit_id=f"target_{target_id:04d}", + scores=dict(zero_scores), + status=ResultStatus.OK, + metadata={ + **shared, + "target_id": target_id, + "phase_status": "first_frame_match_failed", + "V": [], + "D": [], + "R": [], + f"{prefix}_frame_scores": {}, + count_metadata[0]: 0, + count_metadata[1]: 0, + "revisits": _zero_revisits(count_key=count_metadata[1]) if expected_revisits is not None else [], + }, + ) + for target_id in range(len(boxes)) + ] + + +def _aggregate_units( + item: dict, + units: list[UnitResult], + *, + metric_name: str, + prefix: str, + diagnostic_scores: tuple[str, ...], + count_metadata: tuple[str, ...], +) -> ItemResult: + if not units: + raise EvaluationError(f"{metric_name} produced no target units") + non_scored = [unit for unit in units if not unit.status.is_scored] + if non_scored: + first = non_scored[0] + return ItemResult( + **{key: item.get(key) for key in ("item_id", "model_id", "category", "sample_id", "condition_id")}, + status=first.status, + error=first.error, + metadata=first.metadata, + ) + target_ids = sorted(int(unit.metadata["target_id"]) for unit in units) + if target_ids != list(range(len(units))): + raise EvaluationError(f"{metric_name} target units are incomplete") + + scores = {key: sum(unit.scores[key] for unit in units) / len(units) for key in units[0].scores} + metadata = {key: units[0].metadata[key] for key in _SHARED_ITEM_METADATA} + metadata["per_target"] = [ + { + **{key: unit.metadata[key] for key in _SHARED_TARGET_METADATA}, + f"{prefix}_score": unit.scores["score"], + **{key: unit.scores[key] for key in diagnostic_scores}, + **{key: unit.metadata[key] for key in count_metadata}, + } + for unit in units + ] + return ItemResult( + **{key: item.get(key) for key in ("item_id", "model_id", "category", "sample_id", "condition_id")}, + scores=scores, + num_units=len(units), + status=ResultStatus.OK, + metadata=metadata, + ) + + +class HumanIdentityMetric(Metric): + aggregate_summary = staticmethod(aggregate_video_summary) + + name = "human_identity" + description = ( + "InsightFace first-frame-seeded Hungarian face presence with " + "boundary-buffer departure insertion, frozen ArcFace references, " + "revisit splits, per-revisit official macro, and coverage." + ) + unit_type = "target_identity" + + def __init__( + self, + *, + backend: FaceIdentityBackend | None = None, + body_backend: BodyTrackerBackend | None = None, + reader: VideoReader | None = None, + ) -> None: + self.backend = backend + self.body_backend = body_backend + self.reader = reader or VideoReader() + + def evaluate( + self, + item: dict, + context: EvaluationContext, + ) -> list[UnitResult]: + if item.get("category") not in (None, "human"): + return [ + UnitResult( + unit_id="not_applicable", + status=ResultStatus.NOT_APPLICABLE, + metadata={"reason": "human metrics apply to human cases"}, + ) + ] + match_status, match_error = _first_frame_match(item) + if match_status == "failed": + assert match_error is not None + return _zero_units_for_first_frame_failure( + item, + metric_name=self.name, + unit_type=self.unit_type, + prefix="identity", + diagnostic_scores=("identity_similarity", "identity_coverage"), + count_metadata=( + "n_reference_embeddings", + "n_return_embeddings", + ), + error=match_error, + ) + from .evaluator import identity_analysis_for_item + + identity = identity_analysis_for_item( + item, + context, + backend=self.backend, + body_backend=self.body_backend, + reader=self.reader, + ) + shared = _identity_metadata( + identity, + phase_source="rolling_face_tracks", + ) + shared.update( + { + "first_frame_match_status": match_status, + "first_frame_match_error": match_error, + } + ) + return [ + UnitResult( + unit_type=self.unit_type, + unit_id=f"target_{result.target_id:04d}", + scores={ + "score": _official_revisit_score( + result.revisit_scores, + expected_revisits=identity.protocol.expected_revisits, + ), + "identity_similarity": result.similarity, + "identity_coverage": result.coverage, + **_revisit_score_fields( + result.revisit_scores, + result.revisit_coverages, + expected_revisits=identity.protocol.expected_revisits, + ), + }, + status=ResultStatus.OK, + metadata={ + **shared, + "target_id": result.target_id, + **_phase_metadata(result.phase), + "identity_frame_scores": {str(index): score for index, score in result.frame_scores.items()}, + "n_reference_embeddings": result.n_reference_embeddings, + "n_return_embeddings": result.n_return_embeddings, + "revisits": _revisit_metadata( + result.phase, + result.revisit_scores, + result.revisit_coverages, + result.n_revisit_embeddings, + count_key="n_return_embeddings", + expected_revisits=identity.protocol.expected_revisits, + ), + }, + ) + for result in identity.target_results + ] + + def aggregate_item( + self, + item: dict, + units: list[UnitResult], + ) -> ItemResult: + return _aggregate_units( + item, + units, + metric_name=self.name, + prefix="identity", + diagnostic_scores=("identity_similarity", "identity_coverage"), + count_metadata=("n_reference_embeddings", "n_return_embeddings"), + ) + + +class HumanAppearanceMetric(Metric): + aggregate_summary = staticmethod(aggregate_video_summary) + + name = "human_appearance" + description = ( + "Seeded-body presence with boundary-buffer departure insertion, " + "phase-local YOLO-World/SAM2 person masks, component-level revisit " + "splits, per-revisit official macro, and DINOv2-Large similarity." + ) + unit_type = "target_appearance" + + def __init__( + self, + *, + face_backend: FaceIdentityBackend | None = None, + body_backend: BodyTrackerBackend | None = None, + appearance_backend: AppearanceBackend | None = None, + reader: VideoReader | None = None, + ) -> None: + self.face_backend = face_backend + self.body_backend = body_backend + self.appearance_backend = appearance_backend + self.reader = reader or VideoReader() + + def evaluate( + self, + item: dict, + context: EvaluationContext, + ) -> list[UnitResult]: + if item.get("category") not in (None, "human"): + return [ + UnitResult( + unit_id="not_applicable", + status=ResultStatus.NOT_APPLICABLE, + metadata={"reason": "human metrics apply to human cases"}, + ) + ] + match_status, match_error = _first_frame_match(item) + if match_status == "failed": + assert match_error is not None + return _zero_units_for_first_frame_failure( + item, + metric_name=self.name, + unit_type=self.unit_type, + prefix="appearance", + diagnostic_scores=( + "appearance_similarity", + "reappearance_coverage", + ), + count_metadata=( + "n_reference_features", + "n_return_features", + ), + error=match_error, + ) + from .evaluator import appearance_analysis_for_item + + identity, appearance = appearance_analysis_for_item( + item, + context, + face_backend=self.face_backend, + body_backend=self.body_backend, + appearance_backend=self.appearance_backend, + reader=self.reader, + ) + shared = _identity_metadata( + identity, + phase_source="seeded_body_tracks", + ) + shared.update( + { + "first_frame_match_status": match_status, + "first_frame_match_error": match_error, + } + ) + return [ + UnitResult( + unit_type=self.unit_type, + unit_id=f"target_{result.target_id:04d}", + scores={ + "score": _official_revisit_score( + result.revisit_scores, + expected_revisits=identity.protocol.expected_revisits, + ), + "appearance_similarity": result.score, + "reappearance_coverage": result.reappearance_coverage, + **_revisit_score_fields( + result.revisit_scores, + result.revisit_coverages, + expected_revisits=identity.protocol.expected_revisits, + ), + }, + status=ResultStatus.OK, + metadata={ + **shared, + "target_id": result.target_id, + **_phase_metadata(result.phase), + "appearance_frame_scores": {str(index): score for index, score in result.frame_scores.items()}, + "n_reference_features": result.n_reference_features, + "n_return_features": result.n_return_features, + "revisits": _revisit_metadata( + result.phase, + result.revisit_scores, + result.revisit_coverages, + result.n_revisit_features, + count_key="n_return_features", + expected_revisits=identity.protocol.expected_revisits, + ), + }, + ) + for result in appearance.target_results + ] + + def aggregate_item( + self, + item: dict, + units: list[UnitResult], + ) -> ItemResult: + return _aggregate_units( + item, + units, + metric_name=self.name, + prefix="appearance", + diagnostic_scores=("appearance_similarity", "reappearance_coverage"), + count_metadata=("n_reference_features", "n_return_features"), + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/person_mask_tracker.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/person_mask_tracker.py new file mode 100644 index 000000000..2633cfc90 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/person_mask_tracker.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""YOLO-World person anchors and phase-local SAM2 video propagation.""" + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Any, Protocol + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + +from .schemas import BoundingBox + + +class PersonDetector(Protocol): + def detect(self, frame_rgb: np.ndarray) -> list[tuple[BoundingBox, float]]: ... + + +class MaskPropagator(Protocol): + def track( + self, + frames_rgb: Sequence[np.ndarray], + prompts: dict[int, list[tuple[int, BoundingBox]]], + *, + propagate: bool = True, + ) -> dict[int, dict[int, np.ndarray]]: ... + + +class YoloWorldPersonDetector: + """Ultralytics YOLO-World detector restricted to the ``person`` class.""" + + def __init__( + self, + *, + checkpoint: str, + clip_checkpoint: str, + device: str, + confidence: float = 0.05, + ) -> None: + try: + from ultralytics import YOLOWorld + from ultralytics.nn.text_model import CLIP + except ImportError as error: + raise EvaluationError( + "Human Appearance requires ultralytics YOLO-World; install the WorldMemBench metric requirements" + ) from error + checkpoint_path = Path(checkpoint).expanduser().resolve() + if not checkpoint_path.is_file(): + raise EvaluationError(f"YOLO-World checkpoint not found: {checkpoint_path}") + clip_path = Path(clip_checkpoint).expanduser().resolve() + if not clip_path.is_file(): + raise EvaluationError(f"YOLO-World CLIP checkpoint not found: {clip_path}") + self.model = YOLOWorld(str(checkpoint_path)) + # The upstream CLIP wrapper forwards `size` to clip.load, which also + # accepts a local file. Seed its cache so set_classes never downloads. + self.model.model.clip_model = CLIP( + str(clip_path), + device=next(self.model.model.parameters()).device, + ) + self.model.set_classes(["person"]) + self.device = device + self.confidence = confidence + + def detect(self, frame_rgb: np.ndarray) -> list[tuple[BoundingBox, float]]: + results = self.model.predict( + source=frame_rgb, + classes=[0], + conf=self.confidence, + device=self.device, + verbose=False, + ) + detections: list[tuple[BoundingBox, float]] = [] + if not results: + return detections + boxes = getattr(results[0], "boxes", None) + if boxes is None: + return detections + coordinates = boxes.xyxy.detach().float().cpu().numpy() + confidences = boxes.conf.detach().float().cpu().numpy() + for raw_box, confidence in zip(coordinates, confidences): + try: + box = BoundingBox.from_value(raw_box.tolist(), field_name="YOLO person") + except ValueError: + continue + detections.append((box, float(confidence))) + return detections + + +class TransformersSam2MaskPropagator: + """SAM2 video tracking in both temporal directions within one phase.""" + + def __init__( + self, + *, + checkpoint: str, + device: str, + min_mask_area_ratio: float = 0.0, + ) -> None: + try: + from transformers import Sam2VideoModel, Sam2VideoProcessor + except ImportError as error: + raise EvaluationError( + "Human Appearance requires transformers with Sam2VideoModel; " + "install the WorldMemBench metric requirements" + ) from error + checkpoint_path = Path(checkpoint).expanduser().resolve() + if not checkpoint_path.exists(): + raise EvaluationError(f"SAM2 checkpoint not found: {checkpoint_path}") + self.processor = Sam2VideoProcessor.from_pretrained( + checkpoint_path, + local_files_only=True, + ) + self.model = ( + Sam2VideoModel.from_pretrained( + checkpoint_path, + local_files_only=True, + ) + .to(device) + .eval() + ) + self.device = device + self.min_mask_area_ratio = min_mask_area_ratio + + def track( + self, + frames_rgb: Sequence[np.ndarray], + prompts: dict[int, list[tuple[int, BoundingBox]]], + *, + propagate: bool = True, + ) -> dict[int, dict[int, np.ndarray]]: + import torch + from PIL import Image + + if not frames_rgb or not prompts: + return {} + height, width = frames_rgb[0].shape[:2] + pil_frames = [Image.fromarray(frame) for frame in frames_rgb] + session = self.processor.init_video_session( + video=pil_frames, + inference_device=self.device, + # Evicted image features are recomputed on demand. Temporal memory, + # conditioning masks, frame order and numeric precision are unchanged. + max_vision_features_cache_size=1, + ) + known_ids = sorted(prompts) + tracks: dict[int, dict[int, np.ndarray]] = {target_id: {} for target_id in known_ids} + + def store_output( + output: Any, + *, + prompted_target_id: int | None = None, + ) -> None: + masks = getattr(output, "pred_masks", None) + if masks is None: + return + processed = self.processor.post_process_masks( + [masks], + original_sizes=[[height, width]], + binarize=True, + )[0] + binary = processed.detach().cpu().numpy() + if binary.ndim == 4 and binary.shape[1] == 1: + binary = binary[:, 0] + output_ids = list(session.obj_ids) + for index, target_id in enumerate(output_ids): + target_id = int(target_id) + if ( + index >= binary.shape[0] + or target_id not in tracks + or (prompted_target_id is not None and target_id != prompted_target_id) + ): + continue + mask = np.asarray(binary[index], dtype=bool) + if mask.shape != (height, width): + continue + if float(mask.mean()) <= self.min_mask_area_ratio: + continue + frame_index = int(output.frame_idx) + if prompted_target_id is not None: + # A mask inferred directly from this frame's person box is + # more reliable than any propagated mask at the same frame. + tracks[target_id][frame_index] = mask + else: + tracks[target_id].setdefault(frame_index, mask) + + first_prompt = len(frames_rgb) + last_prompt = 0 + with torch.inference_mode(): + for target_id in known_ids: + for frame_index, box in prompts[target_id]: + first_prompt = min(first_prompt, frame_index) + last_prompt = max(last_prompt, frame_index) + self.processor.add_inputs_to_inference_session( + session, + frame_idx=frame_index, + obj_ids=target_id, + input_boxes=[[[float(value) for value in box.as_list()]]], + clear_old_inputs=True, + ) + # Hugging Face SAM2 marks new inputs at object scope rather + # than frame scope. Consume each box immediately so every + # prompted frame becomes a conditioning frame before the + # next prompt overwrites that object-level marker. + output = self.model(session, frame_idx=frame_index) + store_output(output, prompted_target_id=target_id) + + if not propagate: + return tracks + + directions = ((first_prompt, False), (last_prompt, True)) + for start_frame, reverse in directions: + with torch.inference_mode(): + iterator = self.model.propagate_in_video_iterator( + session, + start_frame_idx=start_frame, + reverse=reverse, + ) + for output in iterator: + store_output(output) + return tracks diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/phase_protocol.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/phase_protocol.py new file mode 100644 index 000000000..e49650ad9 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/phase_protocol.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Human input contract and shared revisit allocation.""" + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.revisit_protocol import ( + validate_expected_revisits, +) + +from .schemas import BoundingBox, HumanEntityProtocol + + +def load_human_entity_protocol( + item: dict, + *, + total_frames: int, + frame_width: int, + frame_height: int, +) -> HumanEntityProtocol: + annotations = item.get("annotations", {}) + try: + raw_boxes = annotations["target_face_bboxes_first_frame"] + raw_descriptions = annotations["target_descriptions"] + except KeyError as error: + raise DataInvalidError(f"missing annotations.{error.args[0]}") from error + + expected_revisits = validate_expected_revisits(annotations.get("expected_revisits")) + minimum_frames = (expected_revisits or 1) + 1 + if total_frames < minimum_frames: + raise DataInvalidError( + f"generated video is too short for the requested revisit split; requires at least {minimum_frames} frames" + ) + if not isinstance(raw_boxes, (list, tuple)) or not raw_boxes: + raise DataInvalidError("annotations.target_face_bboxes_first_frame must contain at least one bbox") + try: + boxes = tuple( + BoundingBox.from_value( + value, + field_name="annotations.target_face_bboxes_first_frame", + ) + for value in raw_boxes + ) + except ValueError as error: + raise DataInvalidError(str(error)) from error + for box in boxes: + clipped = box.clipped(width=frame_width, height=frame_height) + if clipped != box: + raise DataInvalidError("target_face_bboxes_first_frame must lie inside the first frame") + + if not isinstance(raw_descriptions, (list, tuple)) or len(raw_descriptions) != len(boxes): + raise DataInvalidError("annotations.target_descriptions must contain one description per first-frame target") + descriptions = tuple(str(value).strip() for value in raw_descriptions) + if any(not value for value in descriptions): + raise DataInvalidError("annotations.target_descriptions cannot contain blanks") + return HumanEntityProtocol( + total_frames=total_frames, + target_face_boxes=boxes, + target_descriptions=descriptions, + expected_revisits=expected_revisits, + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/preparation.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/preparation.py new file mode 100644 index 000000000..31e170697 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/preparation.py @@ -0,0 +1,423 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Match annotated source targets to generated video frames.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Sequence + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError + +from ..revisit_protocol import validate_expected_revisits + + +class ManifestBuildError(RuntimeError): + """A preparation failure that must block formal evaluation.""" + + +def _box(value: Sequence[float], *, field: str) -> list[float]: + if not isinstance(value, (list, tuple)) or len(value) != 4: + raise ManifestBuildError(f"{field} must be [x1,y1,x2,y2]") + try: + box = [float(coordinate) for coordinate in value] + except (TypeError, ValueError) as error: + raise ManifestBuildError(f"{field} contains a non-numeric value") from error + if not np.isfinite(box).all(): + raise ManifestBuildError(f"{field} contains a non-finite value") + if box[2] <= box[0] or box[3] <= box[1]: + raise ManifestBuildError(f"{field} is empty or inverted") + return box + + +def iou(first: Sequence[float], second: Sequence[float]) -> float: + ax1, ay1, ax2, ay2 = first + bx1, by1, bx2, by2 = second + ix1, iy1 = max(ax1, bx1), max(ay1, by1) + ix2, iy2 = min(ax2, bx2), min(ay2, by2) + intersection = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1) + union = (ax2 - ax1) * (ay2 - ay1) + (bx2 - bx1) * (by2 - by1) - intersection + return intersection / union if union > 0.0 else 0.0 + + +def normalized_embedding(value: Any) -> np.ndarray: + embedding = np.asarray(value, dtype=np.float32).reshape(-1) + if embedding.size == 0 or not np.isfinite(embedding).all(): + raise ManifestBuildError("InsightFace returned an invalid embedding") + norm = float(np.linalg.norm(embedding)) + if not np.isfinite(norm) or norm <= 1e-12: + raise ManifestBuildError("InsightFace returned a zero-norm embedding") + return embedding / norm + + +def clip_box(box: Sequence[float], *, width: int, height: int) -> list[float]: + clipped = [ + max(0.0, min(float(width), float(box[0]))), + max(0.0, min(float(height), float(box[1]))), + max(0.0, min(float(width), float(box[2]))), + max(0.0, min(float(height), float(box[3]))), + ] + return _box(clipped, field="generated InsightFace bbox") + + +def load_case( + samples_dir: str, + case_id: str, + *, + description_overrides: dict[str, list[str]], +) -> dict[str, Any]: + case_dir = Path(samples_dir) / case_id + prompt_path = case_dir / "prompt.json" + if not prompt_path.is_file(): + raise ManifestBuildError(f"canonical prompt not found: {prompt_path}") + with open(prompt_path, encoding="utf-8") as handle: + data = json.load(handle) + annotations = data.get("annotations") or {} + raw_boxes = annotations.get("target_face_bboxes_first_frame") or [] + if not raw_boxes: + raise ManifestBuildError(f"{case_id}: canonical target bbox list is empty") + boxes = [ + _box(value, field=f"{case_id}.target_face_bboxes_first_frame[{index}]") for index, value in enumerate(raw_boxes) + ] + descriptions = description_overrides.get( + case_id, + annotations.get("target_descriptions") or [], + ) + descriptions = [str(value).strip() for value in descriptions] + if len(descriptions) != len(boxes) or any(not value for value in descriptions): + raise ManifestBuildError(f"{case_id}: target_descriptions must contain {len(boxes)} non-blank values") + size = data.get("frame_size") or {} + try: + source_width, source_height = int(size["w"]), int(size["h"]) + except (KeyError, TypeError, ValueError) as error: + raise ManifestBuildError(f"{case_id}: invalid frame_size") from error + for target_id, box in enumerate(boxes): + if box[0] < 0.0 or box[1] < 0.0 or box[2] > source_width or box[3] > source_height: + raise ManifestBuildError(f"{case_id}: canonical target_{target_id:03d} bbox lies outside frame_size") + source_image = case_dir / str(data.get("first_frame") or "first_frame.jpg") + if not source_image.is_file(): + raise ManifestBuildError(f"{case_id}: source first frame not found: {source_image}") + return { + "case_id": case_id, + "category": data.get("category"), + "boxes": boxes, + "descriptions": descriptions, + "source_width": source_width, + "source_height": source_height, + "source_image": source_image, + } + + +def detect_faces(app, frame_bgr: np.ndarray, *, min_det_score: float) -> list[dict]: + detections: list[dict] = [] + for face in app.get(frame_bgr): + raw_box = getattr(face, "bbox", None) + raw_embedding = getattr(face, "normed_embedding", None) + confidence = float(getattr(face, "det_score", 0.0)) + if raw_box is None or raw_embedding is None or confidence < min_det_score: + continue + detections.append( + { + "bbox": _box(np.asarray(raw_box).tolist(), field="InsightFace bbox"), + "embedding": normalized_embedding(raw_embedding), + "det_score": confidence, + } + ) + return detections + + +def _hungarian_pairs(scores: np.ndarray) -> list[tuple[int, int]]: + if scores.ndim != 2: + raise ValueError("assignment scores must be two-dimensional") + if scores.shape[0] == 0 or scores.shape[1] == 0: + return [] + from scipy.optimize import linear_sum_assignment + + rows, columns = linear_sum_assignment(scores, maximize=True) + return [(int(row), int(column)) for row, column in zip(rows, columns)] + + +def source_anchors( + canonical_boxes: Sequence[Sequence[float]], + detections: Sequence[dict], + *, + iou_threshold: float, +) -> tuple[list[np.ndarray], list[dict]]: + if len(detections) < len(canonical_boxes): + raise ManifestBuildError( + f"source first frame has {len(detections)} detected faces but {len(canonical_boxes)} canonical targets" + ) + scores = np.asarray( + [[iou(box, detection["bbox"]) for detection in detections] for box in canonical_boxes], + dtype=np.float64, + ) + pairs = _hungarian_pairs(scores) + by_target = {target_id: detection_id for target_id, detection_id in pairs} + if len(by_target) != len(canonical_boxes): + raise ManifestBuildError("source first-frame target assignment is incomplete") + anchors: list[np.ndarray] = [] + audit: list[dict] = [] + for target_id, canonical_box in enumerate(canonical_boxes): + detection_id = by_target[target_id] + overlap = float(scores[target_id, detection_id]) + if overlap < iou_threshold: + raise ManifestBuildError( + f"canonical target_{target_id:03d} has source detection IoU {overlap:.4f} below {iou_threshold:.4f}" + ) + detection = detections[detection_id] + anchors.append(detection["embedding"]) + audit.append( + { + "target_id": target_id, + "canonical_bbox": list(canonical_box), + "source_detection_bbox": detection["bbox"], + "source_detection_score": detection["det_score"], + "source_iou": overlap, + } + ) + return anchors, audit + + +def generated_target_boxes( + anchors: Sequence[np.ndarray], + detections: Sequence[dict], + *, + similarity_threshold: float, + width: int, + height: int, +) -> tuple[list[list[float]], list[dict], list[list[float]]]: + if len(detections) < len(anchors): + raise ManifestBuildError( + f"generated first frame has {len(detections)} detected faces but {len(anchors)} canonical targets" + ) + scores = np.asarray( + [[float(np.dot(anchor, detection["embedding"])) for detection in detections] for anchor in anchors], + dtype=np.float64, + ) + pairs = _hungarian_pairs(scores) + by_target = {target_id: detection_id for target_id, detection_id in pairs} + if len(by_target) != len(anchors): + raise ManifestBuildError("generated first-frame identity assignment is incomplete") + boxes: list[list[float]] = [] + audit: list[dict] = [] + used: set[int] = set() + for target_id in range(len(anchors)): + detection_id = by_target[target_id] + similarity = float(scores[target_id, detection_id]) + if similarity < similarity_threshold: + raise ManifestBuildError( + f"target_{target_id:03d} ArcFace similarity {similarity:.4f} below {similarity_threshold:.4f}" + ) + detection = detections[detection_id] + box = clip_box(detection["bbox"], width=width, height=height) + boxes.append(box) + audit.append( + { + "target_id": target_id, + "generated_bbox": box, + "generated_detection_score": detection["det_score"], + "arcface_similarity": similarity, + } + ) + used.add(detection_id) + extras = [ + clip_box(detection["bbox"], width=width, height=height) + for detection_id, detection in enumerate(detections) + if detection_id not in used + ] + return boxes, audit, extras + + +def scaled_canonical_boxes( + case: dict[str, Any], + *, + width: int, + height: int, +) -> list[list[float]]: + scale_x = float(width) / float(case["source_width"]) + scale_y = float(height) / float(case["source_height"]) + return [ + clip_box( + [ + box[0] * scale_x, + box[1] * scale_y, + box[2] * scale_x, + box[3] * scale_y, + ], + width=width, + height=height, + ) + for box in case["boxes"] + ] + + +def probe_first_frame(video_path: str) -> tuple[int, int, np.ndarray]: + import cv2 + + capture = cv2.VideoCapture(video_path) + if not capture.isOpened(): + capture.release() + raise ManifestBuildError(f"video is unreadable: {video_path}") + width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) + ok, frame = capture.read() + capture.release() + if not ok or frame is None or width <= 0 or height <= 0: + raise ManifestBuildError(f"failed to decode frame 0: {video_path}") + return width, height, frame + + +def load_source_frame(case: dict[str, Any]) -> np.ndarray: + import cv2 + + frame = cv2.imread(str(case["source_image"])) + if frame is None: + raise ManifestBuildError(f"failed to read source image: {case['source_image']}") + height, width = frame.shape[:2] + if (width, height) != (case["source_width"], case["source_height"]): + raise ManifestBuildError( + f"{case['case_id']}: prompt frame_size " + f"{case['source_width']}x{case['source_height']} does not match source image " + f"{width}x{height}" + ) + return frame + + +def prepare_samples( + samples: list[dict], + dataset: Path, + *, + expected_revisits: int | None, + app, + audit_rows: list[dict], +) -> list[dict]: + """Match each generated first frame to frozen source targets, in source order.""" + validate_expected_revisits(expected_revisits) + # Compare identical case sets across models within each trajectory. + matrix: dict[str, dict[str, set[str]]] = {} + for sample in samples: + matrix.setdefault(sample["condition_id"], {}).setdefault(sample["model_id"], set()).add(sample["sample_id"]) + for trajectory, model_cases in matrix.items(): + reference = next(iter(model_cases.values())) + if any(cases != reference for cases in model_cases.values()): + raise ManifestBuildError(f"model case sets differ for trajectory {trajectory}") + cache = {} + prepared = [] + for sample in samples: + case_id = sample["sample_id"] + audit = {"item_id": sample["item_id"], "status": "error"} + try: + if case_id not in cache: + case = load_case(str(dataset), case_id, description_overrides={}) + detections = detect_faces(app, load_source_frame(case), min_det_score=0.0) + anchors, source_audit = source_anchors(case["boxes"], detections, iou_threshold=0.1) + cache[case_id] = case, anchors, source_audit + case, anchors, source_audit = cache[case_id] + width, height, frame = probe_first_frame(sample["video"]) + detections = detect_faces(app, frame, min_det_score=0.0) + match_status, match_error = "ok", None + try: + boxes, generated_audit, extras = generated_target_boxes( + anchors, + detections, + similarity_threshold=0.35, + width=width, + height=height, + ) + target_audit = [{**source, **generated} for source, generated in zip(source_audit, generated_audit)] + except ManifestBuildError as error: + # The reference gives the whole model/case/trajectory zero, + # retaining the canonical target count and explicit reason. + match_status, match_error = "failed", str(error) + boxes = scaled_canonical_boxes(case, width=width, height=height) + extras = [clip_box(detection["bbox"], width=width, height=height) for detection in detections] + target_audit = [ + { + **source, + "generated_bbox": boxes[index], + "assignment_method": "first_frame_match_failed", + "first_frame_match_error": match_error, + } + for index, source in enumerate(source_audit) + ] + prepared.append( + { + **sample, + "annotations": { + "target_face_bboxes_first_frame": boxes, + "target_descriptions": list(case["descriptions"]), + "expected_revisits": expected_revisits, + "first_frame_match_status": match_status, + "first_frame_match_error": match_error, + }, + } + ) + audit.update( + status="ok", + first_frame_match_status=match_status, + first_frame_match_error=match_error, + target_matches=target_audit, + extra_first_frame_face_bboxes=extras, + ) + print(f"{sample['item_id']}: first-frame match {match_status}", flush=True) + except Exception as error: + audit["error"] = f"{type(error).__name__}: {error}" + raise + finally: + audit_rows.append(audit) + return prepared + + +def prepare_for_evaluation( + samples: list[dict], context, *, expected_revisits: int | None, weight_metric: str = "human_identity" +) -> None: + """Annotate standard-layout Human inputs before publishing the run manifest.""" + from .evaluator import _face_backend + + for sample in samples: + if sample.get("category") not in (None, "human"): + continue + if "first_frame_match_status" in (sample.get("annotations") or {}): + if expected_revisits is not None: + sample["annotations"]["expected_revisits"] = validate_expected_revisits(expected_revisits) + continue + if sample.get("category") != "human": + continue + sample.pop("human_preparation_error", None) + try: + selected_revisits = validate_expected_revisits( + expected_revisits + if expected_revisits is not None + else (sample.get("annotations") or {}).get("expected_revisits") + ) + video = Path(sample["video"]) + if video.name != "video.mp4" or video.parent.name != sample.get("sample_id"): + raise DataInvalidError("automatic Human preparation requires video.mp4 in the case folder") + # Source image/annotations and generated video share the case folder. + source_root = video.parents[1] + try: + case = load_case(str(source_root), sample["sample_id"], description_overrides={}) + load_source_frame(case) + except Exception as error: + raise DataInvalidError(str(error)) from error + prepared = prepare_samples( + [sample], + source_root, + expected_revisits=selected_revisits, + app=_face_backend(context, weight_metric).app, + audit_rows=[], + ) + sample.update(prepared[0]) + except Exception as error: + # Persist preparation failures as input provenance and let each + # Human metric report them through the usual per-case error path. + sample["human_preparation_error"] = { + "status": "data_invalid" + if isinstance(error, (DataInvalidError, ManifestBuildError)) + else "eval_failed", + "message": f"{type(error).__name__}: {error}", + } diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/schemas.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/schemas.py new file mode 100644 index 000000000..e6b91bf8f --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/schemas.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Typed records shared by the Human Entity Memory metrics.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.revisit_protocol import TargetPhase + + +@dataclass(frozen=True) +class BoundingBox: + x1: float + y1: float + x2: float + y2: float + + @classmethod + def from_value(cls, value: Any, *, field_name: str) -> BoundingBox: + import numpy as np + + if isinstance(value, dict): + try: + coordinates = [value[key] for key in ("x1", "y1", "x2", "y2")] + except KeyError as error: + raise ValueError(f"{field_name} bbox mappings require x1/y1/x2/y2") from error + elif isinstance(value, (list, tuple)) and len(value) == 4: + coordinates = list(value) + else: + raise ValueError(f"{field_name} must contain [x1, y1, x2, y2] boxes") + try: + box = cls(*(float(coordinate) for coordinate in coordinates)) + except (TypeError, ValueError) as error: + raise ValueError(f"{field_name} contains a non-numeric bbox") from error + if not np.isfinite(box.as_array()).all(): + raise ValueError(f"{field_name} contains a non-finite bbox") + if box.x2 <= box.x1 or box.y2 <= box.y1: + raise ValueError(f"{field_name} contains an empty or inverted bbox") + return box + + def as_array(self) -> np.ndarray: + import numpy as np + + return np.asarray([self.x1, self.y1, self.x2, self.y2], dtype=np.float32) + + def as_list(self) -> list[float]: + return [self.x1, self.y1, self.x2, self.y2] + + @property + def center(self) -> tuple[float, float]: + return ((self.x1 + self.x2) / 2.0, (self.y1 + self.y2) / 2.0) + + @property + def area(self) -> float: + return (self.x2 - self.x1) * (self.y2 - self.y1) + + def contains(self, point: tuple[float, float]) -> bool: + x, y = point + return self.x1 <= x <= self.x2 and self.y1 <= y <= self.y2 + + def clipped(self, *, width: int, height: int) -> BoundingBox: + return BoundingBox( + x1=max(0.0, min(float(width), self.x1)), + y1=max(0.0, min(float(height), self.y1)), + x2=max(0.0, min(float(width), self.x2)), + y2=max(0.0, min(float(height), self.y2)), + ) + + +@dataclass(frozen=True) +class HumanEntityProtocol: + total_frames: int + target_face_boxes: tuple[BoundingBox, ...] + target_descriptions: tuple[str, ...] + expected_revisits: int | None + + @property + def n_targets(self) -> int: + return len(self.target_face_boxes) + + @property + def required_frame_indices(self) -> tuple[int, ...]: + return tuple(range(self.total_frames)) + + +@dataclass +class BodyTrackAnalysis: + target_phases: tuple[TargetPhase, ...] + assignments: dict[int, dict[int, BoundingBox]] + + +@dataclass(frozen=True) +class FaceObservation: + bbox: BoundingBox + embedding: np.ndarray + + +@dataclass +class FaceTrackAnalysis: + anchors: tuple[np.ndarray, ...] + seeded_association: dict[int, FaceObservation] + candidates_by_frame: dict[int, list[FaceObservation]] + track_associations: dict[int, dict[int, FaceObservation]] + + +@dataclass(frozen=True) +class IdentityTargetResult: + target_id: int + phase: TargetPhase + frame_scores: dict[int, float] + similarity: float + coverage: float + n_reference_embeddings: int + n_return_embeddings: int + revisit_scores: tuple[float, ...] + revisit_coverages: tuple[float, ...] + n_revisit_embeddings: tuple[int, ...] + + +@dataclass +class IdentityAnalysis: + protocol: HumanEntityProtocol + target_phases: tuple[TargetPhase, ...] + appearance_target_phases: tuple[TargetPhase, ...] + body_assignments: dict[int, dict[int, BoundingBox]] + target_results: tuple[IdentityTargetResult, ...] + max_revisit_gap_frames: int + revisit_boundary_buffer_frames: int + decoded_frames: dict[int, np.ndarray] = field(repr=False) + video_path: str = "" + + +@dataclass(frozen=True) +class AppearanceTargetResult: + target_id: int + phase: TargetPhase + frame_scores: dict[int, float] + score: float + reappearance_coverage: float + n_reference_features: int + n_return_features: int + revisit_scores: tuple[float, ...] + revisit_coverages: tuple[float, ...] + n_revisit_features: tuple[int, ...] + + +@dataclass(frozen=True) +class AppearanceAnalysis: + target_results: tuple[AppearanceTargetResult, ...] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/target_association.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/target_association.py new file mode 100644 index 000000000..3475ca8f8 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/human/target_association.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""One-to-one target association and cosine helpers.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + +from .schemas import BoundingBox + + +def normalized_embedding(value: np.ndarray | Sequence[float]) -> np.ndarray: + embedding = np.asarray(value, dtype=np.float32).reshape(-1) + if embedding.size == 0 or not np.isfinite(embedding).all(): + raise EvaluationError("embedding must be finite and non-empty") + norm = float(np.linalg.norm(embedding)) + if not np.isfinite(norm) or norm <= 1e-12: + raise EvaluationError("embedding has zero or non-finite norm") + return embedding / norm + + +def cosine_matrix( + references: Sequence[np.ndarray], + candidates: Sequence[np.ndarray], +) -> np.ndarray: + if not references or not candidates: + return np.empty((len(references), len(candidates)), dtype=np.float32) + reference_matrix = np.stack([normalized_embedding(value) for value in references], axis=0) + candidate_matrix = np.stack([normalized_embedding(value) for value in candidates], axis=0) + return reference_matrix @ candidate_matrix.T + + +def hungarian_matches( + scores: np.ndarray, + *, + threshold: float, +) -> dict[int, int]: + matrix = np.asarray(scores, dtype=np.float64) + if matrix.ndim != 2: + raise ValueError("scores must be a 2D matrix") + if matrix.size == 0: + return {} + if not np.isfinite(matrix).all(): + raise EvaluationError("association matrix contains a non-finite value") + try: + from scipy.optimize import linear_sum_assignment + except ImportError as error: + raise EvaluationError( + "Human Entity association requires scipy; install the WorldMemBench metric requirements" + ) from error + rows, columns = linear_sum_assignment(matrix, maximize=True) + return {int(row): int(column) for row, column in zip(rows, columns) if matrix[row, column] >= threshold} + + +def assign_faces_to_person_boxes( + face_boxes: dict[int, BoundingBox], + person_boxes: Sequence[tuple[BoundingBox, float]], +) -> dict[int, BoundingBox]: + """Assign distinct containing person boxes while maximizing confidence.""" + + if not face_boxes or not person_boxes: + return {} + target_ids = sorted(face_boxes) + scores = np.full((len(target_ids), len(person_boxes)), -1e9, dtype=np.float64) + valid = np.zeros_like(scores, dtype=bool) + for row, target_id in enumerate(target_ids): + for column, (person_box, confidence) in enumerate(person_boxes): + if person_box.contains(face_boxes[target_id].center): + scores[row, column] = float(confidence) + valid[row, column] = True + try: + from scipy.optimize import linear_sum_assignment + except ImportError as error: + raise EvaluationError( + "Human Appearance person assignment requires scipy; install the WorldMemBench metric requirements" + ) from error + rows, columns = linear_sum_assignment(scores, maximize=True) + return { + target_ids[int(row)]: person_boxes[int(column)][0] for row, column in zip(rows, columns) if valid[row, column] + } diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/__init__.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/__init__.py new file mode 100644 index 000000000..cb0cfe906 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Object Memory metrics.""" + +from .metric import ObjectAppearanceMetric, ObjectGeometryMetric + +__all__ = ["ObjectAppearanceMetric", "ObjectGeometryMetric"] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/appearance_features.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/appearance_features.py new file mode 100644 index 000000000..5f2d5056e --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/appearance_features.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Masked DINOv2 features and Human-aligned Object Appearance scoring.""" + +from __future__ import annotations + +from typing import Protocol + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.human.target_association import ( + cosine_matrix, + normalized_embedding, +) + +from .schemas import ObjectAppearanceResult, ObjectTrackAnalysis + + +class ObjectAppearanceFeatureExtractor(Protocol): + def embed(self, frame_rgb: np.ndarray, mask: np.ndarray) -> np.ndarray | None: ... + + +def analyze_object_appearance( + track: ObjectTrackAnalysis, + feature_extractor: ObjectAppearanceFeatureExtractor, +) -> ObjectAppearanceResult: + phase = track.phase + frame_scores: dict[int, float] = {frame_index: 0.0 for frame_index in phase.return_indices} + if not phase.is_scorable: + empty_slots = tuple(0.0 for _ in phase.revisit_indices) + empty_counts = tuple(0 for _ in phase.revisit_indices) + return ObjectAppearanceResult( + phase=phase, + frame_scores=frame_scores, + similarity=0.0, + appearance_coverage=0.0, + n_reference_features=0, + n_return_features=0, + revisit_scores=empty_slots, + revisit_coverages=empty_slots, + n_revisit_features=empty_counts, + ) + references: list[np.ndarray] = [] + for frame_index in phase.visible_indices: + observation = track.observations[frame_index] + if not observation.present or observation.mask is None or not observation.mask_valid: + continue + embedding = feature_extractor.embed( + track.decoded_frames[frame_index], + observation.mask, + ) + if embedding is not None: + references.append(normalized_embedding(embedding)) + + returned_by_frame: dict[int, np.ndarray] = {} + for frame_index in phase.return_indices: + observation = track.observations[frame_index] + if not observation.present or observation.mask is None or not observation.mask_valid: + continue + candidate = feature_extractor.embed( + track.decoded_frames[frame_index], + observation.mask, + ) + if candidate is None: + continue + normalized = normalized_embedding(candidate) + returned_by_frame[frame_index] = normalized + if not references: + continue + similarities = cosine_matrix( + references, + [normalized], + )[:, 0] + frame_scores[frame_index] = float(np.clip(similarities.max(), 0.0, 1.0)) + + return_frame_count = len(phase.return_indices) + n_return_features = len(returned_by_frame) + similarity = ( + sum(frame_scores[index] for index in returned_by_frame) / n_return_features + if n_return_features and references and phase.is_scorable + else 0.0 + ) + coverage = n_return_features / return_frame_count if return_frame_count else 0.0 + revisit_scores: list[float] = [] + revisit_coverages: list[float] = [] + n_revisit_features: list[int] = [] + for revisit_indices in phase.revisit_indices: + feature_indices = [index for index in revisit_indices if index in returned_by_frame] + n_features = len(feature_indices) + n_revisit_features.append(n_features) + revisit_coverages.append(n_features / len(revisit_indices) if revisit_indices else 0.0) + revisit_scores.append( + ( + sum(frame_scores[index] for index in feature_indices) / n_features + if n_features and references and phase.is_scorable + else 0.0 + ) + ) + return ObjectAppearanceResult( + phase=phase, + frame_scores=frame_scores, + similarity=similarity, + appearance_coverage=coverage, + n_reference_features=len(references), + n_return_features=n_return_features, + revisit_scores=tuple(revisit_scores), + revisit_coverages=tuple(revisit_coverages), + n_revisit_features=tuple(n_revisit_features), + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/evaluator.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/evaluator.py new file mode 100644 index 000000000..403d69ca4 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/evaluator.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Shared Object tracking and appearance analysis, scoped to one evaluation context.""" + +from pathlib import Path + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.video import VideoReader + +from ..features import MaskedDinoV2Features, get_dino_model +from ..revisit_protocol import MAX_SHORT_MISSED_FRAMES, infer_target_phase +from .appearance_features import analyze_object_appearance +from .mask_alignment import align_mask_to_generated_frame +from .object_redetector import GroundingDinoClipObjectRedetector, object_query_crop +from .object_tracker import TransformersSam2ObjectTracker +from .preparation import load_source, source_path +from .protocol import load_object_entity_protocol +from .schemas import ObjectAppearanceAnalysis, ObjectTrackAnalysis + + +def _tracker(context, weight_metric="object_appearance"): + config = context.metric_config("object_appearance") + checkpoint = str( + config.get("sam2_checkpoint") or context.cache_dir / "metrics" / weight_metric / "sam2-hiera-large" + ) + cache_size = config.get("sam2_vision_cache_size", 1) + return context.get_resource( + "object:tracker", + lambda: TransformersSam2ObjectTracker( + checkpoint=checkpoint, device=context.device, vision_cache_size=cache_size + ), + ) + + +def _redetector(context, weight_metric="object_appearance"): + config = context.metric_config("object_appearance") + checkpoint = str( + config.get("grounding_dino_checkpoint") or context.cache_dir / "metrics" / weight_metric / "grounding-dino-base" + ) + reid = str(config.get("clip_reid_checkpoint") or context.cache_dir / "metrics" / weight_metric / "ViT-B-32.pt") + # Forward explicit overrides; defaults live in the backend. + fields = ( + "search_stride", + "batch_size", + "box_threshold", + "text_threshold", + "nms_threshold", + "candidate_limit", + "reid_batch_size", + "reid_similarity_threshold", + "reid_high_confidence_threshold", + "reid_min_similarity_margin", + "min_confirmations", + "max_sample_gap", + "min_box_area_ratio", + "max_box_area_ratio", + ) + options = {key: config[f"redetection_{key}"] for key in fields if f"redetection_{key}" in config} + if "redetection_reid_crop_context" in config: + options["reid_crop_context_fraction"] = config["redetection_reid_crop_context"] + return context.get_resource( + "object:redetector", + lambda: GroundingDinoClipObjectRedetector( + grounding_dino_checkpoint=checkpoint, reid_checkpoint=reid, device=context.device, **options + ), + ) + + +def _feature_extractor(context): + config = context.metric_config("object_appearance") + checkpoint = str(config.get("dino_checkpoint") or context.cache_dir / "metrics/object_appearance/dinov2-large") + return context.get_resource( + "object:dino", + lambda: MaskedDinoV2Features(model=get_dino_model(context, checkpoint), device=context.device, batch_size=1), + ) + + +def _visibility_segments(observations): + segments = [] + for index in sorted(i for i, observation in observations.items() if observation.present): + if segments and segments[-1][1] + 1 == index: + segments[-1] = (segments[-1][0], index) + else: + segments.append((index, index)) + return tuple(segments) + + +def object_track_analysis_for_item( + item, context, *, tracker=None, redetector=None, reader=None, weight_metric="object_appearance" +): + # The orchestrator clears this case cache after both metrics finish. Context + # configuration/model instances stay fixed throughout the evaluation run. + def create(): + video_path = str(Path(item["video"]).resolve()) + image_path = source_path(item, "condition_image") + mask_path = source_path(item, "target_object_mask_first_frame") + video_reader = reader or VideoReader() + info = video_reader.probe(video_path) + protocol = load_object_entity_protocol( + item, total_frames=info.total_frames, condition_image_path=image_path, target_mask_path=mask_path + ) + condition_rgb, source_mask = load_source(image_path, mask_path) + decoded = video_reader.read_indices(video_path, list(protocol.required_frame_indices)) + decoded_frames = {record.frame_idx: record.rgb for record in decoded} + frames = [decoded_frames[index] for index in protocol.required_frame_indices] + alignment = align_mask_to_generated_frame(condition_rgb, source_mask, frames[0]) + if not alignment.mask.any(): + raise DataInvalidError("source mask is empty after alignment to the generated first frame") + query = object_query_crop(frames[0], alignment.mask) + prompts = (redetector or _redetector(context, weight_metric)).detect(protocol.target_description, query, frames) + observations = (tracker or _tracker(context, weight_metric)).track(frames, alignment.mask, prompts) + gap = context.metric_config("object_appearance").get("max_revisit_gap_frames", MAX_SHORT_MISSED_FRAMES) + phase = infer_target_phase( + target_id=0, + total_frames=protocol.total_frames, + matched_frame_indices=[i for i, observation in observations.items() if observation.present], + max_revisit_gap_frames=gap, + expected_revisits=protocol.expected_revisits, + ) + return ObjectTrackAnalysis( + protocol=protocol, + phase=phase, + observations=observations, + decoded_frames=decoded_frames, + video_path=video_path, + alignment_method=alignment.method, + frame0_alignment_score=alignment.score, + redetection_prompts=prompts, + visibility_segments=_visibility_segments(observations), + max_revisit_gap_frames=gap, + ) + + return context.get_cached(("object:track", item["item_id"]), create) + + +def object_appearance_analysis_for_item( + item, context, *, tracker=None, redetector=None, feature_extractor=None, reader=None +): + track = object_track_analysis_for_item(item, context, tracker=tracker, redetector=redetector, reader=reader) + result = analyze_object_appearance(track, feature_extractor or _feature_extractor(context)) + return ObjectAppearanceAnalysis(track=track, result=result) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/geometry_consistency.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/geometry_consistency.py new file mode 100644 index 000000000..986d17a4b --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/geometry_consistency.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Canonical silhouette comparison for returned object shapes.""" + +from __future__ import annotations + +import math + +import numpy as np + +from .schemas import ObjectGeometryResult, ObjectTrackAnalysis + +DEFAULT_CANONICAL_SIZE = 128 +DEFAULT_CANONICAL_PADDING = 8 + + +def canonicalize_object_mask( + mask: np.ndarray, + *, + canvas_size: int = DEFAULT_CANONICAL_SIZE, + padding: int = DEFAULT_CANONICAL_PADDING, +) -> np.ndarray | None: + """Center a binary mask on a square canvas using uniform-only scaling. + + Translation and global image scale are removed, while aspect ratio and + local silhouette deformation are preserved. + """ + + if mask.ndim != 2: + raise ValueError("object shape mask must be 2-D") + if canvas_size <= 0: + raise ValueError("canonical canvas size must be positive") + if padding < 0 or 2 * padding >= canvas_size: + raise ValueError("canonical padding must leave a non-empty canvas interior") + + binary = np.asarray(mask, dtype=bool) + if not binary.any(): + return None + y_coordinates, x_coordinates = np.where(binary) + crop = binary[ + int(y_coordinates.min()) : int(y_coordinates.max()) + 1, + int(x_coordinates.min()) : int(x_coordinates.max()) + 1, + ] + crop_height, crop_width = crop.shape + usable_size = canvas_size - 2 * padding + scale = usable_size / float(max(crop_height, crop_width)) + target_height = max(1, min(usable_size, int(round(crop_height * scale)))) + target_width = max(1, min(usable_size, int(round(crop_width * scale)))) + + source_y = np.minimum( + (np.arange(target_height) * crop_height // target_height), + crop_height - 1, + ) + source_x = np.minimum( + (np.arange(target_width) * crop_width // target_width), + crop_width - 1, + ) + resized = crop[np.ix_(source_y, source_x)] + canvas = np.zeros((canvas_size, canvas_size), dtype=bool) + top = (canvas_size - target_height) // 2 + left = (canvas_size - target_width) // 2 + canvas[top : top + target_height, left : left + target_width] = resized + return canvas + + +def mask_iou(first: np.ndarray, second: np.ndarray) -> float: + if first.shape != second.shape: + raise ValueError("canonical masks must have the same shape") + union = np.logical_or(first, second).sum() + if union == 0: + return 0.0 + intersection = np.logical_and(first, second).sum() + return float(intersection / union) + + +def analyze_object_geometry( + track: ObjectTrackAnalysis, + *, + canonical_size: int = DEFAULT_CANONICAL_SIZE, + canonical_padding: int = DEFAULT_CANONICAL_PADDING, +) -> ObjectGeometryResult: + phase = track.phase + frame_scores = {index: 0.0 for index in phase.return_indices} + reference_indices: list[int] = [] + reference_masks: list[np.ndarray] = [] + reference_areas: list[float] = [] + for frame_index in phase.visible_indices: + observation = track.observations[frame_index] + if not observation.present or observation.mask is None or not observation.mask_valid: + continue + canonical = canonicalize_object_mask( + observation.mask, + canvas_size=canonical_size, + padding=canonical_padding, + ) + if canonical is None: + continue + reference_indices.append(frame_index) + reference_masks.append(canonical) + reference_areas.append(float(np.asarray(observation.mask, dtype=bool).sum())) + + returned_masks: dict[int, np.ndarray] = {} + area_drifts_by_frame: dict[int, float] = {} + for frame_index in phase.return_indices: + observation = track.observations[frame_index] + if not observation.present or observation.mask is None or not observation.mask_valid: + continue + canonical_candidate = canonicalize_object_mask( + observation.mask, + canvas_size=canonical_size, + padding=canonical_padding, + ) + if canonical_candidate is None: + continue + returned_masks[frame_index] = canonical_candidate + if not reference_masks: + continue + similarities = [mask_iou(reference, canonical_candidate) for reference in reference_masks] + best_reference = int(np.argmax(similarities)) + score = similarities[best_reference] + candidate_area = float(np.asarray(observation.mask, dtype=bool).sum()) + area_ratio = max( + candidate_area / reference_areas[best_reference], + np.finfo(float).tiny, + ) + frame_scores[frame_index] = score + area_drifts_by_frame[frame_index] = abs(math.log(area_ratio)) + + return_frame_count = len(phase.return_indices) + n_return_masks = len(returned_masks) + similarity = ( + sum(frame_scores[index] for index in returned_masks) / n_return_masks + if n_return_masks and reference_masks and phase.is_scorable + else 0.0 + ) + coverage = n_return_masks / return_frame_count if return_frame_count else 0.0 + all_area_drifts = list(area_drifts_by_frame.values()) + area_ratio_drift = float(np.median(all_area_drifts)) if all_area_drifts else 0.0 + revisit_scores: list[float] = [] + revisit_coverages: list[float] = [] + n_revisit_masks: list[int] = [] + revisit_area_ratio_drifts: list[float] = [] + for revisit_indices in phase.revisit_indices: + valid_indices = [index for index in revisit_indices if index in returned_masks] + n_masks = len(valid_indices) + n_revisit_masks.append(n_masks) + revisit_coverages.append(n_masks / len(revisit_indices) if revisit_indices else 0.0) + revisit_scores.append( + ( + sum(frame_scores[index] for index in valid_indices) / n_masks + if n_masks and reference_masks and phase.is_scorable + else 0.0 + ) + ) + revisit_drifts = [area_drifts_by_frame[index] for index in valid_indices if index in area_drifts_by_frame] + revisit_area_ratio_drifts.append(float(np.median(revisit_drifts)) if revisit_drifts else 0.0) + return ObjectGeometryResult( + phase=phase, + reference_frame_indices=tuple(reference_indices), + frame_scores=frame_scores, + similarity=similarity, + geometry_coverage=coverage, + area_ratio_drift=area_ratio_drift, + n_reference_masks=len(reference_masks), + n_return_masks=n_return_masks, + revisit_scores=tuple(revisit_scores), + revisit_coverages=tuple(revisit_coverages), + n_revisit_masks=tuple(n_revisit_masks), + revisit_area_ratio_drifts=tuple(revisit_area_ratio_drifts), + ) + + +__all__ = [ + "DEFAULT_CANONICAL_PADDING", + "DEFAULT_CANONICAL_SIZE", + "analyze_object_geometry", + "canonicalize_object_mask", + "mask_iou", +] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/mask_alignment.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/mask_alignment.py new file mode 100644 index 000000000..56366e420 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/mask_alignment.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Align a source-frame object mask to a generated video's first frame.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + + +@dataclass(frozen=True) +class MaskAlignment: + method: str + score: float + mask: np.ndarray + + +def _resize_pair( + image: np.ndarray, + mask: np.ndarray, + *, + width: int, + height: int, +) -> tuple[np.ndarray, np.ndarray]: + import cv2 + + resized_image = cv2.resize(image, (width, height), interpolation=cv2.INTER_LINEAR) + resized_mask = cv2.resize( + mask.astype(np.uint8), + (width, height), + interpolation=cv2.INTER_NEAREST, + ).astype(bool) + return resized_image, resized_mask + + +def _center_crop( + image: np.ndarray, + mask: np.ndarray, + *, + destination_aspect: float, +) -> tuple[np.ndarray, np.ndarray]: + height, width = image.shape[:2] + source_aspect = width / float(height) + if source_aspect > destination_aspect: + crop_width = max(1, min(width, int(round(height * destination_aspect)))) + left = (width - crop_width) // 2 + return ( + image[:, left : left + crop_width], + mask[:, left : left + crop_width], + ) + crop_height = max(1, min(height, int(round(width / destination_aspect)))) + top = (height - crop_height) // 2 + return ( + image[top : top + crop_height], + mask[top : top + crop_height], + ) + + +def _letterbox_pair( + image: np.ndarray, + mask: np.ndarray, + *, + width: int, + height: int, +) -> tuple[np.ndarray, np.ndarray]: + source_height, source_width = image.shape[:2] + scale = min(width / float(source_width), height / float(source_height)) + scaled_width = max(1, int(round(source_width * scale))) + scaled_height = max(1, int(round(source_height * scale))) + scaled_image, scaled_mask = _resize_pair( + image, + mask, + width=scaled_width, + height=scaled_height, + ) + canvas = np.zeros((height, width, 3), dtype=np.uint8) + mask_canvas = np.zeros((height, width), dtype=bool) + left = (width - scaled_width) // 2 + top = (height - scaled_height) // 2 + canvas[top : top + scaled_height, left : left + scaled_width] = scaled_image + mask_canvas[top : top + scaled_height, left : left + scaled_width] = scaled_mask + return canvas, mask_canvas + + +def _alignment_score(candidate: np.ndarray, generated: np.ndarray) -> float: + difference = np.abs(candidate.astype(np.float32) - generated.astype(np.float32)) + return float(np.clip(1.0 - difference.mean() / 255.0, 0.0, 1.0)) + + +def align_mask_to_generated_frame( + condition_rgb: np.ndarray, + source_mask: np.ndarray, + generated_frame_rgb: np.ndarray, +) -> MaskAlignment: + if condition_rgb.ndim != 3 or condition_rgb.shape[2] != 3: + raise ValueError("condition image must be an RGB image") + if source_mask.shape != condition_rgb.shape[:2]: + raise ValueError("target object mask must match the condition image size") + if generated_frame_rgb.ndim != 3 or generated_frame_rgb.shape[2] != 3: + raise ValueError("generated first frame must be an RGB image") + + height, width = generated_frame_rgb.shape[:2] + resized = _resize_pair(condition_rgb, source_mask, width=width, height=height) + cropped_source = _center_crop( + condition_rgb, + source_mask, + destination_aspect=width / float(height), + ) + center_cropped = _resize_pair( + *cropped_source, + width=width, + height=height, + ) + letterboxed = _letterbox_pair( + condition_rgb, + source_mask, + width=width, + height=height, + ) + candidates = ( + ("resize", resized), + ("center_crop_resize", center_cropped), + ("letterbox_resize", letterboxed), + ) + method, (_, aligned_mask), score = max( + ((name, pair, _alignment_score(pair[0], generated_frame_rgb)) for name, pair in candidates), + key=lambda value: value[2], + ) + return MaskAlignment(method=method, score=score, mask=aligned_mask) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/metric.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/metric.py new file mode 100644 index 000000000..0fdc756b0 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/metric.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Single-target Object Appearance and Geometry with shared tracking.""" + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import ( + DataInvalidError, + EvaluationError, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.aggregation import ( + aggregate_video_summary, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.base import Metric +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.results import ( + ItemResult, + ResultStatus, + UnitResult, +) +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.video import VideoReader + +from ..revisit_protocol import MAX_SHORT_MISSED_FRAMES, MIN_REAPPEAR_MATCH_RUN, REVISIT_BOUNDARY_BUFFER_FRAMES +from .evaluator import object_appearance_analysis_for_item, object_track_analysis_for_item +from .geometry_consistency import DEFAULT_CANONICAL_PADDING, DEFAULT_CANONICAL_SIZE, analyze_object_geometry +from .object_tracker import ( + MAX_MASK_AREA_RATIO, + MIN_MASK_AREA_RATIO, + MIN_MASK_AREA_RELATIVE_TO_SEED, + PRESENCE_LOGIT_THRESHOLD, +) +from .revisit import official_revisit_score, revisit_metadata, revisit_score_fields + + +def _frame_boundary_contact_ratio(track, indices) -> float: + contacts = [] + for index in indices: + observation = track.observations[index] + if not observation.present or observation.mask is None: + continue + mask = observation.mask + contacts.append(bool(mask[0].any() or mask[-1].any() or mask[:, 0].any() or mask[:, -1].any())) + return float(sum(contacts) / len(contacts)) if contacts else 0.0 + + +class _ObjectMetric(Metric): + geometry = False + + def __init__(self, *, tracker=None, redetector=None, feature_extractor=None, reader=None): + self.tracker = tracker + self.redetector = redetector + self.feature_extractor = feature_extractor + self.reader = reader or VideoReader() + + def evaluate(self, sample: dict, context) -> list[UnitResult]: + if sample.get("category") not in (None, "object"): + return [UnitResult(unit_id="target_0000", status=ResultStatus.NOT_APPLICABLE)] + error = sample.get("object_preparation_error") + if error: + exception = DataInvalidError if error["status"] == "data_invalid" else EvaluationError + raise exception(error["message"]) + kwargs = dict(tracker=self.tracker, redetector=self.redetector, reader=self.reader) + extra_scores, extra_metadata = {}, {} + if self.geometry: + track = object_track_analysis_for_item(sample, context, weight_metric=self.name, **kwargs) + config = context.metric_config(self.name) + size = int(config.get("canonical_canvas_size", DEFAULT_CANONICAL_SIZE)) + padding = int(config.get("canonical_padding", DEFAULT_CANONICAL_PADDING)) + result = analyze_object_geometry(track, canonical_size=size, canonical_padding=padding) + coverage = result.geometry_coverage + counts = result.n_revisit_masks + count_key = "n_return_masks" + extra_scores = { + "area_ratio_drift": result.area_ratio_drift, + "frame_boundary_contact_ratio": _frame_boundary_contact_ratio(track, result.phase.return_indices), + } + extra_metadata = { + "reference_frame_indices": list(result.reference_frame_indices), + "n_reference_masks": result.n_reference_masks, + "n_return_masks": result.n_return_masks, + "geometry_protocol": { + "canonical_canvas_size": size, + "canonical_padding": padding, + "alignment": "bbox_crop_center_uniform_scale", + "reference_matching": "max_iou_over_frozen_v_library", + "aggregation": "mean_valid_frames_per_revisit_macro", + }, + } + else: + analysis = object_appearance_analysis_for_item( + sample, context, feature_extractor=self.feature_extractor, **kwargs + ) + track, result = analysis.track, analysis.result + coverage = result.appearance_coverage + counts = result.n_revisit_features + count_key = "n_return_features" + extra_metadata = { + "n_reference_features": result.n_reference_features, + "n_return_features": result.n_return_features, + } + + phase = result.phase + expected = track.protocol.expected_revisits + kind = "geometry" if self.geometry else "appearance" + scores = { + "score": official_revisit_score(result.revisit_scores, expected_revisits=expected), + f"{kind}_similarity": result.similarity, + f"{kind}_coverage": coverage, + **extra_scores, + **revisit_score_fields(result.revisit_scores, result.revisit_coverages, expected_revisits=expected), + } + revisits = revisit_metadata( + phase, + result.revisit_scores, + result.revisit_coverages, + counts, + count_key=count_key, + expected_revisits=expected, + ) + if self.geometry: + for index, revisit in enumerate(revisits): + indices = phase.revisit_indices[index] if index < len(phase.revisit_indices) else () + revisit["frame_boundary_contact_ratio"] = _frame_boundary_contact_ratio(track, indices) + revisit["area_ratio_drift"] = ( + result.revisit_area_ratio_drifts[index] if index < len(result.revisit_area_ratio_drifts) else 0.0 + ) + phase_roles = {index: "V" for index in phase.visible_indices} + for index, indices in enumerate(phase.revisit_indices, 1): + phase_roles.update({frame: f"R{index}" for frame in indices}) + events = [] + for prompt in track.redetection_prompts: + outcome = "rejected_or_unstable" + for index, (start, end) in enumerate(track.visibility_segments): + if start <= prompt.frame_index <= end: + outcome = "already_tracked" if index == 0 else "accepted_reacquisition" + break + events.append( + { + "frame_index": prompt.frame_index, + "box_xyxy": list(prompt.box_xyxy), + "detector_score": prompt.detector_score, + "reid_similarity": prompt.reid_similarity, + "runner_up_reid_similarity": prompt.runner_up_reid_similarity, + "candidate_count": prompt.candidate_count, + "phase_role": phase_roles.get(prompt.frame_index, "D"), + "tracking_outcome": outcome, + } + ) + target = { + "target_id": 0, + "phase_status": phase.status, + "V": list(phase.visible_indices), + "D": list(phase.departure_indices), + "R": list(phase.return_indices), + "revisits": revisits, + f"{kind}_score": scores["score"], + f"{kind}_similarity": result.similarity, + f"{kind}_coverage": coverage, + **extra_scores, + **extra_metadata, + } + return [ + UnitResult( + unit_id="target_0000", + unit_type=f"target_object_{kind}", + scores=scores, + metadata={ + "video_path": track.video_path, + "n_targets": 1, + "condition_image_path": track.protocol.condition_image_path, + "target_mask_path": track.protocol.target_mask_path, + "target_description": track.protocol.target_description, + "expected_revisits": expected, + "revisit_mode": "auto" if expected is None else "fixed", + "alignment_method": track.alignment_method, + "frame0_alignment_score": track.frame0_alignment_score, + "phase_source": "grounding_dino_clip_redetection_with_multitracklet_sam2", + "phase_inference": { + "presence_logit_threshold": PRESENCE_LOGIT_THRESHOLD, + "min_mask_area_ratio": MIN_MASK_AREA_RATIO, + "max_mask_area_ratio": MAX_MASK_AREA_RATIO, + "min_mask_area_relative_to_seed": MIN_MASK_AREA_RELATIVE_TO_SEED, + "max_short_missed_frames": MAX_SHORT_MISSED_FRAMES, + "min_reappear_match_run": MIN_REAPPEAR_MATCH_RUN, + "max_revisit_gap_frames": track.max_revisit_gap_frames, + "revisit_allocator": "observed_return_components" + if expected is None + else "boundary_presence_components", + "revisit_boundary_buffer_frames": None if expected is None else REVISIT_BOUNDARY_BUFFER_FRAMES, + }, + "visibility_segments": [list(segment) for segment in track.visibility_segments], + "redetection_events": events, + "per_target": [target], + }, + ) + ] + + def aggregate_item(self, sample, units): + if len(units) != 1: + raise EvaluationError(f"{self.name} requires exactly one target unit") + unit = units[0] + return ItemResult( + item_id=sample["item_id"], + model_id=sample.get("model_id"), + scores=dict(unit.scores), + num_units=1, + status=unit.status, + error=unit.error, + metadata=dict(unit.metadata), + ) + + def aggregate_summary(self, items): + return aggregate_video_summary(items) + + +class ObjectAppearanceMetric(_ObjectMetric): + name = "object_appearance" + description = "Masked DINOv2-Large appearance consistency across object revisits." + + +class ObjectGeometryMetric(_ObjectMetric): + name = "object_geometry" + description = "Canonical silhouette IoU across object revisits, preserving the reference shape score." + geometry = True diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/object_redetector.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/object_redetector.py new file mode 100644 index 000000000..de3b3da79 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/object_redetector.py @@ -0,0 +1,527 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Text-prompted object re-detection independent from appearance scoring.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + +from .schemas import ObjectRedetectionPrompt + +REDETECTION_SEARCH_STRIDE = 4 +REDETECTION_BATCH_SIZE = 8 +REDETECTION_BOX_THRESHOLD = 0.30 +REDETECTION_TEXT_THRESHOLD = 0.25 +REDETECTION_NMS_THRESHOLD = 0.50 +REDETECTION_CANDIDATE_LIMIT = 64 +REID_BATCH_SIZE = 64 +REID_SIMILARITY_THRESHOLD = 0.75 +REID_HIGH_CONFIDENCE_THRESHOLD = 0.90 +REID_MIN_SIMILARITY_MARGIN = 0.04 +REID_CROP_CONTEXT_FRACTION = 0.10 +REDETECTION_MIN_CONFIRMATIONS = 2 +REDETECTION_MAX_SAMPLE_GAP = 2 +REDETECTION_MIN_BOX_AREA_RATIO = 0.0005 +REDETECTION_MAX_BOX_AREA_RATIO = 0.70 +REDETECTION_MAX_CENTER_DISTANCE_RATIO = 0.35 +REDETECTION_MAX_AREA_CHANGE = 10.0 +QUERY_CONTEXT_FRACTION = 0.10 + + +@dataclass(frozen=True) +class _DetectionCandidate: + detector_score: float + box_xyxy: tuple[float, float, float, float] + + +def reid_match_is_accepted( + top_similarity: float, + runner_up_similarity: float | None, + *, + similarity_threshold: float = REID_SIMILARITY_THRESHOLD, + high_confidence_threshold: float = REID_HIGH_CONFIDENCE_THRESHOLD, + min_similarity_margin: float = REID_MIN_SIMILARITY_MARGIN, +) -> bool: + """Accept an absolute-high match or a lower match with usable evidence. + + A lone candidate has no relative-margin evidence, but still has absolute + similarity evidence and must pass the downstream temporal confirmation. + """ + + if top_similarity < similarity_threshold: + return False + if top_similarity >= high_confidence_threshold: + return True + return runner_up_similarity is None or top_similarity - runner_up_similarity >= min_similarity_margin + + +class ObjectRedetectorBackend(Protocol): + def detect( + self, + text_prompt: str, + query_rgb: np.ndarray, + frames_rgb: Sequence[np.ndarray], + ) -> tuple[ObjectRedetectionPrompt, ...]: ... + + +def object_query_crop( + image_rgb: np.ndarray, + mask: np.ndarray, + *, + context_fraction: float = QUERY_CONTEXT_FRACTION, +) -> np.ndarray: + """Crop the prompted object while retaining a small amount of natural context.""" + + if image_rgb.ndim != 3 or image_rgb.shape[2] != 3: + raise ValueError("object query image must be RGB") + if mask.shape != image_rgb.shape[:2] or not mask.any(): + raise ValueError("object query mask must be non-empty and match the image") + if context_fraction < 0.0: + raise ValueError("context_fraction must be non-negative") + + ys, xs = np.nonzero(mask) + x1, x2 = int(xs.min()), int(xs.max()) + 1 + y1, y2 = int(ys.min()), int(ys.max()) + 1 + padding = max(2, int(round(max(x2 - x1, y2 - y1) * context_fraction))) + height, width = mask.shape + x1 = max(0, x1 - padding) + y1 = max(0, y1 - padding) + x2 = min(width, x2 + padding) + y2 = min(height, y2 + padding) + return np.ascontiguousarray(image_rgb[y1:y2, x1:x2]) + + +def _box_center_and_area( + box: tuple[float, float, float, float], +) -> tuple[float, float, float]: + x1, y1, x2, y2 = box + return (x1 + x2) * 0.5, (y1 + y2) * 0.5, (x2 - x1) * (y2 - y1) + + +def _boxes_share_object_hypothesis( + left: tuple[float, float, float, float], + right: tuple[float, float, float, float], + *, + intersection_over_smaller_threshold: float = 0.50, +) -> bool: + """Return whether two boxes are alternate crops of one spatial object.""" + + left_x1, left_y1, left_x2, left_y2 = left + right_x1, right_y1, right_x2, right_y2 = right + intersection_width = max(0.0, min(left_x2, right_x2) - max(left_x1, right_x1)) + intersection_height = max( + 0.0, + min(left_y2, right_y2) - max(left_y1, right_y1), + ) + intersection_area = intersection_width * intersection_height + if intersection_area <= 0.0: + return False + left_area = max(0.0, left_x2 - left_x1) * max(0.0, left_y2 - left_y1) + right_area = max(0.0, right_x2 - right_x1) * max( + 0.0, + right_y2 - right_y1, + ) + smaller_area = min(left_area, right_area) + return smaller_area > 0.0 and intersection_area / smaller_area >= intersection_over_smaller_threshold + + +def _temporally_compatible( + previous: ObjectRedetectionPrompt, + current: ObjectRedetectionPrompt, + *, + frame_height: int, + frame_width: int, + max_center_distance_ratio: float, + max_area_change: float, +) -> bool: + previous_x, previous_y, previous_area = _box_center_and_area(previous.box_xyxy) + current_x, current_y, current_area = _box_center_and_area(current.box_xyxy) + diagonal = max(1.0, float(np.hypot(frame_width, frame_height))) + center_distance = float(np.hypot(current_x - previous_x, current_y - previous_y) / diagonal) + if center_distance > max_center_distance_ratio: + return False + smaller_area = min(previous_area, current_area) + larger_area = max(previous_area, current_area) + return smaller_area > 0.0 and larger_area / smaller_area <= max_area_change + + +def confirm_redetection_prompts( + prompts: Sequence[ObjectRedetectionPrompt], + *, + frame_height: int, + frame_width: int, + search_stride: int = REDETECTION_SEARCH_STRIDE, + min_confirmations: int = REDETECTION_MIN_CONFIRMATIONS, + max_sample_gap: int = REDETECTION_MAX_SAMPLE_GAP, + max_center_distance_ratio: float = REDETECTION_MAX_CENTER_DISTANCE_RATIO, + max_area_change: float = REDETECTION_MAX_AREA_CHANGE, +) -> tuple[ObjectRedetectionPrompt, ...]: + """Keep only temporally persistent ReID matches, without using score-time DINO.""" + + if search_stride <= 0: + raise ValueError("search_stride must be positive") + if min_confirmations <= 0: + raise ValueError("min_confirmations must be positive") + if max_sample_gap <= 0: + raise ValueError("max_sample_gap must be positive") + ordered = sorted(prompts, key=lambda prompt: prompt.frame_index) + if not ordered: + return () + + runs: list[list[ObjectRedetectionPrompt]] = [] + current_run = [ordered[0]] + maximum_gap = search_stride * max_sample_gap + for prompt in ordered[1:]: + previous = current_run[-1] + compatible = 0 < prompt.frame_index - previous.frame_index <= maximum_gap and _temporally_compatible( + previous, + prompt, + frame_height=frame_height, + frame_width=frame_width, + max_center_distance_ratio=max_center_distance_ratio, + max_area_change=max_area_change, + ) + if compatible: + current_run.append(prompt) + else: + runs.append(current_run) + current_run = [prompt] + runs.append(current_run) + return tuple(prompt for run in runs if len(run) >= min_confirmations for prompt in run) + + +class GroundingDinoClipObjectRedetector: + """Generate candidates with Grounding DINO text-prompt detection and confirm with independent CLIP ReID.""" + + def __init__( + self, + *, + grounding_dino_checkpoint: str, + reid_checkpoint: str, + device: str, + search_stride: int = REDETECTION_SEARCH_STRIDE, + batch_size: int = REDETECTION_BATCH_SIZE, + box_threshold: float = REDETECTION_BOX_THRESHOLD, + text_threshold: float = REDETECTION_TEXT_THRESHOLD, + nms_threshold: float = REDETECTION_NMS_THRESHOLD, + candidate_limit: int = REDETECTION_CANDIDATE_LIMIT, + reid_batch_size: int = REID_BATCH_SIZE, + reid_similarity_threshold: float = REID_SIMILARITY_THRESHOLD, + reid_high_confidence_threshold: float = (REID_HIGH_CONFIDENCE_THRESHOLD), + reid_min_similarity_margin: float = REID_MIN_SIMILARITY_MARGIN, + reid_crop_context_fraction: float = REID_CROP_CONTEXT_FRACTION, + min_confirmations: int = REDETECTION_MIN_CONFIRMATIONS, + max_sample_gap: int = REDETECTION_MAX_SAMPLE_GAP, + min_box_area_ratio: float = REDETECTION_MIN_BOX_AREA_RATIO, + max_box_area_ratio: float = REDETECTION_MAX_BOX_AREA_RATIO, + ) -> None: + try: + import torch + from torchvision.ops import roi_align + from transformers import ( + AutoProcessor, + GroundingDinoForObjectDetection, + ) + + from worldfoundry.base_models.perception_core.general_perception import openai_clip as clip + except ImportError as error: + raise EvaluationError( + "Object Appearance re-detection requires Grounding DINO and " + "openai-clip; install the WorldMemBench metric requirements" + ) from error + + grounding_dino_checkpoint_path = Path(grounding_dino_checkpoint).expanduser().resolve() + reid_checkpoint_path = Path(reid_checkpoint).expanduser().resolve() + if not grounding_dino_checkpoint_path.exists(): + raise EvaluationError(f"Grounding DINO checkpoint not found: {grounding_dino_checkpoint_path}") + if not reid_checkpoint_path.is_file(): + raise EvaluationError(f"CLIP ReID checkpoint not found: {reid_checkpoint_path}") + self.processor = AutoProcessor.from_pretrained( + grounding_dino_checkpoint_path, + local_files_only=True, + ) + self.model = ( + GroundingDinoForObjectDetection.from_pretrained( + grounding_dino_checkpoint_path, + local_files_only=True, + ) + .to(device) + .eval() + ) + self.reid_model, self.reid_preprocess = clip.load( + str(reid_checkpoint_path), + device=device, + jit=False, + ) + self.reid_model.eval() + self.roi_align = roi_align + self.reid_mean = torch.tensor( + [0.48145466, 0.4578275, 0.40821073], + device=device, + ).view(1, 3, 1, 1) + self.reid_std = torch.tensor( + [0.26862954, 0.26130258, 0.27577711], + device=device, + ).view(1, 3, 1, 1) + + self.device = device + self.search_stride = search_stride + self.batch_size = batch_size + self.box_threshold = box_threshold + self.text_threshold = text_threshold + self.nms_threshold = nms_threshold + self.candidate_limit = candidate_limit + self.reid_batch_size = reid_batch_size + self.reid_similarity_threshold = reid_similarity_threshold + self.reid_high_confidence_threshold = reid_high_confidence_threshold + self.reid_min_similarity_margin = reid_min_similarity_margin + self.reid_crop_context_fraction = reid_crop_context_fraction + self.min_confirmations = min_confirmations + self.max_sample_gap = max_sample_gap + self.min_box_area_ratio = min_box_area_ratio + self.max_box_area_ratio = max_box_area_ratio + + def _candidates( + self, + result: dict[str, object], + *, + frame_height: int, + frame_width: int, + ) -> tuple[_DetectionCandidate, ...]: + boxes_tensor = result["boxes"] + scores_tensor = result["scores"] + boxes = boxes_tensor.detach().float().cpu().numpy() + scores = scores_tensor.detach().float().cpu().numpy() + frame_area = float(frame_height * frame_width) + candidates: list[_DetectionCandidate] = [] + for raw_box, raw_score in zip(boxes, scores, strict=True): + x1, y1, x2, y2 = (float(value) for value in raw_box) + x1 = min(max(x1, 0.0), float(frame_width)) + x2 = min(max(x2, 0.0), float(frame_width)) + y1 = min(max(y1, 0.0), float(frame_height)) + y2 = min(max(y2, 0.0), float(frame_height)) + area_ratio = max(0.0, x2 - x1) * max(0.0, y2 - y1) / frame_area + if np.isfinite(raw_score) and self.min_box_area_ratio <= area_ratio <= self.max_box_area_ratio: + candidates.append( + _DetectionCandidate( + detector_score=float(raw_score), + box_xyxy=(x1, y1, x2, y2), + ) + ) + if len(candidates) > 1: + import torch + from torchvision.ops import nms + + boxes_tensor = torch.tensor( + [candidate.box_xyxy for candidate in candidates], + dtype=torch.float32, + ) + scores_tensor = torch.tensor( + [candidate.detector_score for candidate in candidates], + dtype=torch.float32, + ) + keep = nms(boxes_tensor, scores_tensor, self.nms_threshold) + candidates = [candidates[int(index)] for index in keep.tolist()] + candidates.sort( + key=lambda candidate: candidate.detector_score, + reverse=True, + ) + return tuple(candidates[: self.candidate_limit]) + + def _reidentify( + self, + *, + query_embedding: object, + batch_indices: Sequence[int], + frames_rgb: Sequence[np.ndarray], + candidates_by_frame: Sequence[Sequence[_DetectionCandidate]], + ) -> list[ObjectRedetectionPrompt]: + import torch + + owners: list[tuple[int, _DetectionCandidate]] = [] + rois: list[list[float]] = [] + for owner, (frame_index, candidates) in enumerate(zip(batch_indices, candidates_by_frame, strict=True)): + height, width = frames_rgb[frame_index].shape[:2] + for candidate in candidates: + x1, y1, x2, y2 = candidate.box_xyxy + padding = max(x2 - x1, y2 - y1) * self.reid_crop_context_fraction + rois.append( + [ + float(owner), + max(0.0, x1 - padding), + max(0.0, y1 - padding), + min(float(width), x2 + padding), + min(float(height), y2 + padding), + ] + ) + owners.append((owner, candidate)) + if not rois: + return [] + + frame_tensor = torch.stack( + [torch.from_numpy(np.ascontiguousarray(frames_rgb[index])).permute(2, 0, 1) for index in batch_indices] + ).to(device=self.device, dtype=torch.float32) + frame_tensor = frame_tensor / 255.0 + crop_tensor = self.roi_align( + frame_tensor, + torch.tensor(rois, device=self.device, dtype=torch.float32), + output_size=(224, 224), + spatial_scale=1.0, + sampling_ratio=2, + aligned=True, + ) + crop_tensor = (crop_tensor - self.reid_mean) / self.reid_std + + similarities: list[float] = [] + with torch.inference_mode(): + for offset in range(0, len(crop_tensor), self.reid_batch_size): + features = self.reid_model.encode_image(crop_tensor[offset : offset + self.reid_batch_size]) + features = features / features.norm( + dim=-1, + keepdim=True, + ).clamp_min(1e-12) + batch_similarities = (features @ query_embedding.T).squeeze(1) + similarities.extend(batch_similarities.float().cpu().tolist()) + + ranked_by_owner: list[list[tuple[float, _DetectionCandidate]]] = [[] for _ in batch_indices] + for (owner, candidate), similarity in zip( + owners, + similarities, + strict=True, + ): + ranked_by_owner[owner].append((similarity, candidate)) + + prompts: list[ObjectRedetectionPrompt] = [] + for owner, ranked in enumerate(ranked_by_owner): + ranked.sort(key=lambda entry: entry[0], reverse=True) + frame_index = batch_indices[owner] + if not ranked: + continue + top_similarity, top_candidate = ranked[0] + runner_up_similarity = next( + ( + similarity + for similarity, candidate in ranked[1:] + if not _boxes_share_object_hypothesis( + top_candidate.box_xyxy, + candidate.box_xyxy, + ) + ), + None, + ) + if not reid_match_is_accepted( + top_similarity, + runner_up_similarity, + similarity_threshold=self.reid_similarity_threshold, + high_confidence_threshold=self.reid_high_confidence_threshold, + min_similarity_margin=self.reid_min_similarity_margin, + ): + continue + prompts.append( + ObjectRedetectionPrompt( + frame_index=frame_index, + box_xyxy=top_candidate.box_xyxy, + detector_score=top_candidate.detector_score, + reid_similarity=top_similarity, + runner_up_reid_similarity=runner_up_similarity, + candidate_count=len(ranked), + ) + ) + return prompts + + def detect( + self, + text_prompt: str, + query_rgb: np.ndarray, + frames_rgb: Sequence[np.ndarray], + ) -> tuple[ObjectRedetectionPrompt, ...]: + import torch + from PIL import Image + + if not frames_rgb: + return () + frame_height, frame_width = frames_rgb[0].shape[:2] + if any(frame.shape[:2] != (frame_height, frame_width) for frame in frames_rgb): + raise EvaluationError("Object Appearance requires constant-size video frames") + + text = text_prompt.strip() + if not text: + raise EvaluationError("Grounding DINO requires a non-empty text prompt") + sampled_indices = list(range(self.search_stride, len(frames_rgb), self.search_stride)) + raw_prompts: list[ObjectRedetectionPrompt] = [] + query_image = Image.fromarray(query_rgb) + with torch.inference_mode(): + query_embedding = self.reid_model.encode_image( + self.reid_preprocess(query_image).unsqueeze(0).to(self.device) + ) + query_embedding = query_embedding / query_embedding.norm( + dim=-1, + keepdim=True, + ).clamp_min(1e-12) + + for offset in range(0, len(sampled_indices), self.batch_size): + batch_indices = sampled_indices[offset : offset + self.batch_size] + target_images = [Image.fromarray(frames_rgb[index]) for index in batch_indices] + # transformers >= 5 treats a plain str or list[str] as one sample + # ("one image, many phrases"); list[list[str]] is the only form + # that yields one text batch entry per image. + inputs = self.processor( + images=target_images, + text=[[text]] * len(batch_indices), + return_tensors="pt", + ).to(self.device) + with torch.inference_mode(): + outputs = self.model(**inputs) + target_sizes = torch.tensor( + [[frame_height, frame_width]] * len(batch_indices), + device=self.device, + ) + results = self.processor.post_process_grounded_object_detection( + outputs=outputs, + input_ids=inputs.input_ids, + threshold=self.box_threshold, + text_threshold=self.text_threshold, + target_sizes=target_sizes, + ) + candidates_by_frame = [ + self._candidates( + result, + frame_height=frame_height, + frame_width=frame_width, + ) + for result in results + ] + raw_prompts.extend( + self._reidentify( + query_embedding=query_embedding, + batch_indices=batch_indices, + frames_rgb=frames_rgb, + candidates_by_frame=candidates_by_frame, + ) + ) + + return confirm_redetection_prompts( + raw_prompts, + frame_height=frame_height, + frame_width=frame_width, + search_stride=self.search_stride, + min_confirmations=self.min_confirmations, + max_sample_gap=self.max_sample_gap, + ) + + +__all__ = [ + "ObjectRedetectorBackend", + "GroundingDinoClipObjectRedetector", + "confirm_redetection_prompts", + "object_query_crop", + "reid_match_is_accepted", +] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/object_tracker.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/object_tracker.py new file mode 100644 index 000000000..868b804c4 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/object_tracker.py @@ -0,0 +1,317 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""First-frame-mask SAM2 tracking with independent object-presence evidence.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from pathlib import Path +from typing import Any, Protocol + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + +from ..revisit_protocol import MAX_SHORT_MISSED_FRAMES +from .schemas import ObjectRedetectionPrompt, ObjectTrackObservation + +PRESENCE_LOGIT_THRESHOLD = 0.0 +MIN_MASK_AREA_RATIO = 0.0005 +MAX_MASK_AREA_RATIO = 0.70 +MIN_MASK_AREA_RELATIVE_TO_SEED = 0.10 +MIN_REDETECTED_TRACKLET_FRAMES = 3 + + +def _anchor_connected_presence_indices( + observations: dict[int, ObjectTrackObservation], + *, + total_frames: int, + anchor_index: int, + max_short_missed_frames: int = MAX_SHORT_MISSED_FRAMES, +) -> tuple[int, ...]: + presence = [bool(observations.get(index) and observations[index].present) for index in range(total_frames)] + index = 0 + while index < total_frames: + if presence[index]: + index += 1 + continue + start = index + while index < total_frames and not presence[index]: + index += 1 + if start > 0 and index < total_frames and index - start <= max_short_missed_frames: + presence[start:index] = [True] * (index - start) + if not presence[anchor_index]: + return () + start = anchor_index + while start > 0 and presence[start - 1]: + start -= 1 + end = anchor_index + while end + 1 < total_frames and presence[end + 1]: + end += 1 + return tuple(range(start, end + 1)) + + +class ObjectTrackerBackend(Protocol): + def track( + self, + frames_rgb: Sequence[np.ndarray], + first_frame_mask: np.ndarray, + redetection_prompts: Sequence[ObjectRedetectionPrompt] = (), + ) -> dict[int, ObjectTrackObservation]: ... + + +class TransformersSam2ObjectTracker: + """Track one seeded object and preserve SAM2's occlusion/presence logit.""" + + def __init__( + self, + *, + checkpoint: str, + device: str, + presence_logit_threshold: float = PRESENCE_LOGIT_THRESHOLD, + min_mask_area_ratio: float = MIN_MASK_AREA_RATIO, + max_mask_area_ratio: float = MAX_MASK_AREA_RATIO, + min_mask_area_relative_to_seed: float = MIN_MASK_AREA_RELATIVE_TO_SEED, + dtype: str = "bfloat16", + vision_cache_size: int | None = 1, + ) -> None: + try: + import torch + from transformers import Sam2VideoModel, Sam2VideoProcessor + except ImportError as error: + raise EvaluationError( + "Object Appearance requires transformers with Sam2VideoModel; " + "install the WorldMemBench metric requirements" + ) from error + checkpoint_path = Path(checkpoint).expanduser().resolve() + if not checkpoint_path.exists(): + raise EvaluationError(f"SAM2 checkpoint not found: {checkpoint_path}") + if vision_cache_size is not None and ( + isinstance(vision_cache_size, bool) or not isinstance(vision_cache_size, int) or vision_cache_size < 1 + ): + raise ValueError("sam2_vision_cache_size must be a positive integer or null (all frames)") + self.vision_cache_size = vision_cache_size + self._dtype = getattr(torch, dtype) + self.processor = Sam2VideoProcessor.from_pretrained( + checkpoint_path, + local_files_only=True, + ) + self.model = ( + Sam2VideoModel.from_pretrained( + checkpoint_path, + local_files_only=True, + torch_dtype=self._dtype, + ) + .to(device) + .eval() + ) + self.device = device + self.presence_logit_threshold = presence_logit_threshold + self.min_mask_area_ratio = min_mask_area_ratio + self.max_mask_area_ratio = max_mask_area_ratio + self.min_mask_area_relative_to_seed = min_mask_area_relative_to_seed + + def _observation( + self, + output: Any, + *, + height: int, + width: int, + minimum_mask_area_ratio: float, + ) -> ObjectTrackObservation: + masks, logits = output.pred_masks, output.object_score_logits + if masks is None or logits is None: + raise EvaluationError("SAM2 must return masks and object_score_logits; check the installed version") + frame_index = int(output.frame_idx) + + processed = self.processor.post_process_masks( + [masks], + original_sizes=[[height, width]], + binarize=True, + )[0] + binary = processed.detach().cpu().numpy() + while binary.ndim > 2: + binary = binary[0] + mask = np.asarray(binary, dtype=bool) + if mask.shape != (height, width): + raise EvaluationError(f"SAM2 returned mask shape {mask.shape}, expected {(height, width)}") + + presence_logit = float(logits.detach().float().cpu().numpy().reshape(-1)[0]) + if not math.isfinite(presence_logit): + raise EvaluationError("SAM2 returned a non-finite object score") + probability = 1.0 / (1.0 + math.exp(-max(-60.0, min(60.0, presence_logit)))) + area_ratio = float(mask.mean()) + mask_valid = bool(mask.any()) and minimum_mask_area_ratio <= area_ratio <= self.max_mask_area_ratio + present = presence_logit > self.presence_logit_threshold and mask_valid + return ObjectTrackObservation( + frame_index=frame_index, + mask=mask, + presence_logit=presence_logit, + presence_probability=probability, + mask_area_ratio=area_ratio, + mask_valid=mask_valid, + present=present, + ) + + def track( + self, + frames_rgb: Sequence[np.ndarray], + first_frame_mask: np.ndarray, + redetection_prompts: Sequence[ObjectRedetectionPrompt] = (), + ) -> dict[int, ObjectTrackObservation]: + import torch + from PIL import Image + + if not frames_rgb: + return {} + height, width = frames_rgb[0].shape[:2] + if first_frame_mask.shape != (height, width) or not first_frame_mask.any(): + raise EvaluationError("the aligned first-frame object mask is empty or has the wrong size") + ordered_prompts = tuple(sorted(redetection_prompts, key=lambda prompt: prompt.frame_index)) + for prompt in ordered_prompts: + if prompt.frame_index <= 0 or prompt.frame_index >= len(frames_rgb): + raise EvaluationError("SAM2 redetection prompt frame is out of range") + x1, y1, x2, y2 = prompt.box_xyxy + if not (x2 > x1 and y2 > y1): + raise EvaluationError("SAM2 redetection prompt box is invalid") + + minimum_mask_area_ratio = max( + self.min_mask_area_ratio, + float(first_frame_mask.mean()) * self.min_mask_area_relative_to_seed, + ) + pil_frames = [Image.fromarray(frame) for frame in frames_rgb] + # None keeps all frames for large GPUs; the default recomputes evicted + # image features. Neither mode drops temporal tracking state or frames. + session = self.processor.init_video_session( + video=pil_frames, + inference_device=self.device, + max_vision_features_cache_size=self.vision_cache_size or len(pil_frames), + dtype=self._dtype, + ) + + def store( + destination: dict[int, ObjectTrackObservation], + observation: ObjectTrackObservation, + ) -> None: + previous = destination.get(observation.frame_index) + if ( + previous is None + or (observation.present and not previous.present) + or ( + observation.present + and previous.present + and observation.presence_probability > previous.presence_probability + ) + ): + destination[observation.frame_index] = observation + + def run_tracklet( + *, + anchor_index: int, + input_mask: np.ndarray | None = None, + input_box: tuple[float, float, float, float] | None = None, + propagate_reverse: bool, + ) -> dict[int, ObjectTrackObservation]: + if input_mask is not None: + self.processor.add_inputs_to_inference_session( + session, + frame_idx=anchor_index, + obj_ids=0, + input_masks=np.asarray(input_mask, dtype=bool), + clear_old_inputs=True, + ) + elif input_box is not None: + x1, y1, x2, y2 = input_box + center_point = [(x1 + x2) * 0.5, (y1 + y2) * 0.5] + self.processor.add_inputs_to_inference_session( + session, + frame_idx=anchor_index, + obj_ids=0, + input_points=[[[center_point]]], + input_labels=[[[1]]], + input_boxes=[[[float(value) for value in input_box]]], + clear_old_inputs=True, + ) + else: + raise EvaluationError("SAM2 tracklet requires a mask or box prompt") + + tracklet: dict[int, ObjectTrackObservation] = {} + with torch.inference_mode(): + conditioned = self.model(session, frame_idx=anchor_index) + store( + tracklet, + self._observation( + conditioned, + height=height, + width=width, + minimum_mask_area_ratio=minimum_mask_area_ratio, + ), + ) + directions = [False] + if propagate_reverse: + directions.insert(0, True) + for reverse in directions: + iterator = self.model.propagate_in_video_iterator( + session, + start_frame_idx=anchor_index, + reverse=reverse, + ) + for output in iterator: + store( + tracklet, + self._observation( + output, + height=height, + width=width, + minimum_mask_area_ratio=minimum_mask_area_ratio, + ), + ) + + connected_indices = _anchor_connected_presence_indices( + tracklet, + total_frames=len(frames_rgb), + anchor_index=anchor_index, + ) + return { + index: tracklet[index] for index in connected_indices if index in tracklet and tracklet[index].present + } + + observations = run_tracklet( + anchor_index=0, + input_mask=first_frame_mask, + propagate_reverse=False, + ) + accepted_indices = set(observations) + for prompt in ordered_prompts: + existing = observations.get(prompt.frame_index) + if existing is not None and existing.present: + continue + recovered = run_tracklet( + anchor_index=prompt.frame_index, + input_box=prompt.box_xyxy, + propagate_reverse=True, + ) + if len(recovered) < MIN_REDETECTED_TRACKLET_FRAMES: + continue + if accepted_indices.intersection(recovered): + continue + for observation in recovered.values(): + store(observations, observation) + accepted_indices.update(recovered) + + for frame_index in range(len(frames_rgb)): + observations.setdefault( + frame_index, + ObjectTrackObservation( + frame_index=frame_index, + mask=None, + presence_logit=None, + presence_probability=0.0, + mask_area_ratio=0.0, + mask_valid=False, + present=False, + ), + ) + return observations diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/preparation.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/preparation.py new file mode 100644 index 000000000..662d7b8c7 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/preparation.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Read a frozen source mask and target description, without generating annotations.""" + +import json +from pathlib import Path + +import numpy as np + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError + +SOURCE_PATH_FIELDS = ("condition_image", "target_object_mask_first_frame") + + +def source_path(sample: dict, field: str) -> str: + value = (sample.get("annotations") or {}).get(field) + if not isinstance(value, str) or not value.strip(): + raise DataInvalidError(f"missing annotations.{field}") + path = Path(value).expanduser().resolve() + return str(path) + + +def load_source(image_path: str, mask_path: str) -> tuple[np.ndarray, np.ndarray]: + from PIL import Image + + try: + with Image.open(image_path) as image: + rgb = np.asarray(image.convert("RGB")) + raw = np.load(mask_path, allow_pickle=False) + except Exception as error: + raise DataInvalidError(f"cannot read object source image/mask: {error}") from error + if not isinstance(raw, np.ndarray) or raw.ndim != 2 or raw.shape != rgb.shape[:2]: + raise DataInvalidError("object_mask.npy must be a 2-D array matching the source image size") + # Accept boolean masks, 0/1 arrays, and the common 0/255 binary encoding. + if not (np.isin(raw, [0, 1]).all() or np.isin(raw, [0, 255]).all()): + raise DataInvalidError("object_mask.npy must be binary (bool, 0/1, or 0/255)") + mask = raw.astype(bool) + if not mask.any(): + raise DataInvalidError("object_mask.npy cannot be empty") + return rgb, mask + + +def prepare_for_evaluation(samples: list[dict], *, expected_revisits: int | None) -> None: + """Resolve case sidecars once; source contents are checked when scoring.""" + for sample in samples: + if sample.get("category") not in (None, "object"): + continue + sample.pop("object_preparation_error", None) + try: + annotations = dict(sample.get("annotations") or {}) + if not all(field in annotations for field in (*SOURCE_PATH_FIELDS, "target")): + case_dir = Path(sample["video"]).parent + prompt = json.loads((case_dir / "prompt.json").read_text(encoding="utf-8")) + obj = prompt["object"] + annotations.update( + condition_image=str((case_dir / prompt.get("first_frame", "first_frame.jpg")).resolve()), + target_object_mask_first_frame=str( + (case_dir / obj.get("object_mask", "object_mask.npy")).resolve() + ), + target=obj["target"], + ) + if expected_revisits is not None: + annotations["expected_revisits"] = expected_revisits + else: + annotations.setdefault("expected_revisits", None) + sample["annotations"] = annotations + except (OSError, ValueError, KeyError, TypeError, AttributeError) as error: + sample["object_preparation_error"] = { + "status": "data_invalid", + "message": f"invalid Object prompt/annotations: {error}", + } diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/protocol.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/protocol.py new file mode 100644 index 000000000..11b7b3720 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/protocol.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Single-object input contract and fixed/automatic revisit selection.""" + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError + +from ..revisit_protocol import validate_expected_revisits +from .schemas import ObjectEntityProtocol + + +def load_object_entity_protocol( + item: dict, *, total_frames: int, condition_image_path: str, target_mask_path: str +) -> ObjectEntityProtocol: + annotations = item.get("annotations") or {} + target = annotations.get("target") + if not isinstance(target, str) or not target.strip(): + raise DataInvalidError("annotations.target must be a non-empty object description") + revisits = validate_expected_revisits(annotations.get("expected_revisits")) + if total_frames < (revisits or 1) + 1: + raise DataInvalidError("generated video is too short for the requested object revisit split") + return ObjectEntityProtocol( + total_frames=total_frames, + condition_image_path=condition_image_path, + target_mask_path=target_mask_path, + target_description=target.strip(), + expected_revisits=revisits, + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/revisit.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/revisit.py new file mode 100644 index 000000000..f10e48e4c --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/revisit.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Revisit scores and readable segment metadata for a single object.""" + +from ..revisit_protocol import MAX_EXPECTED_REVISITS + + +def revisit_score_fields(scores, coverages, *, expected_revisits): + if expected_revisits is None: + return {} + return { + f"revisit_{index + 1}_{name}": float(values[index]) if index < len(values) else 0.0 + for index in range(MAX_EXPECTED_REVISITS) + for name, values in (("score", scores), ("coverage", coverages)) + } + + +def official_revisit_score(scores, *, expected_revisits): + count = len(scores) if expected_revisits is None else expected_revisits + return float(sum(scores[:count]) / count) if count else 0.0 + + +def revisit_metadata(phase, scores, coverages, counts, *, count_key, expected_revisits): + count = len(scores) if expected_revisits is None else MAX_EXPECTED_REVISITS + fields = { + "indices": (phase.revisit_indices, ()), + "slot": (phase.revisit_slots, ()), + "nominal_boundary": (phase.revisit_nominal_boundaries, None), + "search_window": (phase.revisit_search_windows, ()), + "window_source": (phase.revisit_sources, "missing"), + "departure_evidence": (phase.revisit_departure_evidence, False), + "score": (scores, 0.0), + "coverage": (coverages, 0.0), + count_key: (counts, 0), + } + output = [] + for index in range(count): + row = {"revisit_id": index + 1} + for name, (values, default) in fields.items(): + value = values[index] if index < len(values) else default + row[name] = list(value) if isinstance(value, tuple) else value + indices = row["indices"] + row["actual_window"] = [indices[0], indices[-1] + 1] if indices else [] + output.append(row) + return output diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/schemas.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/schemas.py new file mode 100644 index 000000000..346ecbd61 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/object/schemas.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Typed records for Object Entity evaluation.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np + +from ..revisit_protocol import TargetPhase + + +@dataclass(frozen=True) +class ObjectEntityProtocol: + total_frames: int + condition_image_path: str + target_mask_path: str + target_description: str = "" + expected_revisits: int | None = None + + @property + def required_frame_indices(self) -> tuple[int, ...]: + return tuple(range(self.total_frames)) + + +@dataclass(frozen=True) +class ObjectTrackObservation: + frame_index: int + mask: np.ndarray | None = field(repr=False) + presence_logit: float | None + presence_probability: float + mask_area_ratio: float + mask_valid: bool + present: bool + + +@dataclass(frozen=True) +class ObjectRedetectionPrompt: + frame_index: int + box_xyxy: tuple[float, float, float, float] + detector_score: float + reid_similarity: float = 0.0 + runner_up_reid_similarity: float | None = None + candidate_count: int = 1 + + +@dataclass +class ObjectTrackAnalysis: + protocol: ObjectEntityProtocol + phase: TargetPhase + observations: dict[int, ObjectTrackObservation] + decoded_frames: dict[int, np.ndarray] = field(repr=False) + video_path: str = "" + alignment_method: str = "" + frame0_alignment_score: float = 0.0 + redetection_prompts: tuple[ObjectRedetectionPrompt, ...] = () + visibility_segments: tuple[tuple[int, int], ...] = () + max_revisit_gap_frames: int = 0 + + +@dataclass(frozen=True) +class ObjectAppearanceResult: + phase: TargetPhase + frame_scores: dict[int, float] + similarity: float + appearance_coverage: float + n_reference_features: int + n_return_features: int + revisit_scores: tuple[float, ...] + revisit_coverages: tuple[float, ...] + n_revisit_features: tuple[int, ...] + + +@dataclass +class ObjectAppearanceAnalysis: + track: ObjectTrackAnalysis + result: ObjectAppearanceResult + + +@dataclass(frozen=True) +class ObjectGeometryResult: + phase: TargetPhase + reference_frame_indices: tuple[int, ...] + frame_scores: dict[int, float] + similarity: float + geometry_coverage: float + area_ratio_drift: float + n_reference_masks: int + n_return_masks: int + revisit_scores: tuple[float, ...] + revisit_coverages: tuple[float, ...] + n_revisit_masks: tuple[int, ...] + revisit_area_ratio_drifts: tuple[float, ...] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/registry.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/registry.py new file mode 100644 index 000000000..1c43dcd7b --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/registry.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Explicit name-based registration of metric implementations.""" + +import re + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.base import Metric + + +class MetricRegistry: + def __init__(self) -> None: + self._metrics: dict[str, Metric] = {} + + def register(self, metric: Metric) -> None: + # Metric names are also output-directory components. + if not re.fullmatch(r"[a-z][a-z0-9_-]*", metric.name): + raise ValueError("metric names must match [a-z][a-z0-9_-]*") + if metric.name in self._metrics: + raise ValueError(f"metric already registered: {metric.name}") + self._metrics[metric.name] = metric + + def names(self) -> list[str]: + return sorted(self._metrics) + + def get(self, name: str) -> Metric: + if name not in self._metrics: + raise ValueError(f"unknown metric: {name}; available: {', '.join(self.names())}") + return self._metrics[name] diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/revisit_protocol.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/revisit_protocol.py new file mode 100644 index 000000000..7567fc254 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/metrics/revisit_protocol.py @@ -0,0 +1,426 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Shared Human/Object phase allocation from target-presence runs. + +The protocol has ``K`` possible departure gaps:: + + V [D?] R1 [D?] R2 ... [D?] RK + +Presence components are matched to the fixed revisit boundaries. A component +that covers several consecutive boundaries is split evenly across those +revisits. This represents all ``2**K`` departure-insertion patterns without a +separate never-departed fallback. +""" + +from __future__ import annotations + +from collections.abc import Collection, Sequence +from dataclasses import dataclass +from itertools import combinations + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import DataInvalidError + +MAX_SHORT_MISSED_FRAMES = 2 +MIN_REAPPEAR_MATCH_RUN = 3 +MAX_EXPECTED_REVISITS = 3 +REVISIT_BOUNDARY_BUFFER_FRAMES = 30 + + +@dataclass(frozen=True) +class TargetPhase: + target_id: int + visible_indices: tuple[int, ...] + departure_indices: tuple[int, ...] + return_indices: tuple[int, ...] + status: str + revisit_indices: tuple[tuple[int, ...], ...] = () + revisit_slots: tuple[tuple[int, int], ...] = () + revisit_nominal_boundaries: tuple[int, ...] = () + revisit_search_windows: tuple[tuple[int, int], ...] = () + revisit_sources: tuple[str, ...] = () + revisit_departure_evidence: tuple[bool, ...] = () + + @property + def is_scorable(self) -> bool: + return self.status == "ok" or (self.status == "not_departed" and bool(self.return_indices)) + + +def validate_expected_revisits(value: object) -> int | None: + if value is not None and ( + isinstance(value, bool) or not isinstance(value, int) or not 1 <= value <= MAX_EXPECTED_REVISITS + ): + raise DataInvalidError("expected_revisits must be null (automatic) or an integer in [1, 3]") + return value + + +def _balanced_ranges( + start: int, + end: int, + parts: int, +) -> tuple[tuple[int, ...], ...]: + """Split ``[start, end)`` exactly once, with earlier parts one frame longer.""" + + if parts < 1: + raise ValueError("parts must be positive") + length = max(0, end - start) + base, remainder = divmod(length, parts) + ranges: list[tuple[int, ...]] = [] + cursor = start + for part_index in range(parts): + size = base + int(part_index < remainder) + ranges.append(tuple(range(cursor, cursor + size))) + cursor += size + return tuple(ranges) + + +def _close_short_internal_gaps( + presence: list[bool], + *, + max_gap: int, +) -> list[bool]: + smoothed = list(presence) + index = 0 + while index < len(smoothed): + if smoothed[index]: + index += 1 + continue + start = index + while index < len(smoothed) and not smoothed[index]: + index += 1 + if start > 0 and index < len(smoothed) and index - start <= max_gap: + smoothed[start:index] = [True] * (index - start) + return smoothed + + +def _true_runs(values: Sequence[bool], *, start: int = 0) -> list[tuple[int, int]]: + runs: list[tuple[int, int]] = [] + index = start + while index < len(values): + if not values[index]: + index += 1 + continue + run_start = index + while index < len(values) and values[index]: + index += 1 + runs.append((run_start, index - 1)) + return runs + + +def _merge_nearby_runs( + runs: list[tuple[int, int]], + *, + max_gap: int, +) -> list[tuple[int, int]]: + merged: list[tuple[int, int]] = [] + for start, end in runs: + if merged and start - merged[-1][1] - 1 <= max_gap: + merged[-1] = (merged[-1][0], end) + else: + merged.append((start, end)) + return merged + + +def _slot_geometry( + total_frames: int, + expected_revisits: int, + buffer_frames: int, +) -> tuple[ + tuple[tuple[int, int], ...], + tuple[int, ...], + tuple[tuple[int, int], ...], +]: + slots = tuple((indices[0], indices[-1] + 1) for indices in _balanced_ranges(0, total_frames, expected_revisits)) + boundaries = tuple(end for _, end in slots) + if expected_revisits == 1: + return slots, boundaries, ((0, total_frames),) + + effective_buffer = min( + buffer_frames, + min(end - start for start, end in slots) // 2, + ) + windows = tuple( + ( + max(slot_start, slot_end - effective_buffer), + min(total_frames, slot_end + effective_buffer), + ) + for slot_start, slot_end in slots + ) + return slots, boundaries, windows + + +def _overlap_length( + component: tuple[int, int], + window: tuple[int, int], +) -> int: + start, end = component + window_start, window_end = window + return max(0, min(end + 1, window_end) - max(start, window_start)) + + +def _boundary_distance(component: tuple[int, int], boundary: int) -> int: + start, end = component + if start <= boundary <= end + 1: + return 0 + return min(abs(start - boundary), abs(end + 1 - boundary)) + + +def _is_monotonic_assignment( + assignments: dict[int, int], +) -> bool: + components = [assignments[index] for index in sorted(assignments)] + return components == sorted(components) + + +def _assign_components( + components: Sequence[tuple[int, int]], + *, + boundaries: Sequence[int], + windows: Sequence[tuple[int, int]], + min_overlap_frames: int, +) -> dict[int, int]: + """Map revisit boundaries to ordered visible components. + + Boundary-window overlap is authoritative. Remaining post-departure + components are matched to the nearest free boundaries in monotonic order. + The initial component is never force-matched: without boundary overlap it + remains the full V reference. + """ + + assignments: dict[int, int] = {} + last_component = 0 + for boundary_index, (boundary, window) in enumerate(zip(boundaries, windows)): + candidates = [] + for component_index in range(last_component, len(components)): + component = components[component_index] + overlap = _overlap_length(component, window) + if overlap < min_overlap_frames: + continue + candidates.append( + ( + overlap, + -_boundary_distance(component, boundary), + -component_index, + component_index, + ) + ) + if candidates: + component_index = max(candidates)[-1] + assignments[boundary_index] = component_index + last_component = component_index + + used_components = set(assignments.values()) + unmatched_post = [index for index in range(1, len(components)) if index not in used_components] + free_boundaries = [index for index in range(len(boundaries)) if index not in assignments] + pair_count = min(len(unmatched_post), len(free_boundaries)) + if not pair_count: + return assignments + + best: tuple | None = None + for component_indices in combinations(unmatched_post, pair_count): + for boundary_indices in combinations(free_boundaries, pair_count): + proposed = dict(assignments) + proposed.update(zip(boundary_indices, component_indices)) + if not _is_monotonic_assignment(proposed): + continue + distances = tuple( + _boundary_distance( + components[component_index], + boundaries[boundary_index], + ) + for boundary_index, component_index in zip( + boundary_indices, + component_indices, + ) + ) + total_length = sum(components[index][1] - components[index][0] + 1 for index in component_indices) + key = ( + sum(distances), + max(distances), + -total_length, + boundary_indices, + component_indices, + ) + if best is None or key < best[0]: + best = (key, proposed) + return best[1] if best is not None else assignments + + +def _flatten_revisits( + revisits: Sequence[Sequence[int]], +) -> tuple[int, ...]: + return tuple(index for revisit in revisits for index in revisit) + + +def phase_from_revisit_runs( + *, + target_id: int, + visible_indices: tuple[int, ...], + total_frames: int, + runs: list[tuple[int, int]], + expected_revisits: int | None, + revisit_boundary_buffer_frames: int, +) -> TargetPhase: + """Allocate initial visibility and post-departure runs to R1..RK.""" + + if not visible_indices or visible_indices[0] != 0: + raise ValueError("visible_indices must be a non-empty frame-0 run") + initial_component = (0, visible_indices[-1]) + post_components = sorted( + (max(start, initial_component[1] + 1), end) for start, end in runs if end > initial_component[1] + ) + post_components = [run for run in post_components if run[0] <= run[1]] + if expected_revisits is None: + # Keep observed return components intact, without fixed-time boundary + # allocation. The reference's no-departure/no-return scoring stays intact. + not_departed = initial_component[1] == total_frames - 1 + reference = (0,) if not_departed else visible_indices + revisits = ( + (tuple(range(1, total_frames)),) + if not_departed + else tuple(tuple(range(start, end + 1)) for start, end in post_components) + ) + returned = _flatten_revisits(revisits) + used = set(reference) | set(returned) + return TargetPhase( + target_id=target_id, + visible_indices=reference, + departure_indices=tuple(index for index in range(total_frames) if index not in used), + return_indices=returned, + status="not_departed" if not_departed else "ok" if returned else "not_returned", + revisit_indices=revisits, + revisit_sources=tuple("uniform_not_departed" if not_departed else "observed_cluster" for _ in revisits), + revisit_departure_evidence=tuple(not not_departed for _ in revisits), + ) + components = [initial_component, *post_components] + slots, boundaries, windows = _slot_geometry( + total_frames, + expected_revisits, + revisit_boundary_buffer_frames, + ) + + if expected_revisits == 1: + if post_components: + component_index = min( + range(1, len(components)), + key=lambda index: ( + _boundary_distance(components[index], boundaries[0]), + -components[index][1], + ), + ) + assignments = {0: component_index} + elif initial_component[1] == total_frames - 1: + assignments = {0: 0} + else: + assignments = {} + else: + assignments = _assign_components( + components, + boundaries=boundaries, + windows=windows, + min_overlap_frames=MIN_REAPPEAR_MATCH_RUN, + ) + + component_boundaries: dict[int, list[int]] = {} + for boundary_index, component_index in assignments.items(): + component_boundaries.setdefault(component_index, []).append(boundary_index) + + revisits: list[tuple[int, ...]] = [() for _ in range(expected_revisits)] + sources = ["missing" for _ in range(expected_revisits)] + departure_evidence = [False for _ in range(expected_revisits)] + initial_carries_revisit = 0 in component_boundaries + for component_index, assigned_boundaries in component_boundaries.items(): + assigned_boundaries.sort() + start, end = components[component_index] + if component_index == 0: + start = max(1, start) + segments = _balanced_ranges(start, end + 1, len(assigned_boundaries)) + for local_index, (boundary_index, segment) in enumerate(zip(assigned_boundaries, segments)): + revisits[boundary_index] = segment + if component_index == 0: + sources[boundary_index] = ( + "uniform_not_departed" + if initial_component[1] == total_frames - 1 + else "continuous_presence_until_departure" + ) + else: + sources[boundary_index] = "observed_cluster" if local_index == 0 else "continuous_presence_split" + departure_evidence[boundary_index] = local_index == 0 + + reference_indices = ( + (0,) if initial_carries_revisit else tuple(range(initial_component[0], initial_component[1] + 1)) + ) + return_indices = _flatten_revisits(revisits) + used = set(reference_indices) | set(return_indices) + departure_indices = tuple(index for index in range(total_frames) if index not in used) + if initial_component[1] == total_frames - 1: + status = "not_departed" + elif return_indices: + status = "ok" + else: + status = "not_returned" + return TargetPhase( + target_id=target_id, + visible_indices=reference_indices, + departure_indices=departure_indices, + return_indices=return_indices, + status=status, + revisit_indices=tuple(revisits), + revisit_slots=slots, + revisit_nominal_boundaries=boundaries, + revisit_search_windows=windows, + revisit_sources=tuple(sources), + revisit_departure_evidence=tuple(departure_evidence), + ) + + +def infer_target_phase( + *, + target_id: int, + total_frames: int, + matched_frame_indices: Collection[int], + max_short_missed_frames: int = MAX_SHORT_MISSED_FRAMES, + min_reappear_match_run: int = MIN_REAPPEAR_MATCH_RUN, + max_revisit_gap_frames: int = MAX_SHORT_MISSED_FRAMES, + expected_revisits: int | None = None, + revisit_boundary_buffer_frames: int | None = None, +) -> TargetPhase: + """Infer V/D/R1..RK from one modality-specific presence sequence.""" + + if max_short_missed_frames < 0 or max_revisit_gap_frames < 0: + raise ValueError("missed-frame gap limits must be non-negative") + if min_reappear_match_run < 1: + raise ValueError("min_reappear_match_run must be positive") + validate_expected_revisits(expected_revisits) + buffer_frames = ( + REVISIT_BOUNDARY_BUFFER_FRAMES if revisit_boundary_buffer_frames is None else revisit_boundary_buffer_frames + ) + if buffer_frames < 0: + raise ValueError("revisit_boundary_buffer_frames must be non-negative") + + matched = {int(index) for index in matched_frame_indices} + matched.add(0) + presence = [index in matched for index in range(total_frames)] + smoothed = _close_short_internal_gaps( + presence, + max_gap=max_short_missed_frames, + ) + initial_runs = _true_runs(smoothed) + if not initial_runs or initial_runs[0][0] != 0: + raise AssertionError("the seeded first-frame presence run is missing") + visible_end = initial_runs[0][1] + visible_indices = tuple(range(visible_end + 1)) + post_runs = [ + (start, end) + for start, end in _true_runs(presence, start=visible_end + 1) + if end - start + 1 >= min_reappear_match_run + ] + post_runs = _merge_nearby_runs(post_runs, max_gap=max_revisit_gap_frames) + return phase_from_revisit_runs( + target_id=target_id, + visible_indices=visible_indices, + total_frames=total_frames, + runs=post_runs, + expected_revisits=expected_revisits, + revisit_boundary_buffer_frames=buffer_frames, + ) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/paths.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/paths.py new file mode 100644 index 000000000..402de22ff --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/paths.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 WorldFoundry contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Default location for explicitly staged WorldMemBench metric assets.""" + +import os +from pathlib import Path + +from worldfoundry.core.io.paths import checkpoint_root_path + +DEFAULT_WEIGHTS_DIR = Path( + os.environ.get("WORLDFOUNDRY_WORLDMEMBENCH_WEIGHTS_DIR", str(checkpoint_root_path("worldmembench"))) +).expanduser() diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/__init__.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/__init__.py new file mode 100644 index 000000000..e82410ffd --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Adapters for the official pose-estimation model runtimes.""" diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/da3.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/da3.py new file mode 100644 index 000000000..1d718c9b6 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/da3.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Prepare DA3 poses and scene geometry for one video.mp4 or a folder of cases.""" + +import os +import subprocess +import sys +from pathlib import Path +from tempfile import TemporaryDirectory, TemporaryFile + + +def streaming_config(weights: Path) -> dict: + """Frozen reference settings; Python Sim3 is the reference runtime's solver.""" + model = weights / "DA3NESTED-GIANT-LARGE-1.1" + return { + "Weights": { + "DA3": str(model / "model.safetensors"), + "DA3_CONFIG": str(model / "config.json"), + "SALAD": str(weights / "dino_salad.ckpt"), + }, + "Model": { + "chunk_size": 120, + "overlap": 60, + "loop_chunk_size": 20, + "loop_enable": True, + "useDBoW": False, + "delete_temp_files": True, + "align_lib": "torch", + "align_method": "sim3", + "scale_compute_method": "auto", + "align_type": "dense", + "ref_view_strategy": "middle", + "ref_view_strategy_loop": "saddle_balanced", + "depth_threshold": 15.0, + "save_depth_conf_result": True, + "save_debug_info": False, + "Sparse_Align": {"keypoint_select": "orb", "keypoint_num": 5000}, + # Upstream evaluates the reference YAML's scientific-notation strings. + "IRLS": {"delta": 0.1, "max_iters": 5, "tol": "1e-9"}, + "Pointcloud_Save": {"sample_ratio": 0.015, "conf_threshold_coef": 0.75}, + }, + "Loop": { + "SALAD": { + "image_size": [336, 336], + "batch_size": 32, + "similarity_threshold": 0.85, + "top_k": 5, + "use_nms": True, + "nms_threshold": 25, + }, + "SIM3_Optimizer": {"lang_version": "python", "max_iterations": 30, "lambda_init": "1e-6"}, + }, + } + + +def extract_frames(video: Path, output: Path) -> int: + import cv2 + + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + + output.mkdir() + capture = cv2.VideoCapture(str(video)) + try: + expected = int(capture.get(cv2.CAP_PROP_FRAME_COUNT)) + if not capture.isOpened() or expected <= 0: + raise EvaluationError(f"cannot decode video: {video}") + count = 0 + while True: + ok, frame = capture.read() + if not ok: + break + if not cv2.imwrite(str(output / f"frame_{count:06d}.png"), frame): + raise EvaluationError(f"cannot write extracted frame {count}") + count += 1 + if count != expected: + raise EvaluationError(f"decoded {count} frames but video reports {expected}") + return count + finally: + capture.release() + + +def convert_outputs(native: Path, output: Path, total_frames: int) -> None: + """Preserve full-frame poses and stride-four float16 geometry, including the final frame.""" + import numpy as np + + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + + c2w = np.loadtxt(native / "camera_poses.txt", dtype=np.float64).reshape(-1, 4, 4) + if len(c2w) != total_frames or not np.isfinite(c2w).all(): + raise EvaluationError("DA3 camera poses must be finite and cover the entire video") + if not np.allclose(c2w[:, 3], [0, 0, 0, 1], atol=1e-6): + raise EvaluationError("DA3 camera poses have invalid homogeneous rows") + frame_dir = native / "results_output" + files = sorted(frame_dir.glob("frame_*.npz"), key=lambda path: int(path.stem.split("_")[-1])) + if [path.name for path in files] != [f"frame_{index}.npz" for index in range(total_frames)]: + raise EvaluationError("DA3 must export depth/confidence for every video frame") + selected = np.arange(total_frames, dtype=np.int64)[::4] + if selected[-1] != total_frames - 1: + selected = np.concatenate([selected, np.asarray([total_frames - 1], dtype=np.int64)]) + depths, confidence, intrinsics = [], [], [] + for index in selected: + with np.load(files[int(index)], allow_pickle=False) as frame: + depth, conf = np.asarray(frame["depth"], dtype=np.float32), np.asarray(frame["conf"], dtype=np.float32) + intrinsic = np.asarray(frame["intrinsics"], dtype=np.float32) + if depth.ndim != 2 or conf.shape != depth.shape or intrinsic.shape != (3, 3): + raise EvaluationError("invalid DA3 depth/confidence/intrinsics shapes") + if ( + not all(np.isfinite(array).all() for array in (depth, conf, intrinsic)) + or np.any(depth < 0) + or not np.any(depth > 0) + or np.any(conf < 0) + ): + raise EvaluationError("invalid DA3 depth/confidence values") + depths.append(depth.astype(np.float16)) + confidence.append(conf.astype(np.float16)) + intrinsics.append(intrinsic) + output.mkdir() + np.savez( + output / "results.npz", + extrinsics=np.linalg.inv(c2w)[:, :3], + frame_indices=np.arange(total_frames, dtype=np.int64), + ) + np.savez_compressed( + output / "scene_geometry.npz", + schema_version=np.asarray([1], dtype=np.int64), + source_frame_count=np.asarray([total_frames], dtype=np.int64), + frame_indices=selected, + depths=np.stack(depths), + confidence=np.stack(confidence), + intrinsics=np.stack(intrinsics), + c2w=c2w[selected], + source=np.asarray(["da3_streaming_depth_conf_v1"]), + sample_stride=np.asarray([4], dtype=np.int64), + ) + + +def ensure_artifacts(sample, context, *, weight_metric=None, require_scene=True) -> None: + """Reuse complete artifacts; generate into an empty standard case/da3 directory.""" + import yaml + + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + + paths = sample["artifacts"] + pose = Path(paths["da3"]) + scene = Path(paths.get("da3_scene") or pose.parent / "scene_geometry.npz") + if pose.is_file() and (not require_scene or scene.is_file()): + return + # Keep outputs beside the case's video link, not beside its external target. + video = Path(sample["video"]).absolute() + destination = video.parent / "da3" + if pose != destination / "results.npz" or scene != destination / "scene_geometry.npz": + raise EvaluationError("missing manifest DA3 artifacts; automatic generation uses the standard case/da3 paths") + if destination.exists() and any(destination.iterdir()): + raise EvaluationError( + f"incomplete DA3 artifacts in {destination}; move the existing directory aside before regenerating" + ) + if not context.device.startswith("cuda"): + raise EvaluationError("DA3 generation requires CUDA; CPU scoring can reuse existing DA3 artifacts") + config = context.config.get("da3") or {} + if not config.get("source") or not config.get("dinov2_source"): + raise EvaluationError("Configure da3.source and da3.dinov2_source for DA3-Streaming, or supply DA3 artifacts") + source = Path(config["source"]).expanduser().resolve() + dinov2 = Path(config["dinov2_source"]).expanduser().resolve() + runner = source / "da3_streaming/da3_streaming.py" + if not runner.is_file() or not (dinov2 / "hubconf.py").is_file(): + raise EvaluationError( + "Configured DA3-Streaming source must contain da3_streaming/da3_streaming.py; dinov2_source must contain hubconf.py" + ) + weights = ( + Path(config["weights"]).expanduser().resolve() + if config.get("weights") + else ( + context.cache_dir / "metrics" / weight_metric if weight_metric is not None else context.cache_dir / "models" + ) + ) + settings = streaming_config(weights) + for path in settings["Weights"].values(): + if not Path(path).is_file(): + raise EvaluationError(f"missing DA3 weight: {path}; stage the official checkpoint or configure da3.weights") + command = [str(Path(config.get("python") or sys.executable).expanduser())] + index = int(context.device.split(":")[1]) if ":" in context.device else 0 + env = dict(os.environ) + visible = env.get("CUDA_VISIBLE_DEVICES") + env["CUDA_VISIBLE_DEVICES"] = visible.split(",")[index].strip() if visible else str(index) + env.update( + HF_HUB_OFFLINE="1", + TRANSFORMERS_OFFLINE="1", + PYTHONDONTWRITEBYTECODE="1", + PYTHONUNBUFFERED="1", + DINOV2_REPO_DIR=str(dinov2), + ) + env["PYTHONPATH"] = os.pathsep.join(filter(None, (str(source / "src"), env.get("PYTHONPATH")))) + with TemporaryDirectory(prefix=".da3-", dir=video.parent) as directory: + work = Path(directory) + count = extract_frames(video, work / "frames") + config_path = work / "config.yaml" + config_path.write_text(yaml.safe_dump(settings, sort_keys=False), encoding="utf-8") + # Keep verbose native output off the terminal and out of memory. + with TemporaryFile(dir=work) as log: + try: + subprocess.run( + [ + *command, + str(runner), + "--image_dir", + str(work / "frames"), + "--config", + str(config_path), + "--output_dir", + str(work / "native"), + ], + cwd=runner.parent, + env=env, + check=True, + stdout=log, + stderr=subprocess.STDOUT, + ) + except subprocess.CalledProcessError as error: + end = log.seek(0, os.SEEK_END) + log.seek(max(0, end - 8192)) + tail = "\n".join(log.read().decode("utf-8", errors="replace").splitlines()[-20:]) + raise EvaluationError(f"DA3 exited with code {error.returncode}\n{tail}".rstrip()) from error + convert_outputs(work / "native", work / "artifacts", count) + # Publish both files together. rename refuses to overwrite a non-empty directory. + (work / "artifacts").rename(destination) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/vipe.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/vipe.py new file mode 100644 index 000000000..7c08c69d6 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/preparation/vipe.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Prepare full-frame ViPE poses for video.mp4 cases using the default pose-only pipeline.""" + +import json +import os +import subprocess +import sys +from pathlib import Path +from tempfile import TemporaryDirectory, TemporaryFile + + +def vipe_command(python, source, video, output): + return [ + *python, + str(source / "run.py"), + "pipeline=default", + "streams=raw_mp4_stream", + "streams.base_path=" + json.dumps(str(video), ensure_ascii=False), + "streams.frame_start=0", + "streams.frame_end=-1", + "streams.frame_skip=1", + "streams.cached=true", + "pipeline.post.depth_align_model=null", + "pipeline.output.path=" + json.dumps(str(output), ensure_ascii=False), + "pipeline.output.save_artifacts=true", + "pipeline.output.save_slam_map=false", + "pipeline.output.save_viz=false", + ] + + +def convert_output(native: Path, output: Path, total_frames: int): + import numpy as np + + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.metrics.camera.schemas import ( + PoseTrajectory, + ) + + with np.load(native, allow_pickle=False) as data: + poses = np.asarray(data["data"], dtype=np.float64) + indices = np.asarray(data["inds"]) + if poses.shape != (total_frames, 4, 4) or not np.array_equal(indices, np.arange(total_frames)): + raise EvaluationError("ViPE must register every frame in exact 0..F-1 order; sparse output is invalid") + PoseTrajectory(poses, source="vipe") + np.savez(output, c2w=poses, frame_indices=indices.astype(np.int64)) + + +def ensure_artifacts(sample, context): + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import EvaluationError + from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.video import VideoReader + + output = Path(sample["artifacts"]["vipe"]) + if output.is_file(): + return + video = Path(sample["video"]).absolute() + if output != video.parent / "vipe/results.npz": + raise EvaluationError("missing manifest ViPE artifact; automatic generation uses case/vipe/results.npz") + if output.parent.exists() and any(output.parent.iterdir()): + raise EvaluationError(f"incomplete ViPE artifacts in {output.parent}; move them aside before regenerating") + if not context.device.startswith("cuda"): + raise EvaluationError("ViPE generation requires CUDA; CPU scoring can reuse existing poses") + config = context.config.get("vipe") or {} + if not config.get("source"): + raise EvaluationError("Configure vipe.source for the official ViPE default pipeline, or supply ViPE poses") + source = Path(config["source"]).expanduser().resolve() + if not (source / "run.py").is_file(): + raise EvaluationError("Configured ViPE model runtime must contain run.py") + weights = ( + Path(config.get("weights") or context.cache_dir / "metrics/camera_translation/vipe").expanduser().resolve() + ) + python = [str(Path(config.get("python") or sys.executable).expanduser())] + count = VideoReader().probe(str(video)).total_frames + env = dict(os.environ) + index = int(context.device.split(":")[1]) if ":" in context.device else 0 + visible = env.get("CUDA_VISIBLE_DEVICES") + env["CUDA_VISIBLE_DEVICES"] = visible.split(",")[index].strip() if visible else str(index) + env.update( + TORCH_HOME=str(weights / "torch"), + HF_HOME=str(weights / "huggingface"), + HF_HUB_CACHE=str(weights / "huggingface/hub"), + TRANSFORMERS_CACHE=str(weights / "huggingface/hub"), + HF_HUB_OFFLINE="1", + TRANSFORMERS_OFFLINE="1", + PYTHONDONTWRITEBYTECODE="1", + ) + with TemporaryDirectory(prefix=".vipe-", dir=video.parent) as directory: + work = Path(directory) + native = work / "native" + with TemporaryFile(dir=work) as log: + try: + subprocess.run( + vipe_command(python, source, video, native), + cwd=source, + env=env, + stdout=log, + stderr=subprocess.STDOUT, + check=True, + ) + except subprocess.CalledProcessError as error: + end = log.seek(0, os.SEEK_END) + log.seek(max(0, end - 8192)) + tail = "\n".join(log.read().decode("utf-8", errors="replace").splitlines()[-20:]) + raise EvaluationError(f"ViPE failed (exit {error.returncode}):\n{tail}") from error + converted = work / "results.npz" + convert_output(native / "pose" / f"{video.stem}.npz", converted, count) + output.parent.mkdir(exist_ok=True) + output.hardlink_to(converted) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/results.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/results.py new file mode 100644 index 000000000..b0b56c9fa --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/results.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Unit, item, and summary results shared by all metrics.""" + +import json +import shlex +import sys +from dataclasses import asdict, dataclass, field +from enum import Enum +from pathlib import Path +from typing import Any + + +class ResultStatus(str, Enum): + OK = "ok" + MODEL_FAILURE = "model_failure" + DATA_INVALID = "data_invalid" + EVAL_FAILED = "eval_failed" + NOT_APPLICABLE = "not_applicable" + + @property + def is_scored(self) -> bool: + return self in {ResultStatus.OK, ResultStatus.MODEL_FAILURE} + + +@dataclass +class UnitResult: + unit_id: str + unit_type: str = "item" + scores: dict[str, float] = field(default_factory=dict) + status: ResultStatus = ResultStatus.OK + error: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class ItemResult: + item_id: str + model_id: str | None = None + sample_id: str | None = None + condition_id: str | None = None + category: str | None = None + scores: dict[str, float] = field(default_factory=dict) + num_units: int = 0 + status: ResultStatus = ResultStatus.OK + error: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class SummaryResult: + num_items: int + status_counts: dict[str, int] = field(default_factory=dict) + scores: dict[str, float] = field(default_factory=dict) + metadata: dict[str, Any] = field(default_factory=dict) + group_by: str = "model_id" + results: dict[str, dict[str, Any]] = field(default_factory=dict) + groups: list[dict[str, Any]] = field(default_factory=list) + + +class ResultWriter: + def __init__(self, output: str | Path) -> None: + self.output = Path(output) + self.output.mkdir(parents=True, exist_ok=False) + + @staticmethod + def _json(value: Any) -> str: + return json.dumps(value, ensure_ascii=False, allow_nan=False) + + def write_command(self, command: list[str] | None) -> None: + invocation = shlex.join(command) if command is not None else "Python API call (no CLI command supplied)" + (self.output / "command.txt").write_text( + f"Working directory: {Path.cwd()}\nPython: {sys.executable}\n\n{invocation}\n", + encoding="utf-8", + ) + + def write_metric(self, name: str, items: list[ItemResult], summary: SummaryResult) -> None: + directory = self.output / "metrics" / name + directory.mkdir(parents=True) + with (directory / "cases.jsonl").open("w", encoding="utf-8") as handle: + for item in items: + handle.write(self._json(asdict(item)) + "\n") + (directory / "summary.json").write_text( + json.dumps(summary_payload(summary), ensure_ascii=False, allow_nan=False, indent=2) + "\n", + encoding="utf-8", + ) + + +def summary_payload(summary: SummaryResult) -> dict: + """Keep successful summaries concise, but retain non-OK coverage details.""" + + def compact_group(value): + value = dict(value) + metadata = dict(value.pop("metadata", {})) + if "n_total" in metadata: + value["num_cases"] = metadata.pop("n_total") + scored = metadata.pop("n_scored", None) + if scored is not None and scored != value.get("num_cases"): + value["num_scored"] = scored + counts = {status.value: metadata.pop(f"n_{status.value}", 0) for status in ResultStatus} + if any(count for status, count in counts.items() if status != "ok"): + value["status_counts"] = {status: count for status, count in counts.items() if count} + if metadata: + value["metadata"] = metadata + return value + + payload = {"num_cases": summary.num_items} + if any(count for status, count in summary.status_counts.items() if status != "ok"): + payload["status_counts"] = {status: count for status, count in summary.status_counts.items() if count} + if summary.scores: + payload["scores"] = summary.scores + if summary.metadata: + payload["metadata"] = summary.metadata + if summary.results: + payload["results"] = {key: compact_group(value) for key, value in summary.results.items()} + if summary.groups: + payload["groups"] = [compact_group(value) for value in summary.groups] + return payload diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/video.py b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/video.py new file mode 100644 index 000000000..f4a571a89 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/worldmembench/runtime/video.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 ZJU-REAL/WorldMemBench contributors +# SPDX-License-Identifier: MIT + +"""Original-frame-index video probing and RGB decoding.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from worldfoundry.evaluation.tasks.execution.runners.worldmembench.runtime.errors import VideoDecodeError + + +@dataclass(frozen=True) +class VideoInfo: + total_frames: int + fps: float + width: int + height: int + + +@dataclass(frozen=True) +class DecodedFrame: + frame_idx: int + rgb: Any + + +def fixed_stride_indices(total_frames: int, stride: int) -> list[int]: + if total_frames <= 0: + raise ValueError("total_frames must be positive") + if stride <= 0: + raise ValueError("stride must be positive") + return list(range(0, total_frames, stride)) + + +class VideoReader: + def probe(self, video_path: str) -> VideoInfo: + import cv2 + + path = str(Path(video_path).expanduser().resolve()) + capture = cv2.VideoCapture(path) + try: + total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT)) + fps = float(capture.get(cv2.CAP_PROP_FPS)) + width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) + if not capture.isOpened() or total_frames <= 0 or width <= 0 or height <= 0: + raise VideoDecodeError(f"cannot probe video: {path}") + return VideoInfo( + total_frames=total_frames, + fps=fps, + width=width, + height=height, + ) + finally: + capture.release() + + def read_indices( + self, + video_path: str, + indices: list[int] | tuple[int, ...], + ) -> list[DecodedFrame]: + import cv2 + + path = str(Path(video_path).expanduser().resolve()) + info = self.probe(path) + for index in indices: + if index < 0 or index >= info.total_frames: + raise VideoDecodeError(f"requested frame {index} outside [0, {info.total_frames - 1}] for {path}") + capture = cv2.VideoCapture(path) + frames: list[DecodedFrame] = [] + try: + previous_index: int | None = None + for index in indices: + # Consecutive all-frame consumers (notably Human Entity) should + # decode sequentially. Seeking before every frame is orders of + # magnitude slower for inter-frame codecs. + if previous_index is None or index != previous_index + 1: + capture.set(cv2.CAP_PROP_POS_FRAMES, int(index)) + ok, bgr = capture.read() + if not ok or bgr is None: + raise VideoDecodeError(f"cannot decode frame {index} from {path}") + frames.append( + DecodedFrame( + frame_idx=int(index), + rgb=cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB), + ) + ) + previous_index = int(index) + finally: + capture.release() + return frames diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldscore/runtime/worldscore/worldscore/benchmark/metrics/base_metrics.py b/worldfoundry/evaluation/tasks/execution/runners/worldscore/runtime/worldscore/worldscore/benchmark/metrics/base_metrics.py index 2e7249ea3..805663914 100755 --- a/worldfoundry/evaluation/tasks/execution/runners/worldscore/runtime/worldscore/worldscore/benchmark/metrics/base_metrics.py +++ b/worldfoundry/evaluation/tasks/execution/runners/worldscore/runtime/worldscore/worldscore/benchmark/metrics/base_metrics.py @@ -1,104 +1,23 @@ -from abc import ABC, abstractmethod -from typing import List, Union +"""WorldScore image metrics backed by the shared IQA implementation.""" -import pyiqa -import pyiqa.models -import pyiqa.models.inference_model -import torch -from PIL import Image -from pyiqa.models.inference_model import InferenceModel -from torchvision import transforms -from torchvision.transforms import ToTensor -from typing import Tuple -import numpy as np +from worldfoundry.evaluation.tasks.execution.runners.worldscore.runtime.worldscore.worldscore.common.gpu_utils import ( + get_torch_device, +) +from worldfoundry.evaluation.tasks.execution.runners.worldscore.runtime.worldscore.worldscore.common.images_utils import ( + open_image, +) +from worldfoundry.evaluation.tasks.metrics._shared.pyiqa import ImageQualityMetric, PyIQAMetric -from worldfoundry.evaluation.tasks.execution.runners.worldscore.runtime.worldscore.worldscore.common.gpu_utils import get_torch_device -from worldfoundry.evaluation.tasks.execution.runners.worldscore.runtime.worldscore.worldscore.common.images_utils import open_image - -class BaseMetric(ABC): - """BaseMetric Class.""" +class BaseMetric(ImageQualityMetric): + """Retain WorldScore's device and local/URL image loading policy.""" def __init__(self) -> None: - self._metric = None - self._device = get_torch_device() - - def _process_image( - self, - rendered_images: List[Union[str, Image.Image]], - ) -> float: - preprocessing = transforms.Compose( - [ - transforms.Resize((512, 512)), - transforms.ToTensor(), - ] - ) - - rendered_images_: List[torch.Tensor] = [] - for image in rendered_images: - # Handle the rendered image input - if isinstance(image, str): - image = preprocessing(open_image(image)) - else: - image = preprocessing(image) - rendered_images_.append(image) - - img: torch.Tensor = torch.stack(rendered_images_).to(self._device) - - return img - - def _process_images( - self, - rendered_images: List[Union[str, Image.Image]], - reference_image: Union[str, Image.Image], - ) -> float: - preprocessing = transforms.Compose( - [ - transforms.Resize((512, 512)), - transforms.ToTensor(), - ] - ) + super().__init__(device=get_torch_device(), image_loader=open_image) - # Handle the reference image input - if isinstance(reference_image, str): - reference_image = preprocessing(open_image(reference_image)) - rendered_images_: List[torch.Tensor] = [] - reference_images_: List[torch.Tensor] = [] - for image in rendered_images: - # Handle the rendered image input - if isinstance(image, str): - image = preprocessing(open_image(image)) - else: - image = preprocessing(image) - rendered_images_.append(image) +class IQAPytorchMetric(PyIQAMetric): + """Apply WorldScore's input policy to the selected PyIQA metric.""" - reference_images_.append(reference_image) - - img1: torch.Tensor = torch.stack(rendered_images_).to(self._device) - img2: torch.Tensor = torch.stack(reference_images_).to(self._device) - - return img1, img2 - - def _process_np_to_tensor( - self, - rendered_image: np.ndarray, - reference_image: np.ndarray, - ) -> Tuple[torch.Tensor, torch.Tensor]: - img1 = ToTensor()(rendered_image).unsqueeze(0).to(self._device) - img2 = ToTensor()(reference_image).unsqueeze(0).to(self._device) - return img1, img2 - - @abstractmethod - def _compute_scores(self, *args): - pass - - -class IQAPytorchMetric(BaseMetric): def __init__(self, metric_name: str) -> None: - super().__init__() - self._metric = self._create_metric(metric_name).to(self._device) - - def _create_metric(self, metric: str) -> InferenceModel: - metric = pyiqa.create_metric(metric) - return metric + super().__init__(metric_name, device=get_torch_device(), image_loader=open_image) diff --git a/worldfoundry/evaluation/tasks/execution/runners/worldscore/runtime/worldscore/worldscore/benchmark/metrics/metric_impls/clip_mlp_aesthetic_metrics.py b/worldfoundry/evaluation/tasks/execution/runners/worldscore/runtime/worldscore/worldscore/benchmark/metrics/metric_impls/clip_mlp_aesthetic_metrics.py index 82e1314da..a963a4797 100755 --- a/worldfoundry/evaluation/tasks/execution/runners/worldscore/runtime/worldscore/worldscore/benchmark/metrics/metric_impls/clip_mlp_aesthetic_metrics.py +++ b/worldfoundry/evaluation/tasks/execution/runners/worldscore/runtime/worldscore/worldscore/benchmark/metrics/metric_impls/clip_mlp_aesthetic_metrics.py @@ -1,110 +1,45 @@ -from typing import Union, List -import pytorch_lightning as pl -import torch.nn as nn -import torch +"""WorldScore's CLIP/LAION aesthetic metric using the shared MLP head.""" -from worldfoundry.base_models.perception_core.general_perception import openai_clip as clip -from PIL import Image +from collections.abc import Sequence -from worldfoundry.evaluation.tasks.execution.runners.worldscore.runtime.worldscore.worldscore.benchmark.metrics.base_metrics import BaseMetric - -# if you changed the MLP architecture during training, change it also here: -class MLP(pl.LightningModule): - def __init__(self, input_size, xcol='emb', ycol='avg_rating'): - super().__init__() - self.input_size = input_size - self.xcol = xcol - self.ycol = ycol - self.layers = nn.Sequential( - nn.Linear(self.input_size, 1024), - #nn.ReLU(), - nn.Dropout(0.2), - nn.Linear(1024, 128), - #nn.ReLU(), - nn.Dropout(0.2), - nn.Linear(128, 64), - #nn.ReLU(), - nn.Dropout(0.1), - - nn.Linear(64, 16), - #nn.ReLU(), - - nn.Linear(16, 1) - ) - - def forward(self, x): - return self.layers(x) +import torch +from PIL import Image - def training_step(self, batch, batch_idx): - x = batch[self.xcol] - y = batch[self.ycol].reshape(-1, 1) - x_hat = self.layers(x) - loss = F.mse_loss(x_hat, y) - return loss - - def validation_step(self, batch, batch_idx): - x = batch[self.xcol] - y = batch[self.ycol].reshape(-1, 1) - x_hat = self.layers(x) - loss = F.mse_loss(x_hat, y) - return loss +from worldfoundry.base_models.perception_core.general_perception import openai_clip as clip +from worldfoundry.evaluation.tasks.execution.runners.worldscore.runtime.worldscore.worldscore.benchmark.metrics.base_metrics import ( + BaseMetric, +) +from worldfoundry.evaluation.tasks.metrics._shared.aesthetic import load_laion_aesthetic_mlp_head - def configure_optimizers(self): - optimizer = torch.optim.Adam(self.parameters(), lr=1e-3) - return optimizer -def normalized(a, axis=-1, order=2): - import numpy as np # pylint: disable=import-outside-toplevel +def normalized(a, axis: int = -1, order: int = 2): + """Retain the official NumPy feature normalization, including zero rows.""" + import numpy as np l2 = np.atleast_1d(np.linalg.norm(a, order, axis)) l2[l2 == 0] = 1 return a / np.expand_dims(l2, axis) + class CLIPMLPAestheticScoreMetric(BaseMetric): - """ - - We use the CLIP+MLP Aesthetic Score Predictor implementation: - https://github.com/christophschuhmann/improved-aesthetic-predictor - - Range: [0, 10] higher the better - """ - + """Score CLIP ViT-L/14 features with the improved-aesthetic-predictor MLP.""" + def __init__(self) -> None: super().__init__() - model = MLP(768) # CLIP embedding dim is 768 for CLIP ViT L 14 - s = torch.load("worldscore/benchmark/metrics/checkpoints/sac+logos+ava1-l14-linearMSE.pth", map_location=torch.device('cpu')) # load the model you trained previously or the model available in this repo + self._model = load_laion_aesthetic_mlp_head( + "worldscore/benchmark/metrics/checkpoints/sac+logos+ava1-l14-linearMSE.pth" + ).to(self._device) + self._model2, self._preprocess = clip.load("ViT-L/14", device=self._device) - model.load_state_dict(s) - - model.to(self._device) - model.eval() - - model2, preprocess = clip.load("ViT-L/14", device=self._device) #RN50x64 - - self._model = model - self._model2 = model2 - self._preprocess = preprocess - - def _compute_scores( - self, - rendered_images: List[Union[str, Image.Image]], - ) -> float: - try: - from torchmetrics.multimodal.clip_score import CLIPScore - except ModuleNotFoundError as e: - handle_module_not_found_error(e, ["worldscore"]) - + def _compute_scores(self, rendered_images: Sequence[str | Image.Image]) -> float: scores = [] - - for i, image in enumerate(rendered_images): + for image in rendered_images: if isinstance(image, str): image = Image.open(image) - image = self._preprocess(image).unsqueeze(0).to(self._device) + tensor = self._preprocess(image).unsqueeze(0).to(self._device) with torch.no_grad(): - image_features = self._model2.encode_image(image) - im_emb_arr = normalized(image_features.cpu().detach().numpy()) - prediction = self._model(torch.from_numpy(im_emb_arr).type(torch.FloatTensor).to(self._device)) + image_features = self._model2.encode_image(tensor) + features = normalized(image_features.cpu().detach().numpy()) + prediction = self._model(torch.from_numpy(features).type(torch.FloatTensor).to(self._device)) scores.append(prediction.detach()[0].item()) - - score = sum(scores) / len(scores) - return score + return sum(scores) / len(scores) diff --git a/worldfoundry/evaluation/tasks/metrics/_shared/pyiqa.py b/worldfoundry/evaluation/tasks/metrics/_shared/pyiqa.py new file mode 100644 index 000000000..1edea4aba --- /dev/null +++ b/worldfoundry/evaluation/tasks/metrics/_shared/pyiqa.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 WorldFoundry contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Image batching and PyIQA inference shared by WorldScore and 4DWorldBench. + +Adapted from the WorldScore and 4DWorldBench image metric bases. +The 512-square resize is part of these benchmarks' protocol, not a default +preprocessing policy for other IQA metrics. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable, Sequence +from importlib import import_module +from typing import TYPE_CHECKING, Any + +from worldfoundry.evaluation.tasks.metrics._shared.images import load_rgb_image + +if TYPE_CHECKING: + import numpy as np + import torch + from PIL import Image + + +class ImageQualityMetric(ABC): + """Batch images using the WorldScore/4DWorldBench 512-square transform.""" + + def __init__(self, *, device: Any, image_loader: Callable[[str], Image.Image] = load_rgb_image) -> None: + self._device = device + self._image_loader = image_loader + self._metric: Any = None + + def _process_image(self, rendered_images: Sequence[str | Image.Image]) -> torch.Tensor: + import torch + from torchvision import transforms + + preprocessing = transforms.Compose([transforms.Resize((512, 512)), transforms.ToTensor()]) + images = [ + preprocessing(self._image_loader(image) if isinstance(image, str) else image) for image in rendered_images + ] + return torch.stack(images).to(self._device) + + def _process_images( + self, + rendered_images: Sequence[str | Image.Image], + reference_image: str | Image.Image, + ) -> tuple[torch.Tensor, torch.Tensor]: + reference = self._process_image([reference_image]) + rendered = self._process_image(rendered_images) + return rendered, reference.repeat(len(rendered_images), 1, 1, 1) + + def _process_np_to_tensor( + self, rendered_image: np.ndarray, reference_image: np.ndarray + ) -> tuple[torch.Tensor, torch.Tensor]: + from torchvision.transforms import ToTensor + + return ( + ToTensor()(rendered_image).unsqueeze(0).to(self._device), + ToTensor()(reference_image).unsqueeze(0).to(self._device), + ) + + @abstractmethod + def _compute_scores(self, *args: Any) -> Any: + """Compute the benchmark's metric-specific score.""" + + +class PyIQAMetric(ImageQualityMetric): + """Load the explicitly selected PyIQA variant without changing its defaults.""" + + def __init__( + self, metric_name: str, *, device: Any, image_loader: Callable[[str], Image.Image] = load_rgb_image + ) -> None: + super().__init__(device=device, image_loader=image_loader) + self._metric = self._create_metric(metric_name).to(self._device) + + def _create_metric(self, metric: str) -> Any: + return import_module("pyiqa").create_metric(metric)