Return rgbd_image (with surface normal) to compare with other methods. Camera should be the orignal camera used to cast the camera rays. Only support shape = (b, q, h, w).
(
self,
camera: 'Camera',
th_hit_prob: float = None,
th_dot_product: float = None, # (less than 60 degree)
use_plane_normal: bool = False,
)
| 1503 | torch.save(self.state_dict(), filename) |
| 1504 | |
| 1505 | def get_rgbd_image( |
| 1506 | self, |
| 1507 | camera: 'Camera', |
| 1508 | th_hit_prob: float = None, |
| 1509 | th_dot_product: float = None, # (less than 60 degree) |
| 1510 | use_plane_normal: bool = False, |
| 1511 | ) -> 'RGBDImage': |
| 1512 | """ |
| 1513 | Return rgbd_image (with surface normal) to compare with other methods. |
| 1514 | Camera should be the orignal camera used to cast the camera rays. |
| 1515 | Only support shape = (b, q, h, w). |
| 1516 | """ |
| 1517 | b, *m_shape, _ = self.intersection_xyz_w.shape # (b, q, h, w, 3) |
| 1518 | assert len(m_shape) == 3 # (q, h, w) |
| 1519 | q, h, w = m_shape |
| 1520 | |
| 1521 | if th_hit_prob is None: |
| 1522 | th_hit_prob = 0 |
| 1523 | if th_dot_product is None: |
| 1524 | th_dot_product = 0 |
| 1525 | |
| 1526 | # xyz_w -> xyz_c |
| 1527 | H_w2c = camera.get_H_w2c() # (b, q, 4, 4) |
| 1528 | H_w2c = H_w2c.reshape(b, q, 1, 1, 4, 4) # (b, q, 1, 1, 4, 4) |
| 1529 | xyz_w = self.intersection_xyz_w.unsqueeze(-1) # (b, q, h, w, 3, 1) |
| 1530 | xyz_c = H_w2c[..., :3, :3] @ xyz_w + H_w2c[..., :3, 3:4] # (b, q, h, w, 3, 1) |
| 1531 | # depth is the z in camera coordinate |
| 1532 | depth = xyz_c[..., 2, 0] # (b, q, h, w) |
| 1533 | assert depth.shape == torch.Size([b, *m_shape]) |
| 1534 | |
| 1535 | # determine hit map to use |
| 1536 | if self.refined_ray_hit is not None: |
| 1537 | hit_map = self.refined_ray_hit # (b, q, h, w) |
| 1538 | else: |
| 1539 | hit_map = self.ray_hit # (b, q, h, w) |
| 1540 | |
| 1541 | if th_dot_product > 1e-6 or th_hit_prob > 1e-6: |
| 1542 | ray = camera.generate_camera_rays(device=camera.H_c2w.device) |
| 1543 | confidence = self.compute_confidence( |
| 1544 | ray_direction_w=ray.directions_w, # (b, q, h, w, 3) |
| 1545 | th_hit_prob=th_hit_prob, |
| 1546 | th_dot_product=th_dot_product, |
| 1547 | ) # (b, q, h, w) |
| 1548 | hit_map = torch.logical_and(hit_map, confidence) |
| 1549 | |
| 1550 | if use_plane_normal and self.intersection_plane_normals_w is not None: |
| 1551 | normal_w = self.intersection_plane_normals_w * hit_map.unsqueeze(-1) # (b, q, h, w, 3) |
| 1552 | else: |
| 1553 | normal_w = self.intersection_surface_normal_w * hit_map.unsqueeze(-1) # (b, q, h, w, 3) |
| 1554 | |
| 1555 | return RGBDImage( |
| 1556 | rgb=self.intersection_rgb * hit_map.unsqueeze(-1), # (b, q, h, w, 3) |
| 1557 | depth=depth * hit_map, # (b, q, h, w) |
| 1558 | normal_w=normal_w, # (b, q, h, w, 3) |
| 1559 | hit_map=hit_map, |
| 1560 | camera=camera, |
| 1561 | ) |
| 1562 |
nothing calls this directly
no test coverage detected