Estimate the focal lengths of N input images. images: N x 3 x H x W masks: N x 1 x H x W
(
self,
images,
masks=None,
use_first_focal=False,
)
| 65 | |
| 66 | @torch.inference_mode() |
| 67 | def estimate_focals( |
| 68 | self, |
| 69 | images, |
| 70 | masks=None, |
| 71 | use_first_focal=False, |
| 72 | ): |
| 73 | """ |
| 74 | Estimate the focal lengths of N input images. |
| 75 | |
| 76 | images: N x 3 x H x W |
| 77 | masks: N x 1 x H x W |
| 78 | """ |
| 79 | assert images.ndim == 4 |
| 80 | N, _, H, W = images.shape |
| 81 | assert H == W, "Non-square images are not supported." |
| 82 | |
| 83 | pp = self.pp.to(images) |
| 84 | # pp = torch.tensor([W/2, H/2]).to(images) |
| 85 | |
| 86 | focals = [] |
| 87 | for i in range(N): |
| 88 | if use_first_focal and i > 0: |
| 89 | break |
| 90 | images_input = torch.cat([images[i:], images[:i]], dim=0) |
| 91 | gaussians = self.forward_gaussians(images_input.unsqueeze(0)) # 1 x (N x H x W) x 14 |
| 92 | points = rearrange(gaussians[0, :H*W, :3], '(h w) c -> h w c', h=H, w=W) |
| 93 | mask = masks[i] if masks is not None else None |
| 94 | focal = estimate_focal(points, pp=pp, mask=mask) |
| 95 | focals.append(focal) |
| 96 | |
| 97 | focals = torch.stack(focals).to(images) |
| 98 | focals = focals.mean().reshape(1).repeat(N) |
| 99 | return focals |
| 100 | |
| 101 | @torch.inference_mode() |
| 102 | def estimate_poses( |
no test coverage detected