def set_project_mat_from_KRT(self, K, R, T):
"""set projection matrix from KRT matrices
Args:
K: (B, 3, 3)
R: (B, 3, 3) as matrices, or (B, 4) as quaternions
T: (B, 3)
"""
if self.camera_mode not in ['projection']:
raise ValueError('Only projection mode requires project mat (P)')
bs = K.shape[0]
if R.ndimension() == 3:
Rot = R
elif R.ndimension() == 2:
if R.shape[1] == 4: # quaternion
try:
from liegroups.torch import SO3
except:
raise ImportError(
f"failed to 'from liegroups.torch import SO3'")
import torch.nn.functional as F
R = F.normalize(R, eps=1e-6)
Rot = SO3.from_quaternion(R).mat
if Rot.ndimension() == 2:
Rot = Rot[
None, :, :] # liegroups tends to squeeze the results
else:
raise RuntimeError(f"invalid R.shape = {R.shape}")
else:
raise RuntimeError(f"invalid R.ndimension() = {R.ndimension()}")
P = torch.bmm(K, torch.cat((Rot, T.view(-1, 3, 1)), dim=2))
self.transformer.P = P
def test_quaternion():
q1 = torch.Tensor([1, 0, 0, 0])
q2 = torch.Tensor([0, 1, 0, 0])
q3 = torch.Tensor([0, 0, 1, 0])
q4 = torch.Tensor([0, 0, 0, 1])
q5 = 0.5 * torch.ones(4)
q6 = -q5
assert utils.allclose(SO3.from_quaternion(q1).to_quaternion(), q1)
assert utils.allclose(SO3.from_quaternion(q2).to_quaternion(), q2)
assert utils.allclose(SO3.from_quaternion(q3).to_quaternion(), q3)
assert utils.allclose(SO3.from_quaternion(q4).to_quaternion(), q4)
assert utils.allclose(SO3.from_quaternion(q5).to_quaternion(), q5)
assert utils.allclose(
SO3.from_quaternion(q5).mat,
SO3.from_quaternion(q6).mat)
def test_quaternion_batch():
quats = torch.Tensor([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0],
[0, 0, 0, 1], [0.5, 0.5, 0.5, 0.5]])
assert utils.allclose(SO3.from_quaternion(quats).to_quaternion(), quats)