Estimate the camera poses of N input images. images: N x 3 x h x W gaussians: K x 14 or 1 x K x 14 masks: N x 1 x H x W focals: N
(
self,
images,
gaussians=None,
masks=None,
focals=None,
use_first_focal=True,
opacity_threshold=5e-2,
pnp_iter=20,
)
| 100 | |
| 101 | @torch.inference_mode() |
| 102 | def estimate_poses( |
| 103 | self, |
| 104 | images, |
| 105 | gaussians=None, |
| 106 | masks=None, |
| 107 | focals=None, |
| 108 | use_first_focal=True, |
| 109 | opacity_threshold=5e-2, |
| 110 | pnp_iter=20, |
| 111 | ): |
| 112 | """ |
| 113 | Estimate the camera poses of N input images. |
| 114 | |
| 115 | images: N x 3 x h x W |
| 116 | gaussians: K x 14 or 1 x K x 14 |
| 117 | masks: N x 1 x H x W |
| 118 | focals: N |
| 119 | """ |
| 120 | assert images.ndim == 4 |
| 121 | N, _, H, W = images.shape |
| 122 | assert H == W, "Non-square images are not supported." |
| 123 | |
| 124 | # predict gaussians from images |
| 125 | if gaussians is None: |
| 126 | gaussians = self.forward_gaussians(images.unsqueeze(0)) # 1 x (N x H x W) x 14 |
| 127 | else: |
| 128 | if gaussians.ndim == 2: |
| 129 | gaussians = gaussians.unsqueeze(0) |
| 130 | assert gaussians.shape[1] == N * H * W |
| 131 | |
| 132 | points = gaussians[..., :3].reshape(1, N, H, W, 3).squeeze(0) # N x H x W x 3 |
| 133 | opacities = gaussians[..., 3+self.sh_dim].reshape(1, N, H, W).squeeze(0) |
| 134 | opacities = torch.sigmoid(opacities) # N x H x W |
| 135 | |
| 136 | # estimate focals if not provided |
| 137 | if focals is None: |
| 138 | focals = self.estimate_focals(images, masks=masks, use_first_focal=use_first_focal) |
| 139 | |
| 140 | # run PnP |
| 141 | c2ws = [] |
| 142 | for i in range(N): |
| 143 | pts3d = points[i].float().detach().cpu().numpy() |
| 144 | # If masks are not provided, we use Gaussian opacities |
| 145 | if masks is None: |
| 146 | mask = (opacities[i] > opacity_threshold).detach().cpu().numpy() |
| 147 | else: |
| 148 | mask = masks[i].reshape(H, W).bool().detach().cpu().numpy() |
| 149 | |
| 150 | focal = focals[i].item() |
| 151 | _, c2w = fast_pnp(pts3d, mask, focal=focal, niter_PnP=pnp_iter) |
| 152 | |
| 153 | c2ws.append(torch.from_numpy(c2w)) |
| 154 | |
| 155 | c2ws = torch.stack(c2ws, dim=0).to(images) |
| 156 | return c2ws, focals |
no test coverage detected