Args: epoch: bidx: batch: input_rgbd_images: RGBDImage, (b, q, h, w), images along q should be used to create point cloud ray: Ray, (b, q=n_target_img, ho, wo) target rays
(
self,
epoch: int,
bidx: int,
batch: T.Any,
update: bool,
)
| 552 | raise NotImplementedError |
| 553 | |
| 554 | def _step( |
| 555 | self, |
| 556 | epoch: int, |
| 557 | bidx: int, |
| 558 | batch: T.Any, |
| 559 | update: bool, |
| 560 | ): |
| 561 | """ |
| 562 | Args: |
| 563 | epoch: |
| 564 | bidx: |
| 565 | batch: |
| 566 | input_rgbd_images: |
| 567 | RGBDImage, (b, q, h, w), images along q should be used to create point cloud |
| 568 | ray: |
| 569 | Ray, (b, q=n_target_img, ho, wo) target rays |
| 570 | ray_gt_dict: |
| 571 | ray_rgbs: (b, q=n_target_img, ho, wo, 3) |
| 572 | ray_ts: (b, q=n_target_img, ho, wo) |
| 573 | surface_normals_w: (b, q=n_target_img, ho, wo, 3) |
| 574 | hit_map: (b, q=n_target_img, ho, wo) 1 if hit a surface, 0 otherwise |
| 575 | update: |
| 576 | |
| 577 | Returns: |
| 578 | |
| 579 | """ |
| 580 | |
| 581 | loss = 0 |
| 582 | total_stime = timer() |
| 583 | with torch.autocast( |
| 584 | device_type='cuda' if self.process_info['n_gpus'] > 0 else 'cpu', |
| 585 | enabled=self.optim_info['use_amp'], |
| 586 | ): |
| 587 | input_rgbd_images: structures.RGBDImage = batch['input_rgbd_images'].to(device=self.device) # (b, q, h, w) |
| 588 | ray: structures.Ray = batch['ray'].to(device=self.device) # (b, qo, ho, wo) |
| 589 | ray_gt_dict = batch['ray_gt_dict'] |
| 590 | gt_ray_rgb: torch.Tensor = ray_gt_dict['ray_rgbs'].to(device=self.device) # (b, qo, ho, wo, 3) |
| 591 | gt_ray_ts: torch.Tensor = ray_gt_dict['ray_ts'].to(device=self.device) # (b, qo, ho, wo) |
| 592 | gt_surface_normals_w: torch.Tensor = ray_gt_dict['surface_normals_w'].to( |
| 593 | device=self.device) # (b, qo, ho, wo, 3) |
| 594 | gt_hit_map: torch.Tensor = ray_gt_dict['hit_map'].to(device=self.device) # (b, qo, ho, wo) |
| 595 | |
| 596 | input_point_cloud: structures.PointCloud = input_rgbd_images.get_pcd( |
| 597 | subsample=1, |
| 598 | remove_background=(input_rgbd_images.rgb.size(0) == 1), |
| 599 | ) # (b, n, 3) |
| 600 | |
| 601 | # add noise to input point cloud |
| 602 | if self.optim_info['pcd_noise_std'] > 1.e-6: |
| 603 | noise = torch.randn_like(input_point_cloud.xyz_w) * self.optim_info['pcd_noise_std'] |
| 604 | input_point_cloud.xyz_w = input_point_cloud.xyz_w + noise |
| 605 | |
| 606 | # determine k to use |
| 607 | r_k_ratio = torch.rand(1).item() * \ |
| 608 | (self.dataset_info.get('max_k_ratio', 1.) - self.dataset_info.get('min_k_ratio', 1.)) + \ |
| 609 | self.dataset_info.get('min_k_ratio', 1.) |
| 610 | k = max(10, int(self.dataset_info.get('k', 40) * r_k_ratio)) |
| 611 | # self.logger.info(f'k = {k}, subsample = {r_subsample}') |