Only cast rays that intersect with the bbox

active_grasp/bbox.py

@@ -1,3 +1,4 @@
+import itertools
 import numpy as np
 
 import active_grasp.msg
@@ -9,6 +10,10 @@ class AABBox:
         self.min = bbox_min
         self.max = bbox_max
 
+    @property
+    def corners(self):
+        return list(itertools.product(*np.vstack((self.min, self.max)).T))
+
     def is_inside(self, p):
         return np.all(p > self.min) and np.all(p < self.max)
 

active_grasp/nbv.py

@@ -1,11 +1,22 @@
 import itertools
 import numpy as np
+from robot_helpers.perception import CameraIntrinsic
+from robot_helpers.spatial import Transform
 import rospy
 
 from .policy import BasePolicy
 from vgn.utils import look_at, spherical_to_cartesian
 
 
+class Ray:
+    def __init__(self, origin, direction):
+        self.o = origin
+        self.d = direction
+
+    def __call__(self, t):
+        return self.o + self.d * t
+
+
 class NextBestView(BasePolicy):
     def __init__(self, rate, filter_grasps):
         super().__init__(rate, filter_grasps)
@@ -19,6 +30,7 @@ class NextBestView(BasePolicy):
 
         # Generate viewpoints
         views = self.generate_viewpoints()
+        self.visualizer.views(views)
 
         # Evaluate viewpoints
         gains = [self.compute_ig(v) for v in views]
@@ -47,18 +59,25 @@ class NextBestView(BasePolicy):
         return views
 
     def compute_ig(self, view, downsample=20):
-        res_x = int(self.intrinsic.width / downsample)
-        res_y = int(self.intrinsic.height / downsample)
         fx = self.intrinsic.fx / downsample
         fy = self.intrinsic.fy / downsample
         cx = self.intrinsic.cx / downsample
         cy = self.intrinsic.cy / downsample
 
-        for x in range(res_x):
-            for y in range(res_y):
-                d = np.r_[(x - cx) / fx, (y - cy) / fy, 1.0]
-                d = d / np.linalg.norm(d)
-                d = view.rotation.apply(d)
+        T_cam_base = view.inv()
+        corners = np.array([T_cam_base.apply(p) for p in self.bbox.corners]).T
+        u = (fx * corners[0] / corners[2] + cx).round().astype(int)
+        v = (fy * corners[1] / corners[2] + cy).round().astype(int)
+        u_min, u_max = u.min(), u.max()
+        v_min, v_max = v.min(), v.max()
+
+        for u in range(u_min, u_max):
+            for v in range(v_min, v_max):
+                direction = np.r_[(u - cx) / fx, (v - cy) / fy, 1.0]
+                direction = direction / np.linalg.norm(direction)
+                direction = view.rotation.apply(direction)
+                ray = Ray(view.translation, direction)
+
         return 1.0
 
     def compute_cost(self, view):

active_grasp/visualization.py

@@ -13,6 +13,7 @@ class Visualizer:
         self.marker_pub = rospy.Publisher(topic, MarkerArray, queue_size=1)
         self.scene_cloud_pub = rospy.Publisher("scene_cloud", PointCloud2, queue_size=1)
         self.quality_pub = rospy.Publisher("quality", PointCloud2, queue_size=1)
+        self.views_pub = rospy.Publisher("views", PoseArray, queue_size=1)
         self.grasps_pub = rospy.Publisher("grasps", PoseArray, queue_size=1)
 
     def clear(self):
@@ -24,6 +25,9 @@ class Visualizer:
         msg.header.frame_id = self.frame
         self.grasps_pub.publish(msg)
 
+    def draw(self, markers):
+        self.marker_pub.publish(MarkerArray(markers=markers))
+
     def bbox(self, bbox):
         pose = Transform.translation((bbox.min + bbox.max) / 2.0)
         scale = bbox.max - bbox.min
@@ -31,6 +35,24 @@ class Visualizer:
         marker = create_cube_marker(self.frame, pose, scale, color, ns="bbox")
         self.draw([marker])
 
+    def grasps(self, grasps):
+        msg = PoseArray()
+        msg.header.frame_id = self.frame
+        msg.poses = [to_pose_msg(grasp.pose) for grasp in grasps]
+        self.grasps_pub.publish(msg)
+
+    def lines(self, lines):
+        marker = create_line_list_marker(
+            self.frame,
+            Transform.identity(),
+            [0.005, 0.0, 0.0],
+            [1.0, 0.0, 0.0],
+            lines,
+            "rays",
+            0,
+        )
+        self.draw([marker])
+
     def path(self, poses):
         color = np.r_[31, 119, 180] / 255.0
         points = [p.translation for p in poses]
@@ -54,20 +76,17 @@ class Visualizer:
         )
         self.draw([spheres, lines])
 
-    def draw(self, markers):
-        self.marker_pub.publish(MarkerArray(markers=markers))
-
-    def scene_cloud(self, frame, cloud):
-        msg = to_cloud_msg(frame, np.asarray(cloud.points))
-        self.scene_cloud_pub.publish(msg)
-
     def quality(self, frame, voxel_size, quality):
         points, values = grid_to_map_cloud(voxel_size, quality, threshold=0.8)
         msg = to_cloud_msg(frame, points, intensities=values)
         self.quality_pub.publish(msg)
 
-    def grasps(self, grasps):
+    def scene_cloud(self, frame, cloud):
+        msg = to_cloud_msg(frame, np.asarray(cloud.points))
+        self.scene_cloud_pub.publish(msg)
+
+    def views(self, views):
         msg = PoseArray()
         msg.header.frame_id = self.frame
-        msg.poses = [to_pose_msg(grasp.pose) for grasp in grasps]
-        self.grasps_pub.publish(msg)
+        msg.poses = [to_pose_msg(view) for view in views]
+        self.views_pub.publish(msg)