| """μ°¨λ κ°μ²΄ κ°μ°μμμ μμΈ ν©μ± (nusim.recon.train._compose_objects, recon.objects) β CPU.""" |
|
|
| import numpy as np |
| import pytest |
|
|
| torch = pytest.importorskip("torch") |
|
|
| from nusim.data.rig import matrix_to_quat, quat_to_matrix |
| from nusim.recon.train import _compose_objects |
|
|
|
|
| def _rand_q(rng, n): |
| q = rng.normal(size=(n, 4)) |
| return q / np.linalg.norm(q, axis=1, keepdims=True) |
|
|
|
|
| def test_compose_objects_moves_only_object_gaussians(): |
| rng = np.random.default_rng(0) |
| n, K = 50, 3 |
| means = rng.normal(size=(n, 3)).astype(np.float32) |
| quats = _rand_q(rng, n).astype(np.float32) |
| obj = np.array([-1] * 20 + [0] * 10 + [1] * 10 + [2] * 10, np.float32) |
| params = {"means": torch.tensor(means), "quats": torch.tensor(quats), |
| "opacities": torch.zeros(n), "obj": torch.tensor(obj)} |
| Qb = _rand_q(rng, K) |
| R = np.stack([quat_to_matrix(q) for q in Qb]).astype(np.float32) |
| T = rng.normal(size=(K, 3)).astype(np.float32) * 10 |
| present = np.array([1, 1, 0], np.float32) |
| m, q, op = _compose_objects(params, tuple(torch.tensor(a) for a in (R, T, Qb.astype(np.float32), present))) |
| m, q, op = m.numpy(), q.numpy(), op.numpy() |
| bg = obj < 0 |
| assert np.allclose(m[bg], means[bg]) and np.allclose(q[bg], quats[bg]) and np.allclose(op[bg], 0.5) |
| for i in np.nonzero(~bg)[0]: |
| k = int(obj[i]) |
| assert np.allclose(m[i], R[k] @ means[i] + T[k], atol=1e-5) |
| |
| assert np.allclose(quat_to_matrix(q[i]), R[k] @ quat_to_matrix(quats[i]), atol=1e-5) |
| assert op[i] == pytest.approx(0.5 * present[k]) |
|
|
|
|
| def test_compose_objects_without_objects_is_identity(): |
| params = {"means": torch.zeros(4, 3), "quats": torch.tensor([[1.0, 0, 0, 0]] * 4), "opacities": torch.zeros(4)} |
| m, q, op = _compose_objects(params, None) |
| assert torch.equal(m, params["means"]) and torch.allclose(op, torch.full((4,), 0.5)) |
|
|
|
|
| def test_matrix_quat_consistency_for_box_rotation(): |
| |
| q = np.array([0.9, 0.1, -0.2, 0.37]) |
| q /= np.linalg.norm(q) |
| assert np.allclose(quat_to_matrix(matrix_to_quat(quat_to_matrix(q))), quat_to_matrix(q)) |
|
|
