Reproject RGBD pixels to world coordinate. Args: subsample: use 1 out of every `subsample` pixels remove_background: whether to remove background pixels (depth > 1e6). only has effect if b == 1. kee
(
self,
subsample: int = 1,
remove_background: bool = True,
keep_img_idxs: bool = False,
)
| 2226 | return RGBDImage(**out) |
| 2227 | |
| 2228 | def get_pcd( |
| 2229 | self, |
| 2230 | subsample: int = 1, |
| 2231 | remove_background: bool = True, |
| 2232 | keep_img_idxs: bool = False, |
| 2233 | ) -> PointCloud: |
| 2234 | """ |
| 2235 | Reproject RGBD pixels to world coordinate. |
| 2236 | |
| 2237 | Args: |
| 2238 | subsample: |
| 2239 | use 1 out of every `subsample` pixels |
| 2240 | remove_background: |
| 2241 | whether to remove background pixels (depth > 1e6). |
| 2242 | only has effect if b == 1. |
| 2243 | keep_img_idxs: |
| 2244 | whether to record the original image index in the point cloud |
| 2245 | """ |
| 2246 | assert self.depth is not None |
| 2247 | # only has effect if b == 1 |
| 2248 | valid_mask = self.hit_map # None or (b, q, h, w) |
| 2249 | if remove_background: |
| 2250 | if valid_mask is None: |
| 2251 | valid_mask = self.depth < 1.e6 # (b, q, h, w) |
| 2252 | else: |
| 2253 | valid_mask = torch.logical_and(valid_mask, self.depth < 1.e6) # (b, q, h, w) |
| 2254 | depth = self.depth.masked_fill( |
| 2255 | mask=torch.logical_not(valid_mask), |
| 2256 | value=INF, |
| 2257 | ) |
| 2258 | |
| 2259 | if keep_img_idxs: |
| 2260 | b, q, h, w = self.depth.shape |
| 2261 | qhw = q * h * w |
| 2262 | img_idxs = torch.arange(qhw, device=self.depth.device).expand(b, -1) # (b, qhw) |
| 2263 | img_idxs = img_idxs.reshape(b, q, h, w, 1) # (b, q, h, w, 1) it saves the linear index of qhw |
| 2264 | else: |
| 2265 | img_idxs = None |
| 2266 | |
| 2267 | other_maps = [ |
| 2268 | self.rgb, |
| 2269 | self.normal_w, |
| 2270 | valid_mask.unsqueeze(-1), |
| 2271 | self.feature, |
| 2272 | img_idxs, |
| 2273 | ] |
| 2274 | |
| 2275 | xyz_dict = utils.compute_3d_xyz( |
| 2276 | z_map=depth, |
| 2277 | intrinsic=self.camera.intrinsic, |
| 2278 | H_c2w=self.camera.H_c2w, |
| 2279 | subsample=subsample, |
| 2280 | other_maps=other_maps, |
| 2281 | ) # (b, q, h', w', 3) can contain nan |
| 2282 | |
| 2283 | points_w = xyz_dict['xyz_w'] # (b, q, h', w', 3) |
| 2284 | b, *qhw, c = points_w.shape |
| 2285 | points_w = points_w.reshape(b, -1, 3) # (b, n, 3) |
no test coverage detected