Fixed multi-view baselines

active_grasp/__init__.py

@@ -2,5 +2,6 @@ from .policy import register
 from .baselines import *
 
 register("initial-view", InitialView)
-register("front-view", FrontView)
 register("top-view", TopView)
+register("top-trajectory", TopTrajectory)
+register("circular-trajectory", CircularTrajectory)

active_grasp/baselines.py

@@ -1,7 +1,8 @@
 import numpy as np
+import rospy
+import scipy.interpolate
 
-
-from .policy import SingleViewPolicy
+from .policy import SingleViewPolicy, MultiViewPolicy
 from vgn.utils import look_at
 
 
@@ -11,22 +12,48 @@ class InitialView(SingleViewPolicy):
         super().update(img, extrinsic)
 
 
-class FrontView(SingleViewPolicy):
-    def activate(self, bbox):
-        super().activate(bbox)
-        l, theta = 0.25, np.deg2rad(30)
-        eye = np.r_[
-            self.center[0] - l * np.sin(theta),
-            self.center[1],
-            self.center[2] + l * np.cos(theta),
-        ]
-        up = np.r_[1.0, 0.0, 0.0]
-        self.target = look_at(eye, self.center, up)
-
-
 class TopView(SingleViewPolicy):
     def activate(self, bbox):
         super().activate(bbox)
-        eye = np.r_[self.center[:2], self.center[2] + 0.25]
+        eye = np.r_[self.center[:2], self.center[2] + 0.3]
         up = np.r_[1.0, 0.0, 0.0]
         self.target = look_at(eye, self.center, up)
+
+
+class TopTrajectory(MultiViewPolicy):
+    def activate(self, bbox):
+        super().activate(bbox)
+        eye = np.r_[self.center[:2], self.center[2] + 0.3]
+        up = np.r_[1.0, 0.0, 0.0]
+        self.target = look_at(eye, self.center, up)
+
+    def update(self, img, extrinsic):
+        self.integrate(img, extrinsic)
+        if np.linalg.norm(extrinsic.translation - self.target.translation) < 0.01:
+            self.done = True
+        else:
+            return self.target
+
+
+class CircularTrajectory(MultiViewPolicy):
+    def __init__(self, rate):
+        super().__init__(rate)
+        self.r = 0.1
+        self.h = 0.3
+        self.duration = 12.0
+        self.m = scipy.interpolate.interp1d([0, self.duration], [np.pi, 3.0 * np.pi])
+
+    def activate(self, bbox):
+        super().activate(bbox)
+        self.tic = rospy.Time.now()
+
+    def update(self, img, extrinsic):
+        self.integrate(img, extrinsic)
+        elapsed_time = (rospy.Time.now() - self.tic).to_sec()
+        if elapsed_time > self.duration:
+            self.done = True
+        else:
+            t = self.m(elapsed_time)
+            eye = self.center + np.r_[self.r * np.cos(t), self.r * np.sin(t), self.h]
+            up = np.r_[1.0, 0.0, 0.0]
+            return look_at(eye, self.center, up)

active_grasp/policy.py

@@ -39,6 +39,7 @@ class Policy:
         self.tsdf = UniformTSDFVolume(0.3, 40)
         self.vgn = VGN(Path(rospy.get_param("vgn/model")))
 
+        self.views = []
         self.best_grasp = None
         self.done = False
 
@@ -65,21 +66,19 @@ class Policy:
         return grasps[indices], scores[indices]
 
     def score_fn(self, grasp):
-        # return grasp.quality
-        return grasp.pose.translation[2]
+        return grasp.quality
+        # return grasp.pose.translation[2]
 
     def update(sekf, img, extrinsic):
         raise NotImplementedError
 
 
 class SingleViewPolicy(Policy):
-    """Plan grasps from a single view of the target object."""
-
     def update(self, img, extrinsic):
         error = extrinsic.translation - self.target.translation
 
         if np.linalg.norm(error) < 0.01:
-            self.views = [extrinsic.inv()]
+            self.views.append(extrinsic.inv())
 
             self.tsdf.integrate(img, self.intrinsic, extrinsic * self.T_base_task)
             tsdf_grid, voxel_size = self.tsdf.get_grid(), self.tsdf.voxel_size
@@ -93,7 +92,6 @@ class SingleViewPolicy(Policy):
             grasps, scores = self.sort_grasps(grasps)
 
             self.vis.grasps(self.base_frame, grasps, scores)
-            rospy.sleep(1.0)
 
             self.best_grasp = grasps[0] if len(grasps) > 0 else None
             self.done = True
@@ -101,6 +99,31 @@ class SingleViewPolicy(Policy):
             return self.target
 
 
+class MultiViewPolicy(Policy):
+    def __init__(self, rate):
+        super().__init__(rate)
+        self.preempt = True
+
+    def integrate(self, img, extrinsic):
+        self.views.append(extrinsic.inv())
+        self.tsdf.integrate(img, self.intrinsic, extrinsic * self.T_base_task)
+
+        self.vis.scene_cloud(self.task_frame, self.tsdf.get_scene_cloud())
+        self.vis.map_cloud(self.task_frame, self.tsdf.get_map_cloud())
+        self.vis.path(self.base_frame, self.views)
+
+        tsdf_grid, voxel_size = self.tsdf.get_grid(), self.tsdf.voxel_size
+        out = self.vgn.predict(tsdf_grid)
+
+        grasps = select_grid(voxel_size, out, threshold=0.95)
+        grasps, scores = self.sort_grasps(grasps)
+
+        if len(grasps) > 0:
+            self.best_grasp = grasps[0]
+
+        self.vis.grasps(self.base_frame, grasps, scores)
+
+
 registry = {}