| 31 | |
| 32 | |
| 33 | class SparseGA(): |
| 34 | def __init__(self, img_paths, pairs_in, res_fine, anchors, canonical_paths=None): |
| 35 | def fetch_img(im): |
| 36 | def torgb(x): return (x[0].permute(1, 2, 0).numpy() * .5 + .5).clip(min=0., max=1.) |
| 37 | for im1, im2 in pairs_in: |
| 38 | if im1['instance'] == im: |
| 39 | return torgb(im1['img']) |
| 40 | if im2['instance'] == im: |
| 41 | return torgb(im2['img']) |
| 42 | self.canonical_paths = canonical_paths |
| 43 | self.img_paths = img_paths |
| 44 | self.imgs = [fetch_img(img) for img in img_paths] |
| 45 | self.intrinsics = res_fine['intrinsics'] |
| 46 | self.cam2w = res_fine['cam2w'] |
| 47 | self.depthmaps = res_fine['depthmaps'] |
| 48 | self.pts3d = res_fine['pts3d'] |
| 49 | self.pts3d_colors = [] |
| 50 | self.working_device = self.cam2w.device |
| 51 | for i in range(len(self.imgs)): |
| 52 | im = self.imgs[i] |
| 53 | x, y = anchors[i][0][..., :2].detach().cpu().numpy().T |
| 54 | self.pts3d_colors.append(im[y, x]) |
| 55 | assert self.pts3d_colors[-1].shape == self.pts3d[i].shape |
| 56 | self.n_imgs = len(self.imgs) |
| 57 | |
| 58 | def get_focals(self): |
| 59 | return torch.tensor([ff[0, 0] for ff in self.intrinsics]).to(self.working_device) |
| 60 | |
| 61 | def get_principal_points(self): |
| 62 | return torch.stack([ff[:2, -1] for ff in self.intrinsics]).to(self.working_device) |
| 63 | |
| 64 | def get_im_poses(self): |
| 65 | return self.cam2w |
| 66 | |
| 67 | def get_sparse_pts3d(self): |
| 68 | return self.pts3d |
| 69 | |
| 70 | def get_dense_pts3d(self, clean_depth=True, subsample=8): |
| 71 | assert self.canonical_paths, 'cache_path is required for dense 3d points' |
| 72 | device = self.cam2w.device |
| 73 | confs = [] |
| 74 | base_focals = [] |
| 75 | anchors = {} |
| 76 | for i, canon_path in enumerate(self.canonical_paths): |
| 77 | (canon, canon2, conf), focal = torch.load(canon_path, map_location=device) |
| 78 | confs.append(conf) |
| 79 | base_focals.append(focal) |
| 80 | |
| 81 | H, W = conf.shape |
| 82 | pixels = torch.from_numpy(np.mgrid[:W, :H].T.reshape(-1, 2)).float().to(device) |
| 83 | idxs, offsets = anchor_depth_offsets(canon2, {i: (pixels, None)}, subsample=subsample) |
| 84 | anchors[i] = (pixels, idxs[i], offsets[i]) |
| 85 | |
| 86 | # densify sparse depthmaps |
| 87 | pts3d, depthmaps = make_pts3d(anchors, self.intrinsics, self.cam2w, [ |
| 88 | d.ravel() for d in self.depthmaps], base_focals=base_focals, ret_depth=True) |
| 89 | |
| 90 | if clean_depth: |
no outgoing calls
no test coverage detected