tests/dynamic/test_dl3dv_benchmark.py
"""Contract tests for the metadata-only and causal DL3DV-140 adapter."""
from __future__ import annotations
import json
from pathlib import Path
import cv2
import numpy as np
import pytest
import torch
from PIL import Image
from mcss.data.dl3dv_benchmark import (
DL3DV_SCENE_COUNT,
DL3DV_SPLIT_SHA256,
Dl3dvBenchmarkMetadata,
Dl3dvFrame,
Dl3dvOnlineObservationSource,
Dl3dvProtocol,
Dl3dvProtocolError,
Dl3dvSceneMetadata,
LongLRMSplitRecord,
compute_warmup_gauge,
load_benchmark_metadata,
load_long_lrm_split,
official_pose_transform,
prepare_scene_metadata,
preprocess_dl3dv_rgb,
target_frame_indices,
validate_complete_coverage,
)
from mcss.dynamic.types import SealedScene, hash_scene_state
from mcss.types import SceneState
PROJECT = Path(__file__).resolve().parents[2]
SPLIT_PATH = PROJECT / "tests/fixtures/long_lrm/dl3dv_fold_8_kmeans_input_idx.json"
def _frame(
scene_id: str,
frame_id: int,
*,
source_size: tuple[int, int] = (4, 6),
x: float | None = None,
rgb_name: str | None = None,
) -> Dl3dvFrame:
pose = np.eye(4, dtype=np.float64)
pose[0, 3] = float(frame_id if x is None else x)
return Dl3dvFrame(
scene_id=scene_id,
frame_id=frame_id,
rgb_relpath=rgb_name or f"images_4/frame_{frame_id:04d}.png",
c2w=tuple(tuple(float(value) for value in row) for row in pose),
raw_fxfycxcy=(5.0, 5.0, 2.5, 1.5),
distortion=(0.0, 0.0, 0.0, 0.0),
source_image_size=source_size,
)
def _scene(
root: Path,
*,
scene_id: str = "scene",
input_indices: tuple[int, ...] = (
1,
2,
3,
4,
5,
6,
7,
9,
10,
11,
12,
13,
14,
15,
17,
18,
),
query_indices: tuple[int, ...] = (0, 8, 16),
source_size: tuple[int, int] = (4, 6),
target_size: tuple[int, int] = (2, 3),
mode: str = "full",
input_count: int = 16,
frame_count: int = 20,
pose_values: tuple[float, ...] | None = None,
) -> Dl3dvSceneMetadata:
values = pose_values or tuple(float(index) for index in range(frame_count))
frames = tuple(
_frame(scene_id, index, source_size=source_size, x=values[index])
for index in range(frame_count)
)
return Dl3dvSceneMetadata(
scene_id=scene_id,
scene_root=root,
frames=frames,
input_indices=input_indices,
query_indices=query_indices,
protocol=Dl3dvProtocol(mode, input_count),
target_size=target_size,
)
def _write_image(path: Path, value: int, size: tuple[int, int]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
Image.new("RGB", (size[1], size[0]), (value, value, value)).save(path)
def _write_transforms(
root: Path,
*,
scene_id: str,
frame_count: int,
input_indices: tuple[int, ...],
source_size: tuple[int, int] = (4, 6),
file_prefix: str = "images_4",
) -> Path:
scene_root = root / scene_id
transforms_dir = scene_root / "nerfstudio"
transforms_dir.mkdir(parents=True)
height, width = source_size
transforms = {
"w": width,
"h": height,
"fl_x": 5.0,
"fl_y": 5.0,
"cx": 2.5,
"cy": 1.5,
"k1": 0.0,
"k2": 0.0,
"p1": 0.0,
"p2": 0.0,
"frames": [],
}
for index in range(frame_count):
pose = np.eye(4, dtype=np.float64)
pose[0, 3] = float(index)
transforms["frames"].append(
{
"file_path": f"{file_prefix}/frame_{index:04d}.png",
"transform_matrix": pose.tolist(),
}
)
transforms_path = transforms_dir / "transforms.json"
transforms_path.write_text(json.dumps(transforms), encoding="utf-8")
return transforms_path
def _dummy_sealed(scene: Dl3dvSceneMetadata) -> SealedScene:
state = SceneState(
torch.zeros((1, 1, 2, 2, 2)),
torch.full((1, 3, 2, 2, 2), 0.5),
torch.zeros((1, 1, 2, 2, 2)),
torch.tensor([[[-6.0, -4.0, 0.05], [6.0, 4.0, 12.05]]]),
)
return SealedScene(
scene.episode_id,
scene.scene_id,
scene.split_id,
scene.query_vault_id,
state,
hash_scene_state(state),
scene.arrival_indices,
"checkpoint",
"config",
"fast",
torch.eye(4),
)
def test_pinned_long_lrm_split_has_complete_140_scene_coverage() -> None:
records = load_long_lrm_split(SPLIT_PATH)
assert len(records) == DL3DV_SCENE_COUNT
assert records[0].input_indices(16)[:4] == (13, 36, 57, 76)
assert all(
len(record.input_indices(count)) == count
for record in records
for count in (16, 32, 64, 128)
)
def test_pinned_long_lrm_split_rejects_modified_bytes(tmp_path: Path) -> None:
modified = tmp_path / "split.json"
original = SPLIT_PATH.read_bytes()
modified.write_bytes(original.replace(b"scene_name", b"scene_namx", 1))
with pytest.raises(Dl3dvProtocolError, match="SHA256"):
load_long_lrm_split(modified)
def test_frame_rejects_non_rigid_camera_pose() -> None:
pose = np.eye(4, dtype=np.float64)
pose[0, 1] = 0.01
with pytest.raises(Dl3dvProtocolError, match="orthonormal"):
Dl3dvFrame(
scene_id="scene",
frame_id=0,
rgb_relpath="images_4/frame_0000.png",
c2w=tuple(tuple(float(value) for value in row) for row in pose),
raw_fxfycxcy=(5.0, 5.0, 2.5, 1.5),
distortion=(0.0, 0.0, 0.0, 0.0),
source_image_size=(4, 6),
)
def test_metadata_round_trip_rebinds_split_and_scene_protocol(tmp_path: Path) -> None:
records = load_long_lrm_split(SPLIT_PATH)
protocol = Dl3dvProtocol("full", 16)
raw_root = tmp_path / "raw"
scenes = []
for record in records:
inputs = record.input_indices(16)
frame_count = max(inputs) + 1
frames = tuple(
_frame(record.scene_name, index, source_size=(540, 960))
for index in range(frame_count)
)
scenes.append(
Dl3dvSceneMetadata(
record.scene_name,
raw_root / record.scene_name,
frames,
inputs,
target_frame_indices(frame_count),
protocol,
)
)
metadata = Dl3dvBenchmarkMetadata(
tuple(scenes), raw_root, SPLIT_PATH, DL3DV_SPLIT_SHA256, protocol
)
path = tmp_path / "metadata.json"
path.write_text(json.dumps(metadata.to_dict(), indent=2), encoding="utf-8")
loaded = load_benchmark_metadata(path)
assert loaded.protocol == protocol
assert loaded.coverage()["scene_count"] == DL3DV_SCENE_COUNT
payload = json.loads(path.read_text(encoding="utf-8"))
payload["scenes"][0]["input_indices"][0] += 1
path.write_text(json.dumps(payload), encoding="utf-8")
with pytest.raises(Dl3dvProtocolError, match="input_indices"):
load_benchmark_metadata(path)
def test_supported_modes_match_paper_conditions() -> None:
assert Dl3dvProtocol("full", 16).split_key == "fold_8_kmeans_16_input"
assert Dl3dvProtocol("ar", 32).name == "ar32"
with pytest.raises(Dl3dvProtocolError, match="supports input counts"):
Dl3dvProtocol("ar", 16)
def test_prepare_scene_reads_metadata_without_opening_any_image(tmp_path: Path) -> None:
scene_id = "scene"
input_indices = tuple(index for index in range(32) if index % 8 != 0)[:16]
_write_transforms(tmp_path, scene_id=scene_id, frame_count=32, input_indices=input_indices)
split = LongLRMSplitRecord(scene_id, {16: input_indices})
metadata = prepare_scene_metadata(
tmp_path / scene_id,
split,
Dl3dvProtocol("full", 16),
require_source_size=False,
)
assert metadata.raw_source_size == (4, 6)
assert metadata.query_indices == target_frame_indices(32)
assert not list((tmp_path / scene_id / "nerfstudio").glob("images_4/*"))
def test_prepare_scene_maps_original_resolution_transforms_to_images4(tmp_path: Path) -> None:
scene_id = "scene"
input_indices = tuple(index for index in range(32) if index % 8 != 0)[:16]
_write_transforms(
tmp_path,
scene_id=scene_id,
frame_count=32,
input_indices=input_indices,
source_size=(2160, 3840),
file_prefix="images",
)
split = LongLRMSplitRecord(scene_id, {16: input_indices})
metadata = prepare_scene_metadata(
tmp_path / scene_id,
split,
Dl3dvProtocol("full", 16),
require_source_size=True,
)
assert metadata.raw_source_size == (540, 960)
assert metadata.frames[0].rgb_relpath == "images_4/frame_0000.png"
assert metadata.frames[0].raw_fxfycxcy == pytest.approx((1.25, 1.25, 0.625, 0.375))
def test_prepare_scene_rejects_context_query_overlap() -> None:
frames = tuple(_frame("scene", index) for index in range(16))
with pytest.raises(Dl3dvProtocolError, match="context/query overlap"):
# A direct protocol check is kept independent of image preparation.
from mcss.data.dl3dv_benchmark import _validate_scene_indices
_validate_scene_indices(16, (0, 1, 2, 8), target_stride=8)
assert frames[0].frame_id == 0
def test_complete_coverage_rejects_duplicate_scene_ids(tmp_path: Path) -> None:
scenes = [
_scene(
tmp_path,
scene_id=f"scene-{index}",
frame_count=20,
input_indices=(
1,
2,
3,
4,
5,
6,
7,
9,
10,
11,
12,
13,
14,
15,
17,
18,
),
query_indices=(0, 8, 16),
)
for index in range(DL3DV_SCENE_COUNT)
]
scenes[-1] = scenes[-2]
with pytest.raises(Dl3dvProtocolError, match="duplicate scene IDs"):
validate_complete_coverage(scenes)
def test_official_pose_transform_matches_pinned_axis_operations() -> None:
raw = np.eye(4, dtype=np.float64)
raw[0, 1] = 2.0
raw[0, 3] = 3.0
expected = raw.copy()
expected[0:3, 1:3] *= -1.0
expected = expected[[1, 0, 2, 3], :]
expected[2, :] *= -1.0
np.testing.assert_allclose(official_pose_transform(raw), expected)
def test_official_rgb_preprocessing_and_intrinsics_transform(tmp_path: Path) -> None:
image_path = tmp_path / "frame.png"
source_size = (4, 6)
image = np.zeros((source_size[0], source_size[1], 3), dtype=np.uint8)
image[:, :, 0] = np.arange(source_size[1], dtype=np.uint8)[None, :]
image[:, :, 1] = 100
cv2.imwrite(str(image_path), cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
frame = _frame("scene", 0, source_size=source_size, rgb_name="images_4/frame.png")
rgb, intrinsic = preprocess_dl3dv_rgb(image_path, frame, target_size=(4, 6))
assert rgb.shape == (3, 4, 6)
assert rgb.dtype == torch.float32
assert float(rgb.min()) >= 0.0 and float(rgb.max()) <= 1.0
expected_raw = np.asarray([[5.0, 0.0, 2.5], [0.0, 5.0, 1.5], [0.0, 0.0, 1.0]])
expected_new, _ = cv2.getOptimalNewCameraMatrix(
expected_raw, np.zeros(4), (6, 4), 0, (6, 4)
)
np.testing.assert_allclose(intrinsic, expected_new, rtol=1e-6, atol=1e-6)
def test_lazy_source_keeps_query_poison_unread_until_after_seal(tmp_path: Path) -> None:
input_indices = tuple(index for index in range(24) if index % 8 != 0)[:16]
query_indices = target_frame_indices(24)
scene = _scene(
tmp_path,
input_indices=input_indices,
query_indices=query_indices,
frame_count=24,
source_size=(4, 6),
target_size=(2, 3),
)
for frame_id in input_indices:
_write_image(
tmp_path / "nerfstudio" / scene.frame(frame_id).rgb_relpath,
17,
(4, 6),
)
source = Dl3dvOnlineObservationSource(scene)
warmup = source.warmup()
while source.remaining_steps:
source.next_observation()
assert [observation.frame_id for observation in warmup] == list(scene.warmup_indices)
assert source.sealed is False
sealed = _dummy_sealed(scene)
queries = source.queries_for_sealed(sealed)
assert source.sealed is True
# Query files did not need to exist while context was consumed. They are
# created only after callback construction to prove the supplier is lazy.
query_frame = scene.frame(scene.query_indices[0])
query_path = tmp_path / "nerfstudio" / query_frame.rgb_relpath
_write_image(query_path, 239, (4, 6))
assert float(queries[0].rgb_supplier().mean()) > 0.8
with pytest.raises(ValueError, match="sealed"):
source.next_observation()
def test_warmup_gauge_is_causal_and_rejects_zero_baseline() -> None:
warmup = []
for value in (0.0, 1.0, 2.0, 3.0):
pose = np.eye(4)
pose[0, 3] = value
warmup.append(pose)
first = compute_warmup_gauge(warmup)
changed_future = np.eye(4)
changed_future[0, 3] = 1000.0
assert first.scale == pytest.approx(1.9428480059945201 / 3.0)
expected_second = np.eye(4)
expected_second[0, 3] = first.scale
np.testing.assert_allclose(first.transform(warmup[1]), expected_second)
assert first.transform(changed_future)[0, 3] == pytest.approx(1000.0 * first.scale)
with pytest.raises(Dl3dvProtocolError, match="too small"):
compute_warmup_gauge([np.eye(4)] * 4)
def test_source_gauge_and_prefix_cameras_ignore_later_and_query_poses(tmp_path: Path) -> None:
input_indices = tuple(index for index in range(24) if index % 8 != 0)[:16]
first_values = tuple(float(index) for index in range(24))
second_values = list(first_values)
second_values[0] = 500.0
second_values[5] = 100.0
second_values[6] = 200.0
second_values[8] = 800.0
scene_a = _scene(
tmp_path / "a",
input_indices=input_indices,
query_indices=target_frame_indices(24),
frame_count=24,
pose_values=first_values,
source_size=(4, 6),
target_size=(2, 3),
)
scene_b = _scene(
tmp_path / "b",
input_indices=input_indices,
query_indices=target_frame_indices(24),
frame_count=24,
pose_values=tuple(second_values),
source_size=(4, 6),
target_size=(2, 3),
)
source_a = Dl3dvOnlineObservationSource(scene_a)
source_b = Dl3dvOnlineObservationSource(scene_b)
assert source_a.gauge.scale == pytest.approx(source_b.gauge.scale)
for index in scene_a.warmup_indices:
np.testing.assert_allclose(
source_a.gauge.transform(scene_a.frame(index).c2w_array),
source_b.gauge.transform(scene_b.frame(index).c2w_array),
)