Fix a minor bug: use point to point depth for supervising neural rend…

…ering
OpenGVLab · Mar 15, 2024 · bd8cda8 · bd8cda8
1 parent 6576cc1
commit bd8cda8
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/ponder/models/ponder/ponder_indoor_base.py b/ponder/models/ponder/ponder_indoor_base.py
@@ -557,6 +557,14 @@ def ray_sample(self, data_dict):
                 ray_o = ray_o[pixels_y, pixels_x, :]
                 ray_d = ray_d[pixels_y, pixels_x, :]
 
+                # convert plane-to-plane depth to point-to-point depth
+                cam2lidar = torch.linalg.inv(RT)
+                plane_dir = (
+                    cam2lidar @ torch.FloatTensor([0, 0, 1, 1]).to(RT.device)[:, None]
+                )[:3, 0] - ray_o[0]
+                plane_dir = plane_dir / torch.linalg.norm(plane_dir)
+                depth = depth / torch.linalg.multi_dot((ray_d, plane_dir))
+
                 mask_at_box = torch.from_numpy(
                     self.get_mask_at_box(ray_o.cpu().numpy(), ray_d.cpu().numpy())
                 ).to(ray_o.device)