tests/test_rgbd_replica_data.py
"""Replica cohort preparation: camera model, depth semantics, geometry rule and a synthetic run."""
import importlib.util
import io
import json
import types
import zipfile
from pathlib import Path
import numpy as np
import pytest
import torch
from PIL import Image
from mcss.geometry import generate_rays
from mcss.types import Cameras
SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "prepare_rgbd_replica_data.py"
spec = importlib.util.spec_from_file_location("prepare_rgbd_replica_data", SCRIPT)
replica = importlib.util.module_from_spec(spec)
spec.loader.exec_module(replica)
SOURCE = np.array([[600.0, 0.0, 599.5], [0.0, 600.0, 339.5], [0.0, 0.0, 1.0]])
ROOM = (np.array([-4.0, -3.0, -4.0]), np.array([4.0, 1.5, 4.0]))
def test_target_camera_matches_the_hypersim_field_of_view():
target = replica.target_intrinsics(SOURCE)
assert target[0, 2] == pytest.approx(79.5) and target[1, 2] == pytest.approx(63.5)
assert target[0, 0] == pytest.approx(138.56, rel=5e-3) and target[1, 1] == pytest.approx(
147.80, rel=5e-3
)
assert replica.CROP["y0"] * 2 + replica.CROP["height"] == 680
assert replica.CROP["x0"] * 2 + replica.CROP["width"] == 1200
def look_at(position, target):
forward = (target - position) / np.linalg.norm(target - position)
right = np.cross(forward, np.array([0.0, -1.0, 0.0]))
right /= np.linalg.norm(right)
down = np.cross(forward, right)
pose = np.eye(4)
pose[:3, :3] = np.stack((right, down, forward), -1) # OpenCV: x right, y down, z forward
pose[:3, 3] = position
return pose
def render(pose, height=680, width=1200, intrinsics=SOURCE):
"""Planar z-depth and a textured RGB of the inside of an axis-aligned box room."""
y, x = np.indices((height, width))
rays = np.stack(
(
(x - intrinsics[0, 2]) / intrinsics[0, 0],
(y - intrinsics[1, 2]) / intrinsics[1, 1],
np.ones_like(x, float),
),
-1,
)
world = rays @ pose[:3, :3].T
origin = pose[:3, 3]
with np.errstate(divide="ignore", invalid="ignore"):
t = np.stack(
[(ROOM[j][i] - origin[i]) / world[..., i] for j in range(2) for i in range(3)], -1
)
t = np.where(t > 0, t, np.inf).min(-1)
points = origin + world * t[..., None]
z = t # rays have unit z, so the ray parameter is the planar depth
checker = (
np.floor(points[..., 0] * 2) + np.floor(points[..., 1] * 2) + np.floor(points[..., 2] * 2)
) % 2
rgb = np.stack((60 + 150 * checker, 90 + 100 * (1 - checker), 140 + 0 * checker), -1).astype(
np.uint8
)
return rgb, z
def test_ray_distance_matches_the_project_ray_convention():
target = replica.target_intrinsics(SOURCE)
camera = Cameras(torch.tensor(target), torch.eye(4, dtype=torch.float64), (128, 160))
_, directions = generate_rays(camera)
z = np.full((128, 160), 2.0)
expected = 2.0 / directions[..., 2].numpy()
assert np.allclose(replica.ray_distance(z, target), expected, rtol=1e-9)
def test_geometry_rule_accepts_opencv_poses_and_rejects_flipped_ones():
a = look_at(np.array([0.0, -1.0, 0.0]), np.array([2.0, -1.0, 3.0]))
b = look_at(np.array([0.3, -1.0, 0.2]), np.array([2.5, -1.0, 3.0]))
_, za = render(a)
_, zb = render(b)
error, landing = replica.reprojection_error(za, a, zb, b, SOURCE)
assert error < 1e-3 and landing > 0.3
flip = np.diag([1.0, -1.0, -1.0, 1.0])
bad, _ = replica.reprojection_error(za, a @ flip, zb, b @ flip, SOURCE)
assert bad is None or bad > 0.05
def trajectory_text(n=2000):
poses = []
for k in range(n):
angle = np.radians(0.3 * k) # 15 degrees every 50 frames
position = np.array([1.5 * np.cos(angle), -1.0, 1.5 * np.sin(angle)])
target = np.array([-1.5 * np.cos(angle), -1.0, -1.5 * np.sin(angle)]) # opposite side
poses.append(look_at(position, target))
return "\n".join(" ".join(f"{v:.9f}" for v in pose.ravel()) for pose in poses) + "\n", poses
def synthetic_archive(path, scenes):
params = {
"camera": {
"w": 1200,
"h": 680,
"fx": 600.0,
"fy": 600.0,
"cx": 599.5,
"cy": 339.5,
"scale": 6553.5,
}
}
text, poses = trajectory_text()
with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_STORED) as archive:
archive.writestr("Replica/cam_params.json", json.dumps(params))
for scene in scenes:
archive.writestr(f"Replica/{scene}/traj.txt", text)
for fid in range(0, 751, 50):
rgb, z = render(poses[fid])
buffer = io.BytesIO()
Image.fromarray(rgb).save(buffer, format="JPEG", quality=95)
archive.writestr(f"Replica/{scene}/results/frame{fid:06d}.jpg", buffer.getvalue())
buffer = io.BytesIO()
Image.fromarray(np.round(z * 6553.5).clip(0, 65535).astype(np.uint16)).save(
buffer, format="PNG"
)
archive.writestr(f"Replica/{scene}/results/depth{fid:06d}.png", buffer.getvalue())
return poses
def test_lock_then_prepare_on_a_synthetic_archive(tmp_path, monkeypatch):
scenes = ("office0",)
monkeypatch.setattr(replica, "SCENES", scenes)
monkeypatch.setattr(replica, "MIN_VALID_SCENES", 1)
bounds = tmp_path / "bounds_contract.json"
bounds.write_text(json.dumps({"near_m": 0.214, "far_m": 5.672, "source_sha256": "0" * 64}))
monkeypatch.setattr(replica, "BOUNDS", bounds)
poses = synthetic_archive(tmp_path / "Replica.zip", scenes)
reader = types.SimpleNamespace(
etag='"synthetic"', size=(tmp_path / "Replica.zip").stat().st_size, transferred=0
)
def opener(budget=None):
return reader, zipfile.ZipFile(tmp_path / "Replica.zip")
root = tmp_path / replica.EXPERIMENT
root.mkdir()
_, archive = opener()
with archive:
lock = replica.build_lock(archive, reader)
assert lock["frame_rule"]["spacing"] == 50 and lock["image_or_depth_member_read"] is False
(root / "candidate_lock.json").write_text(json.dumps(lock))
replica.prepare(root, opener)
integrity = json.loads((root / "preparation_integrity.json").read_text())
assert integrity["status"] == "PASS", json.loads((root / "scene_failures.json").read_text())
manifest = json.loads((root / "manifest_replica.json").read_text())
record = manifest["scenes"][0]
assert record["split"] == "REPLICA" and set(record["roles"]) >= {
"context_a",
"context_b",
"primary_query",
}
frame = record["frames"][0]
depth = np.load(frame["depth"])
assert depth.shape == (128, 160) and manifest["depth_semantics"] == "ray_distance_meters"
target = np.array(frame["intrinsics"])
camera = Cameras(torch.tensor(target), torch.tensor(np.array(frame["c2w"])), (128, 160))
_, z = render(poses[0], 128, 160, target)
_, directions = generate_rays(camera)
local = directions.numpy() @ np.array(frame["c2w"])[:3, :3]
expected = z / local[..., 2]
inner = (slice(8, -8), slice(8, -8))
assert np.median(np.abs(depth[inner] - expected[inner]) / expected[inner]) < 0.02
assert json.loads((root / "geometry_checks.json").read_text())["office0"]["pass"]