| 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] |
|
|
| |
| 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) |
|
|
| |
| self.r = avg_r.moveaxis(2,1).contiguous() |
| self.t = (-self.r @ avg_t.unsqueeze(-1)).squeeze(-1).unsqueeze(1).contiguous() |
|
|
| |
| 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): |
| |
| 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 |
|
|
