Wrapper class for the open3d fuser. This wrapper does not support fusion of tensors with higher than batch 1.
| 609 | |
| 610 | |
| 611 | class Open3DFuser(DepthFuser): |
| 612 | """ |
| 613 | Wrapper class for the open3d fuser. |
| 614 | |
| 615 | This wrapper does not support fusion of tensors with higher than batch 1. |
| 616 | """ |
| 617 | |
| 618 | def __init__( |
| 619 | self, |
| 620 | gt_path="", |
| 621 | fusion_resolution=0.04, |
| 622 | max_fusion_depth=3, |
| 623 | fuse_color=False, |
| 624 | use_upsample_depth=False, |
| 625 | ): |
| 626 | super().__init__( |
| 627 | gt_path, |
| 628 | fusion_resolution, |
| 629 | max_fusion_depth, |
| 630 | fuse_color, |
| 631 | ) |
| 632 | |
| 633 | self.fuse_color = fuse_color |
| 634 | self.use_upsample_depth = use_upsample_depth |
| 635 | self.fusion_max_depth = max_fusion_depth |
| 636 | |
| 637 | voxel_size = fusion_resolution * 100 |
| 638 | self.volume = o3d.pipelines.integration.ScalableTSDFVolume( |
| 639 | voxel_length=float(voxel_size) / 100, |
| 640 | sdf_trunc=3 * float(voxel_size) / 100, |
| 641 | color_type=o3d.pipelines.integration.TSDFVolumeColorType.RGB8 |
| 642 | ) |
| 643 | |
| 644 | def fuse_frames( |
| 645 | self, |
| 646 | depths_b1hw, |
| 647 | K_b44, |
| 648 | cam_T_world_b44, |
| 649 | color_b3hw, |
| 650 | ): |
| 651 | |
| 652 | width = depths_b1hw.shape[-1] |
| 653 | height = depths_b1hw.shape[-2] |
| 654 | |
| 655 | if self.fuse_color: |
| 656 | color_b3hw = torch.nn.functional.interpolate( |
| 657 | color_b3hw, |
| 658 | size=(height, width), |
| 659 | ) |
| 660 | color_b3hw = reverse_imagenet_normalize(color_b3hw) |
| 661 | |
| 662 | for batch_index in range(depths_b1hw.shape[0]): |
| 663 | if self.fuse_color: |
| 664 | image_i = color_b3hw[batch_index].permute(1, 2, 0) |
| 665 | |
| 666 | color_im = (image_i * 255).cpu().numpy().astype( |
| 667 | np.uint8 |
| 668 | ).copy(order='C') |