Returns: input_rgbd_images: RGBDImage, (b, q, h, w), ray: Ray, (b, q=n_target_img, ho, wo) target rays ray_gt_dict: ray_rgbs: (b, q=n_target_img, ho, wo, 3) ray_ts: (b, q=n_target_img, ho, w
(
self,
input_list: T.List[T.Dict[str, T.Any]],
)
| 338 | class MeshDatasetCollate: |
| 339 | |
| 340 | def __call__( |
| 341 | self, |
| 342 | input_list: T.List[T.Dict[str, T.Any]], |
| 343 | ): |
| 344 | """ |
| 345 | Returns: |
| 346 | input_rgbd_images: |
| 347 | RGBDImage, (b, q, h, w), |
| 348 | ray: |
| 349 | Ray, (b, q=n_target_img, ho, wo) target rays |
| 350 | ray_gt_dict: |
| 351 | ray_rgbs: (b, q=n_target_img, ho, wo, 3) |
| 352 | ray_ts: (b, q=n_target_img, ho, wo) |
| 353 | surface_normals_w: (b, q=n_target_img, ho, wo, 3) |
| 354 | hit_map: (b, q=n_target_img, ho, wo) 1 if hit a surface, 0 otherwise |
| 355 | """ |
| 356 | |
| 357 | # input rgbd images |
| 358 | input_rgbd_images = structures.RGBDImage.cat([p['input_rgbd_images'] for p in input_list], dim=0) # (b, n) |
| 359 | |
| 360 | # ray |
| 361 | ray = structures.Ray.cat([p['ray'] for p in input_list], dim=0) # (b, q, h, w) |
| 362 | |
| 363 | # ray_gt_dict |
| 364 | ray_gt_dict = dict() |
| 365 | for key in ['ray_rgbs', 'ray_ts', 'surface_normals_w', 'hit_map']: |
| 366 | arr = torch.cat([p['ray_gt_dict'][key] for p in input_list], dim=0) |
| 367 | ray_gt_dict[key] = arr |
| 368 | |
| 369 | return dict( |
| 370 | input_rgbd_images=input_rgbd_images, |
| 371 | ray=ray, |
| 372 | ray_gt_dict=ray_gt_dict, |
| 373 | ) |
| 374 | |
| 375 | |
| 376 | class MeshConcatDataset(torch.utils.data.Dataset): |