|
"""Test gradient implementations.""" |
|
|
|
import logging |
|
import unittest |
|
|
|
import torch |
|
from torch.func import jacfwd, vmap |
|
|
|
from siclib.geometry.camera import camera_models |
|
from siclib.geometry.gravity import Gravity |
|
from siclib.geometry.jacobians import J_up_projection |
|
from siclib.geometry.manifolds import SphericalManifold |
|
from siclib.geometry.perspective_fields import J_perspective_field, get_perspective_field |
|
from siclib.models.optimization.lm_optimizer import LMOptimizer |
|
from siclib.utils.conversions import deg2rad, fov2focal |
|
|
|
|
|
|
|
|
|
H, W = 320, 320 |
|
|
|
K1 = -0.1 |
|
|
|
|
|
CAMERA_MODEL = "simple_radial" |
|
|
|
|
|
Camera = camera_models[CAMERA_MODEL] |
|
|
|
|
|
torch.autograd.set_detect_anomaly(True) |
|
|
|
|
|
logger = logging.getLogger("geocalib.models.base_model") |
|
logger.setLevel("ERROR") |
|
|
|
|
|
def get_toy_rpf(roll=None, pitch=None, vfov=None) -> torch.Tensor: |
|
"""Return a random roll, pitch, focal length if not specified.""" |
|
if roll is None: |
|
roll = deg2rad((torch.rand(1) - 0.5) * 90) |
|
elif not isinstance(roll, torch.Tensor): |
|
roll = torch.tensor(deg2rad(roll)).unsqueeze(0) |
|
|
|
if pitch is None: |
|
pitch = deg2rad((torch.rand(1) - 0.5) * 90) |
|
elif not isinstance(pitch, torch.Tensor): |
|
pitch = torch.tensor(deg2rad(pitch)).unsqueeze(0) |
|
|
|
if vfov is None: |
|
vfov = deg2rad(5 + torch.rand(1) * 75) |
|
elif not isinstance(vfov, torch.Tensor): |
|
vfov = torch.tensor(deg2rad(vfov)).unsqueeze(0) |
|
|
|
return torch.stack([roll, pitch, fov2focal(vfov, H)], dim=-1).float() |
|
|
|
|
|
class TestJacobianFunctions(unittest.TestCase): |
|
"""Test the jacobian functions.""" |
|
|
|
eps = 5e-3 |
|
|
|
def validate(self, J: torch.Tensor, J_auto: torch.Tensor): |
|
"""Check if the jacobians are close and finite.""" |
|
self.assertTrue(torch.all(torch.isfinite(J)), "found nan in numerical") |
|
self.assertTrue(torch.all(torch.isfinite(J_auto)), "found nan in auto") |
|
|
|
text_j = f" > {self.eps}\nJ:\n{J[0, 0].numpy()}\nJ_auto:\n{J_auto[0, 0].numpy()}" |
|
max_diff = torch.max(torch.abs(J - J_auto)) |
|
text = f"Overall - max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J, J_auto, atol=self.eps), text) |
|
|
|
def test_spherical_plus(self): |
|
"""Test the spherical plus operator.""" |
|
rpf = get_toy_rpf() |
|
gravity = Gravity.from_rp(rpf[..., 0], rpf[..., 1]) |
|
J = gravity.J_update(spherical=True) |
|
|
|
|
|
delta = gravity.vec3d.new_zeros(gravity.vec3d.shape)[..., :-1] |
|
|
|
def spherical_plus(delta: torch.Tensor) -> torch.Tensor: |
|
"""Plus operator.""" |
|
return SphericalManifold.plus(gravity.vec3d, delta) |
|
|
|
J_auto = vmap(jacfwd(spherical_plus))(delta).squeeze(0) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_up_projection_uv(self): |
|
"""Test the up projection jacobians.""" |
|
rpf = get_toy_rpf() |
|
|
|
r, p, f = rpf.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": [K1]}) |
|
gravity = Gravity.from_rp(r, p) |
|
uv = camera.normalize(camera.pixel_coordinates()) |
|
|
|
J = J_up_projection(uv, gravity.vec3d, "uv") |
|
|
|
|
|
def projection_uv(uv: torch.Tensor) -> torch.Tensor: |
|
"""Projection.""" |
|
abc = gravity.vec3d |
|
projected_up2d = abc[..., None, :2] - abc[..., 2, None, None] * uv |
|
return projected_up2d[0, 0] |
|
|
|
J_auto = vmap(jacfwd(projection_uv))(uv[0])[None] |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_up_projection_abc(self): |
|
"""Test the up projection jacobians.""" |
|
rpf = get_toy_rpf() |
|
|
|
r, p, f = rpf.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": [K1]}) |
|
gravity = Gravity.from_rp(r, p) |
|
uv = camera.normalize(camera.pixel_coordinates()) |
|
J = J_up_projection(uv, gravity.vec3d, "abc") |
|
|
|
|
|
def projection_abc(abc: torch.Tensor) -> torch.Tensor: |
|
"""Projection.""" |
|
return abc[..., None, :2] - abc[..., 2, None, None] * uv |
|
|
|
J_auto = vmap(jacfwd(projection_abc))(gravity.vec3d)[0] |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_undistort_pts(self): |
|
"""Test the undistortion jacobians.""" |
|
if CAMERA_MODEL == "pinhole": |
|
return |
|
|
|
rpf = get_toy_rpf() |
|
_, _, f = rpf.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": [K1]}) |
|
uv = camera.normalize(camera.pixel_coordinates()) |
|
J = camera.J_undistort(uv, "pts") |
|
|
|
|
|
def func_pts(pts): |
|
return camera.undistort(pts)[0][0] |
|
|
|
J_auto = vmap(jacfwd(func_pts))(uv[0])[None].squeeze(-3) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_undistort_k1(self): |
|
"""Test the undistortion jacobians.""" |
|
if CAMERA_MODEL == "pinhole": |
|
return |
|
|
|
rpf = get_toy_rpf() |
|
_, _, f = rpf.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": [K1]}) |
|
uv = camera.normalize(camera.pixel_coordinates()) |
|
J = camera.J_undistort(uv, "dist") |
|
|
|
|
|
def func_k1(k1): |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
return camera.undistort(uv)[0][0] |
|
|
|
J_auto = vmap(jacfwd(func_k1))(camera.dist[..., :1]).squeeze(-1) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_up_projection_offset(self): |
|
"""Test the up projection offset jacobians.""" |
|
if CAMERA_MODEL == "pinhole": |
|
return |
|
|
|
rpf = get_toy_rpf() |
|
|
|
_, _, f = rpf.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": [K1]}) |
|
uv = camera.normalize(camera.pixel_coordinates()) |
|
J = camera.up_projection_offset(uv) |
|
|
|
|
|
def projection_uv(uv: torch.Tensor) -> torch.Tensor: |
|
"""Projection.""" |
|
s, _ = camera.distort(uv, return_scale=True) |
|
return s[0, 0, 0] |
|
|
|
J_auto = vmap(jacfwd(projection_uv))(uv[0])[None].squeeze(-2) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_J_up_projection_offset_uv(self): |
|
"""Test the up projection offset jacobians.""" |
|
if CAMERA_MODEL == "pinhole": |
|
return |
|
|
|
rpf = get_toy_rpf() |
|
_, _, f = rpf.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": [K1]}) |
|
uv = camera.normalize(camera.pixel_coordinates()) |
|
J = camera.J_up_projection_offset(uv, "uv") |
|
|
|
|
|
def projection_uv(uv: torch.Tensor) -> torch.Tensor: |
|
"""Projection.""" |
|
return camera.up_projection_offset(uv)[0, 0] |
|
|
|
J_auto = vmap(jacfwd(projection_uv))(uv[0])[None] |
|
|
|
|
|
|
|
self.validate(J, J_auto) |
|
|
|
|
|
class TestEuclidean(unittest.TestCase): |
|
"""Test the Euclidean manifold jacobians.""" |
|
|
|
eps = 5e-3 |
|
|
|
def validate(self, J: torch.Tensor, J_auto: torch.Tensor): |
|
"""Check if the jacobians are close and finite.""" |
|
self.assertTrue(torch.all(torch.isfinite(J)), "found nan in numerical") |
|
self.assertTrue(torch.all(torch.isfinite(J_auto)), "found nan in auto") |
|
|
|
|
|
|
|
text_j = f" > {self.eps}\nJ:\n{J[0, 0, 0].numpy()}\nJ_auto:\n{J_auto[0, 0, 0].numpy()}" |
|
|
|
J_up2grav = J[..., :2, :2] |
|
J_up2grav_auto = J_auto[..., :2, :2] |
|
max_diff = torch.max(torch.abs(J_up2grav - J_up2grav_auto)) |
|
text = f"UP - GRAV max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J_up2grav, J_up2grav_auto, atol=self.eps), text) |
|
|
|
J_up2focal = J[..., :2, 2] |
|
J_up2focal_auto = J_auto[..., :2, 2] |
|
max_diff = torch.max(torch.abs(J_up2focal - J_up2focal_auto)) |
|
text = f"UP - FOCAL max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J_up2focal, J_up2focal_auto, atol=self.eps), text) |
|
|
|
if CAMERA_MODEL != "pinhole": |
|
J_up2k1 = J[..., :2, 3] |
|
J_up2k1_auto = J_auto[..., :2, 3] |
|
max_diff = torch.max(torch.abs(J_up2k1 - J_up2k1_auto)) |
|
text = f"UP - K1 max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J_up2k1, J_up2k1_auto, atol=self.eps), text) |
|
|
|
J_lat2grav = J[..., 2:, :2] |
|
J_lat2grav_auto = J_auto[..., 2:, :2] |
|
max_diff = torch.max(torch.abs(J_lat2grav - J_lat2grav_auto)) |
|
text = f"LAT - GRAV max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J_lat2grav, J_lat2grav_auto, atol=self.eps), text) |
|
|
|
J_lat2focal = J[..., 2:, 2] |
|
J_lat2focal_auto = J_auto[..., 2:, 2] |
|
max_diff = torch.max(torch.abs(J_lat2focal - J_lat2focal_auto)) |
|
text = f"LAT - FOCAL max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J_lat2focal, J_lat2focal_auto, atol=self.eps), text) |
|
|
|
if CAMERA_MODEL != "pinhole": |
|
J_lat2k1 = J[..., 2:, 3] |
|
J_lat2k1_auto = J_auto[..., 2:, 3] |
|
max_diff = torch.max(torch.abs(J_lat2k1 - J_lat2k1_auto)) |
|
text = f"LAT - K1 max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J_lat2k1, J_lat2k1_auto, atol=self.eps), text) |
|
|
|
max_diff = torch.max(torch.abs(J - J_auto[..., : J.shape[-1]])) |
|
text = f"Overall - max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J, J_auto[..., : J.shape[-1]], atol=self.eps), text) |
|
|
|
def local_pf_calc(self, rpfk: torch.Tensor): |
|
"""Calculate the perspective field.""" |
|
r, p, f, k1 = rpfk.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
gravity = Gravity.from_rp(r, p) |
|
up, lat = get_perspective_field(camera, gravity) |
|
persp = torch.cat([up, torch.sin(lat)], dim=-3) |
|
return persp.permute(0, 2, 3, 1).reshape(1, -1, 3) |
|
|
|
def test_random(self): |
|
"""Random rpf.""" |
|
rpf = get_toy_rpf() |
|
rpfk = torch.cat([rpf, torch.tensor([[K1]])], dim=-1) |
|
r, p, f, k1 = rpfk.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
gravity = Gravity.from_rp(r, p) |
|
|
|
J = torch.cat(J_perspective_field(camera, gravity, spherical=False), -2) |
|
J_auto = jacfwd(self.local_pf_calc)(rpfk).squeeze(-2, -3).reshape(1, H, W, 3, 4) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_zero_roll(self): |
|
"""Roll = 0.""" |
|
rpf = get_toy_rpf(roll=0) |
|
rpfk = torch.cat([rpf, torch.tensor([[K1]])], dim=-1) |
|
r, p, f, k1 = rpfk.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
gravity = Gravity.from_rp(r, p) |
|
|
|
J = torch.cat(J_perspective_field(camera, gravity, spherical=False), -2) |
|
J_auto = jacfwd(self.local_pf_calc)(rpfk).squeeze(-2, -3).reshape(1, H, W, 3, 4) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_zero_pitch(self): |
|
"""Pitch = 0.""" |
|
rpf = get_toy_rpf(pitch=0) |
|
rpfk = torch.cat([rpf, torch.tensor([[K1]])], dim=-1) |
|
r, p, f, k1 = rpfk.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
gravity = Gravity.from_rp(r, p) |
|
|
|
J = torch.cat(J_perspective_field(camera, gravity, spherical=False), -2) |
|
J_auto = jacfwd(self.local_pf_calc)(rpfk).squeeze(-2, -3).reshape(1, H, W, 3, 4) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_max_roll(self): |
|
"""Roll = -45, 45.""" |
|
for roll in [-45, 45]: |
|
rpf = get_toy_rpf(roll=roll) |
|
rpfk = torch.cat([rpf, torch.tensor([[K1]])], dim=-1) |
|
r, p, f, k1 = rpfk.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
gravity = Gravity.from_rp(r, p) |
|
|
|
J = torch.cat(J_perspective_field(camera, gravity, spherical=False), -2) |
|
J_auto = jacfwd(self.local_pf_calc)(rpfk).squeeze(-2, -3).reshape(1, H, W, 3, 4) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_max_pitch(self): |
|
"""Pitch = -45, 45.""" |
|
for pitch in [-45, 45]: |
|
rpf = get_toy_rpf(pitch=pitch) |
|
rpfk = torch.cat([rpf, torch.tensor([[K1]])], dim=-1) |
|
r, p, f, k1 = rpfk.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
gravity = Gravity.from_rp(r, p) |
|
|
|
J = torch.cat(J_perspective_field(camera, gravity, spherical=False), -2) |
|
J_auto = jacfwd(self.local_pf_calc)(rpfk).squeeze(-2, -3).reshape(1, H, W, 3, 4) |
|
|
|
self.validate(J, J_auto) |
|
|
|
|
|
class TestSpherical(unittest.TestCase): |
|
"""Test the spherical manifold jacobians.""" |
|
|
|
eps = 5e-3 |
|
|
|
def validate(self, J: torch.Tensor, J_auto: torch.Tensor): |
|
"""Check if the jacobians are close and finite.""" |
|
self.assertTrue(torch.all(torch.isfinite(J)), "found nan in numerical") |
|
self.assertTrue(torch.all(torch.isfinite(J_auto)), "found nan in auto") |
|
|
|
text_j = f" > {self.eps}\nJ:\n{J[0, 0, 0].numpy()}\nJ_auto:\n{J_auto[0, 0, 0].numpy()}" |
|
|
|
J_up2grav = J[..., :2, :2] |
|
J_up2grav_auto = J_auto[..., :2, :2] |
|
max_diff = torch.max(torch.abs(J_up2grav - J_up2grav_auto)) |
|
text = f"UP - GRAV max diff is {max_diff:.4f}" + text_j |
|
|
|
self.assertTrue(torch.allclose(J_up2grav, J_up2grav_auto, atol=self.eps), text) |
|
|
|
J_up2focal = J[..., :2, 2] |
|
J_up2focal_auto = J_auto[..., :2, 2] |
|
max_diff = torch.max(torch.abs(J_up2focal - J_up2focal_auto)) |
|
text = f"UP - FOCAL max diff is {max_diff:.4f}" + text_j |
|
|
|
self.assertTrue(torch.allclose(J_up2focal, J_up2focal_auto, atol=self.eps), text) |
|
|
|
if CAMERA_MODEL != "pinhole": |
|
J_up2k1 = J[..., :2, 3] |
|
J_up2k1_auto = J_auto[..., :2, 3] |
|
max_diff = torch.max(torch.abs(J_up2k1 - J_up2k1_auto)) |
|
text = f"UP - K1 max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J_up2k1, J_up2k1_auto, atol=self.eps), text) |
|
|
|
J_lat2grav = J[..., 2:, :2] |
|
J_lat2grav_auto = J_auto[..., 2:, :2] |
|
max_diff = torch.max(torch.abs(J_lat2grav - J_lat2grav_auto)) |
|
text = f"LAT - GRAV max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J_lat2grav, J_lat2grav_auto, atol=self.eps), text) |
|
|
|
J_lat2focal = J[..., 2:, 2] |
|
J_lat2focal_auto = J_auto[..., 2:, 2] |
|
max_diff = torch.max(torch.abs(J_lat2focal - J_lat2focal_auto)) |
|
text = f"LAT - FOCAL max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J_lat2focal, J_lat2focal_auto, atol=self.eps), text) |
|
|
|
if CAMERA_MODEL != "pinhole": |
|
J_lat2k1 = J[..., 2:, 3] |
|
J_lat2k1_auto = J_auto[..., 2:, 3] |
|
max_diff = torch.max(torch.abs(J_lat2k1 - J_lat2k1_auto)) |
|
text = f"LAT - K1 max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J_lat2k1, J_lat2k1_auto, atol=self.eps), text) |
|
|
|
max_diff = torch.max(torch.abs(J - J_auto[..., : J.shape[-1]])) |
|
text = f"Overall - max diff is {max_diff:.4f}" + text_j |
|
self.assertTrue(torch.allclose(J, J_auto[..., : J.shape[-1]], atol=self.eps), text) |
|
|
|
def local_pf_calc(self, uvfk: torch.Tensor, gravity: Gravity): |
|
"""Calculate the perspective field.""" |
|
delta, f, k1 = uvfk[..., :2], uvfk[..., 2], uvfk[..., 3] |
|
cam = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
up, lat = get_perspective_field(cam, gravity.update(delta, spherical=True)) |
|
persp = torch.cat([up, torch.sin(lat)], dim=-3) |
|
return persp.permute(0, 2, 3, 1).reshape(1, -1, 3) |
|
|
|
def test_random(self): |
|
"""Test random rpf.""" |
|
rpf = get_toy_rpf() |
|
rpfk = torch.cat([rpf, torch.tensor([[K1]])], dim=-1) |
|
r, p, f, k1 = rpfk.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
gravity = Gravity.from_rp(r, p) |
|
|
|
J = torch.cat(J_perspective_field(camera, gravity, spherical=True), -2) |
|
|
|
uvfk = torch.zeros_like(rpfk) |
|
uvfk[..., 2] = f |
|
uvfk[..., 3] = k1 |
|
func = lambda uvfk: self.local_pf_calc(uvfk, gravity) |
|
J_auto = jacfwd(func)(uvfk).squeeze(-2).reshape(1, H, W, 3, 4) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_zero_roll(self): |
|
"""Test roll = 0.""" |
|
rpf = get_toy_rpf(roll=0) |
|
rpfk = torch.cat([rpf, torch.tensor([[K1]])], dim=-1) |
|
r, p, f, k1 = rpfk.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
gravity = Gravity.from_rp(r, p) |
|
|
|
J = torch.cat(J_perspective_field(camera, gravity, spherical=True), -2) |
|
|
|
uvfk = torch.zeros_like(rpfk) |
|
uvfk[..., 2] = f |
|
uvfk[..., 3] = k1 |
|
func = lambda uvfk: self.local_pf_calc(uvfk, gravity) |
|
J_auto = jacfwd(func)(uvfk).squeeze(-2).reshape(1, H, W, 3, 4) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_zero_pitch(self): |
|
"""Test pitch = 0.""" |
|
rpf = get_toy_rpf(pitch=0) |
|
rpfk = torch.cat([rpf, torch.tensor([[K1]])], dim=-1) |
|
r, p, f, k1 = rpfk.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
gravity = Gravity.from_rp(r, p) |
|
|
|
J = torch.cat(J_perspective_field(camera, gravity, spherical=True), -2) |
|
|
|
uvfk = torch.zeros_like(rpfk) |
|
uvfk[..., 2] = f |
|
uvfk[..., 3] = k1 |
|
func = lambda uvfk: self.local_pf_calc(uvfk, gravity) |
|
J_auto = jacfwd(func)(uvfk).squeeze(-2).reshape(1, H, W, 3, 4) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_max_roll(self): |
|
"""Test roll = -45, 45.""" |
|
for roll in [-45, 45]: |
|
rpf = get_toy_rpf(roll=roll) |
|
rpfk = torch.cat([rpf, torch.tensor([[K1]])], dim=-1) |
|
r, p, f, k1 = rpfk.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
gravity = Gravity.from_rp(r, p) |
|
|
|
J = torch.cat(J_perspective_field(camera, gravity, spherical=True), -2) |
|
|
|
uvfk = torch.zeros_like(rpfk) |
|
uvfk[..., 2] = f |
|
uvfk[..., 3] = k1 |
|
func = lambda uvfk: self.local_pf_calc(uvfk, gravity) |
|
J_auto = jacfwd(func)(uvfk).squeeze(-2).reshape(1, H, W, 3, 4) |
|
|
|
self.validate(J, J_auto) |
|
|
|
def test_max_pitch(self): |
|
"""Test pitch = -45, 45.""" |
|
for pitch in [-45, 45]: |
|
rpf = get_toy_rpf(pitch=pitch) |
|
rpfk = torch.cat([rpf, torch.tensor([[K1]])], dim=-1) |
|
r, p, f, k1 = rpfk.unbind(dim=-1) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": f, "k1": k1}) |
|
gravity = Gravity.from_rp(r, p) |
|
|
|
J = torch.cat(J_perspective_field(camera, gravity, spherical=True), -2) |
|
|
|
uvfk = torch.zeros_like(rpfk) |
|
uvfk[..., 2] = f |
|
uvfk[..., 3] = k1 |
|
func = lambda uvfk: self.local_pf_calc(uvfk, gravity) |
|
J_auto = jacfwd(func)(uvfk).squeeze(-2).reshape(1, H, W, 3, 4) |
|
|
|
self.validate(J, J_auto) |
|
|
|
|
|
class TestLM(unittest.TestCase): |
|
"""Test the LM optimizer.""" |
|
|
|
eps = 1e-3 |
|
|
|
def test_random_spherical(self): |
|
"""Test random rpf.""" |
|
rpf = get_toy_rpf() |
|
gravity = Gravity.from_rp(rpf[..., 0], rpf[..., 1]) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": rpf[..., 2], "k1": [K1]}) |
|
|
|
up, lat = get_perspective_field(camera, gravity) |
|
|
|
lm = LMOptimizer({"use_spherical_manifold": True, "camera_model": CAMERA_MODEL}) |
|
|
|
out = lm({"up_field": up, "latitude_field": lat}) |
|
|
|
cam_opt = out["camera"] |
|
gravity_opt = out["gravity"] |
|
|
|
if hasattr(cam_opt, "k1"): |
|
text = f"cam_opt: {cam_opt.k1.numpy()} | rpf: {[K1]}" |
|
self.assertTrue( |
|
torch.allclose(cam_opt.k1, torch.tensor([K1]).float(), atol=self.eps), text |
|
) |
|
|
|
text = f"cam_opt: {cam_opt.f[..., 1].numpy()} | rpf: {rpf[..., 2].numpy()}" |
|
self.assertTrue(torch.allclose(cam_opt.f[..., 1], rpf[..., 2], atol=self.eps), text) |
|
|
|
text = f"gravity_opt.roll: {gravity_opt.roll.numpy()} | rpf: {rpf[..., 0].numpy()}" |
|
self.assertTrue(torch.allclose(gravity_opt.roll, rpf[..., 0], atol=self.eps), text) |
|
|
|
text = f"gravity_opt.pitch: {gravity_opt.pitch.numpy()} | rpf: {rpf[..., 1].numpy()}" |
|
self.assertTrue(torch.allclose(gravity_opt.pitch, rpf[..., 1], atol=self.eps), text) |
|
|
|
def test_random(self): |
|
"""Test random rpf.""" |
|
rpf = get_toy_rpf() |
|
gravity = Gravity.from_rp(rpf[..., 0], rpf[..., 1]) |
|
camera = Camera.from_dict({"height": [H], "width": [W], "f": rpf[..., 2], "k1": [K1]}) |
|
|
|
up, lat = get_perspective_field(camera, gravity) |
|
|
|
lm = LMOptimizer({"use_spherical_manifold": False, "camera_model": CAMERA_MODEL}) |
|
out = lm({"up_field": up, "latitude_field": lat}) |
|
|
|
cam_opt = out["camera"] |
|
gravity_opt = out["gravity"] |
|
|
|
if hasattr(cam_opt, "k1"): |
|
text = f"cam_opt: {cam_opt.k1.numpy()} | rpf: {[K1]}" |
|
self.assertTrue( |
|
torch.allclose(cam_opt.k1, torch.tensor([K1]).float(), atol=self.eps), text |
|
) |
|
|
|
text = f"cam_opt: {cam_opt.f[..., 1].numpy()} | rpf: {rpf[..., 2].numpy()}" |
|
self.assertTrue(torch.allclose(cam_opt.f[..., 1], rpf[..., 2], atol=self.eps), text) |
|
|
|
text = f"gravity_opt.roll: {gravity_opt.roll.numpy()} | rpf: {rpf[..., 0].numpy()}" |
|
self.assertTrue(torch.allclose(gravity_opt.roll, rpf[..., 0], atol=self.eps), text) |
|
|
|
text = f"gravity_opt.pitch: {gravity_opt.pitch.numpy()} | rpf: {rpf[..., 1].numpy()}" |
|
self.assertTrue(torch.allclose(gravity_opt.pitch, rpf[..., 1], atol=self.eps), text) |
|
|
|
|
|
if __name__ == "__main__": |
|
unittest.main() |
|
|