(self, i)
| 539 | return len(self.idx_to_cam_name_and_frame_idx) |
| 540 | |
| 541 | def __getitem__(self, i): |
| 542 | cam_name_and_frame_idx_pairs = self.idx_to_cam_name_and_frame_idx[i] # list of (cam_name, frame_idx) |
| 543 | surface_normal_w = [] |
| 544 | z_map = [] |
| 545 | rgb = [] |
| 546 | H_c2ws = [] |
| 547 | for i in range(len(cam_name_and_frame_idx_pairs)): |
| 548 | cam_name, frame_idx = cam_name_and_frame_idx_pairs[i] |
| 549 | # print(f'cam_name = {cam_name}, frame_idx = {frame_idx}', flush=True) |
| 550 | d_dict = self._get_frame( |
| 551 | camera_name=cam_name, |
| 552 | frame_id=frame_idx, |
| 553 | ) |
| 554 | surface_normal_w.append(d_dict['surface_normal_w']) # (h, w, 3) |
| 555 | z_map.append(d_dict['z_map']) # (h, w) |
| 556 | rgb.append(d_dict['rgb']) # (h, w, 3) |
| 557 | H_c2ws.append(self.H_c2ws[cam_name][self.cam_name_to_frame_idxs_to_idxs[cam_name][frame_idx]]) # (4, 4) |
| 558 | |
| 559 | surface_normal_w = torch.stack(surface_normal_w, dim=0) # (num_images_per_item, h, w, 3) |
| 560 | z_map = torch.stack(z_map, dim=0) # (num_images_per_item, h, w) |
| 561 | rgb = torch.stack(rgb, dim=0) # (num_images_per_item, h, w, 3) |
| 562 | H_c2ws = torch.stack(H_c2ws, dim=0) # (num_images_per_item, 4, 4) |
| 563 | |
| 564 | # handle image subsample |
| 565 | ori_w = self.width_px |
| 566 | h = self.height_px // self.image_subsample |
| 567 | w = self.width_px // self.image_subsample |
| 568 | intrinsic = self.intrinsic.clone() |
| 569 | intrinsic[:2, :] = intrinsic[:2, :] * w / ori_w |
| 570 | # intrinsic = self.intrinsic * w / ori_w |
| 571 | # intrinsic[..., 2, 2] = ori22 |
| 572 | rgb = rgb[..., ::self.image_subsample, ::self.image_subsample, :] |
| 573 | z_map = z_map[..., ::self.image_subsample, ::self.image_subsample] |
| 574 | surface_normal_w = surface_normal_w[..., ::self.image_subsample, ::self.image_subsample, :] |
| 575 | |
| 576 | camera = structures.Camera( |
| 577 | H_c2w=H_c2ws.unsqueeze(0), # (1, num_images_per_item, 4, 4) |
| 578 | intrinsic=intrinsic.view(1, 1, 3, 3).expand(1, rgb.size(0), 3, 3), |
| 579 | width_px=w, |
| 580 | height_px=h, |
| 581 | ) |
| 582 | |
| 583 | return dict( |
| 584 | rgbd_image=structures.RGBDImage( |
| 585 | rgb=rgb.unsqueeze(0), # (1, num_images_per_item, h, w, 3) |
| 586 | depth=z_map.unsqueeze(0), # (1, num_images_per_item, h, w) |
| 587 | camera=camera, |
| 588 | normal_w=surface_normal_w.unsqueeze(0), # (1, num_images_per_item, h, w, 3) |
| 589 | hit_map=z_map.unsqueeze(0) < 1e6, # (1, num_images_per_item, h, w) |
| 590 | ) |
| 591 | ) |
nothing calls this directly
no test coverage detected