File size: 2,070 Bytes
cebb76b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | import torch
class PosePreprocessor:
def __init__(self, c2w_r, c2w_t, subset_idx=None, scale_factor=0.75):
"""
Compute a series of [offset -> scale -> rotation] for normalizing poses.
c2w_r: [b, v, 3, 3]
c2w_t: [b, v, 3]
subset_idx: [v] or None. Use a subset of view to compute normalization.
scale_factor: Rescale factor for the normalized camera
"""
assert len(c2w_r.shape) == 4 and c2w_r.shape[-2:] == (3, 3)
assert len(c2w_t.shape) == 3 and c2w_t.shape[2] == 3
assert c2w_r.shape[:2] == c2w_t.shape[:2]
# Compute SE(3) normalizing factor
if subset_idx is not None:
assert len(subset_idx.shape) == 1 and len(subset_idx) == c2w_r.shape[1]
avg_t = c2w_t[:, subset_idx].mean(1)
avg_r = compute_avg_rotation(c2w_r[:, subset_idx])
else:
avg_t = c2w_t.mean(1)
avg_r = compute_avg_rotation(c2w_r)
# Target transformation for normalization
self.r = avg_r.moveaxis(2,1).contiguous()
self.t = (-self.r @ avg_t.unsqueeze(-1)).squeeze(-1).unsqueeze(1).contiguous()
# Target scene rescaling
new_t = torch.einsum('bij, bvj -> bvi', self.r, c2w_t) + self.t
post_scale = new_t.abs().amax((1,2), keepdim=True)
self.s = scale_factor / post_scale
def __call__(self, c2w_r, c2w_t):
c2w_r = torch.einsum('bij, bvjk -> bvik', self.r, c2w_r)
c2w_t = torch.einsum('bij, bvj -> bvi', self.r, c2w_t) + self.t
c2w_t = c2w_t * self.s
return c2w_r, c2w_t
def compute_avg_rotation(c2w_r):
# c2w_r: [b, v, 3, 3]
avg_down = torch.nn.functional.normalize(c2w_r[..., 1].mean(1), dim=-1)
avg_forward = c2w_r[..., 2].mean(1)
avg_right = torch.nn.functional.normalize(torch.cross(avg_down, avg_forward, dim=-1), dim=-1)
avg_forward = torch.nn.functional.normalize(torch.cross(avg_right, avg_down, dim=-1), dim=-1)
avg_r = torch.stack([avg_right, avg_down, avg_forward], dim=-1)
return avg_r
|