Improve visualization of ray casting

src/active_grasp/nbv.py

@@ -76,7 +76,7 @@ class NextBestView(MultiViewPolicy):
             with Timer("cost_computation"):
                 costs = [self.cost_fn(v) for v in views]
             utilities = gains / np.sum(gains) - costs / np.sum(costs)
-            self.vis.views(self.base_frame, self.intrinsic, views, utilities)
+            self.vis.ig_views(self.base_frame, self.intrinsic, views, utilities)
             i = np.argmax(utilities)
             nbv, gain = views[i], gains[i]
 
@@ -104,7 +104,7 @@ class NextBestView(MultiViewPolicy):
                 view_candidates.append(view)
         return view_candidates
 
-    def ig_fn(self, view, downsample):
+    def ig_fn(self, view, downsample, vis=False):
         tsdf_grid, voxel_size = self.tsdf.get_grid(), self.tsdf.voxel_size
         tsdf_grid = -1.0 + 2.0 * tsdf_grid  # Open3D maps tsdf to [0,1]
 
@@ -157,6 +157,14 @@ class NextBestView(MultiViewPolicy):
         tsdfs = tsdf_grid[i, j, k]
         ig = np.logical_and(tsdfs > -1.0, tsdfs < 0.0).sum()
 
+        if vis:
+            dirs = []
+            for u in range(u_min, u_max):
+                for v in range(v_min, v_max):
+                    d = np.asarray([(u - cx) / fx, (v - cy) / fy, 1.0])
+                    dirs.append(ori @ (d / np.linalg.norm(d)))
+            self.vis.rays(self.task_frame, pos, dirs, t_max)
+
         return ig
 
     def cost_fn(self, view):

src/active_grasp/rviz.py

@@ -3,11 +3,13 @@ import numpy as np
 from robot_helpers.ros.rviz import *
 from robot_helpers.spatial import Transform
 import vgn.rviz
-from vgn.utils import box_lines
+from vgn.utils import *
+
 
 cm = lambda s: tuple([float(1 - s), float(s), float(0)])
 red = [1.0, 0.0, 0.0]
 blue = [0, 0.6, 1.0]
+grey = [0.9, 0.9, 0.9]
 
 
 class Visualizer(vgn.rviz.Visualizer):
@@ -19,17 +21,26 @@ class Visualizer(vgn.rviz.Visualizer):
         marker = create_line_list_marker(frame, pose, scale, color, lines, "bbox")
         self.draw([marker])
 
-    def rays(self, frame, origin, directions, t_max=1.0):
+    def ig_views(self, frame, intrinsic, views, values):
-        lines = [[origin, origin + t_max * direction] for direction in directions]
+        vmin, vmax = min(values), max(values)
-        marker = create_line_list_marker(
+        scale = [0.002, 0.0, 0.0]
+        near, far = 0.0, 0.02
+        markers = []
+        for i, (view, value) in enumerate(zip(views, values)):
+            color = cm((value - vmin) / (vmax - vmin))
+            marker = create_view_marker(
                 frame,
-            Transform.identity(),
+                view,
-            [0.001, 0.0, 0.0],
+                scale,
-            [0.9, 0.9, 0.9],
+                color,
-            lines,
+                intrinsic,
-            "rays",
+                near,
+                far,
+                ns="ig_views",
+                id=i,
             )
-        self.draw([marker])
+            markers.append(marker)
+        self.draw(markers)
 
     def path(self, frame, poses):
         color = blue
@@ -67,25 +78,34 @@ class Visualizer(vgn.rviz.Visualizer):
         )
         self.draw([marker])
 
-    def views(self, frame, intrinsic, views, values):
+    def rays(self, frame, origin, directions, t_max=1.0):
-        vmin, vmax = min(values), max(values)
+        lines = [[origin, origin + t_max * direction] for direction in directions]
-        scale = [0.002, 0.0, 0.0]
+        marker = create_line_list_marker(
-        near, far = 0.0, 0.02
+            frame,
+            Transform.identity(),
+            [0.001, 0.0, 0.0],
+            grey,
+            lines,
+            "rays",
+        )
+        self.draw([marker])
+
+    def views(self, frame, intrinsic, views):
         markers = []
-        for i, (view, value) in enumerate(zip(views, values)):
+        for i, view in enumerate(views):
-            color = cm((value - vmin) / (vmax - vmin))
+            markers.append(
-            marker = create_view_marker(
+                create_view_marker(
                     frame,
                     view,
-                scale,
+                    [0.002, 0.0, 0.0],
-                color,
+                    blue,
                     intrinsic,
-                near,
+                    0.0,
-                far,
+                    0.02,
                     ns="views",
                     id=i,
                 )
-            markers.append(marker)
+            )
         self.draw(markers)