50 lines
1.9 KiB
Python
50 lines
1.9 KiB
Python
import json
|
|
|
|
import numpy as np
|
|
|
|
from iiwa_genesis.camera import D455, Frame, calibration, optical_pose, optical_rays, process_depth
|
|
from iiwa_genesis.config import ROOT, load_config
|
|
|
|
|
|
def camera_config():
|
|
return load_config(ROOT / "configs/scene.yaml")["cameras"]["overhead"]
|
|
|
|
|
|
def test_downward_pose():
|
|
cfg = camera_config()
|
|
pose = optical_pose(cfg["position"], cfg["lookat"], cfg["up"])
|
|
np.testing.assert_allclose(pose[:3, :3].T @ pose[:3, :3], np.eye(3))
|
|
np.testing.assert_allclose(pose[:3, 2], [0, 0, -1])
|
|
assert np.linalg.det(pose[:3, :3]) > 0.999
|
|
|
|
|
|
def test_plane_range_converts_to_constant_z():
|
|
cfg = camera_config()
|
|
distances = 2 * np.linalg.norm(optical_rays(cfg), axis=-1)
|
|
depth, valid = process_depth(distances, cfg, np.random.default_rng(1))
|
|
np.testing.assert_allclose(depth, 2)
|
|
assert valid.all()
|
|
distances[0, :3] = [np.inf, 0, 100]
|
|
depth, valid = process_depth(distances, cfg, np.random.default_rng(1))
|
|
assert not valid[0, :3].any()
|
|
assert not depth[0, :3].any()
|
|
|
|
|
|
def test_seeded_noise_and_serialization(tmp_path):
|
|
cfg = camera_config()
|
|
cfg["noise"]["enabled"] = True
|
|
distances = 2 * np.linalg.norm(optical_rays(cfg), axis=-1)
|
|
depth, valid = process_depth(distances, cfg, np.random.default_rng(7))
|
|
second, _ = process_depth(distances, cfg, np.random.default_rng(7))
|
|
np.testing.assert_array_equal(depth, second)
|
|
assert np.std(depth) > 0
|
|
camera = object.__new__(D455)
|
|
camera.name, camera.cfg = "overhead", cfg
|
|
frame = Frame(np.zeros((480, 640, 3), dtype=np.uint8), depth, valid,
|
|
{"intrinsics": calibration(cfg).tolist(), "world_from_optical": np.eye(4).tolist()})
|
|
camera.save(frame, tmp_path, 0, pointcloud=True)
|
|
data = np.load(tmp_path / "overhead_000000.npz")
|
|
assert data["depth"].shape == (480, 640)
|
|
assert data["xyz_optical"].shape == (480, 640, 3)
|
|
assert json.loads((tmp_path / "overhead_000000.json").read_text()) == frame.metadata
|