single-view initialized

cfg/active_grasp.rviz

@@ -7,7 +7,7 @@ Panels:
         - /TF1/Frames1
         - /Markers1/Namespaces1
       Splitter Ratio: 0.6881287693977356
-    Tree Height: 371
+    Tree Height: 222
   - Class: rviz/Selection
     Name: Selection
   - Class: rviz/Tool Properties
@@ -25,7 +25,7 @@ Panels:
   - Class: rviz/Time
     Name: Time
     SyncMode: 0
-    SyncSource: ""
+    SyncSource: Image
 Preferences:
   PromptSaveOnExit: true
 Toolbars:
@@ -249,42 +249,6 @@ Visualization Manager:
       Use Fixed Frame: true
       Use rainbow: true
       Value: true
-    - Class: rviz/Image
-      Enabled: false
-      Image Topic: /camera/color/image_raw
-      Max Value: 1
-      Median window: 5
-      Min Value: 0
-      Name: ColorImage
-      Normalize Range: true
-      Queue Size: 2
-      Transport Hint: raw
-      Unreliable: false
-      Value: false
-    - Class: rviz/Image
-      Enabled: false
-      Image Topic: /camera/depth/image_rect_raw
-      Max Value: 1
-      Median window: 5
-      Min Value: 0
-      Name: DepthImage
-      Normalize Range: true
-      Queue Size: 2
-      Transport Hint: raw
-      Unreliable: false
-      Value: false
-    - Class: rviz/Image
-      Enabled: false
-      Image Topic: /camera/infra1/image_rect_raw
-      Max Value: 1
-      Median window: 5
-      Min Value: 0
-      Name: InfraImage
-      Normalize Range: true
-      Queue Size: 2
-      Transport Hint: raw
-      Unreliable: false
-      Value: false
     - Class: rviz/MarkerArray
       Enabled: true
       Marker Topic: visualization_marker_array
@@ -381,6 +345,18 @@ Visualization Manager:
       Use Fixed Frame: true
       Use rainbow: false
       Value: true
+    - Class: rviz/Image
+      Enabled: true
+      Image Topic: /camera/color/image_rect_raw
+      Max Value: 1
+      Median window: 5
+      Min Value: 0
+      Name: Image
+      Normalize Range: true
+      Queue Size: 2
+      Transport Hint: raw
+      Unreliable: false
+      Value: true
   Enabled: true
   Global Options:
     Background Color: 48; 48; 48
@@ -409,7 +385,7 @@ Visualization Manager:
   Views:
     Current:
       Class: rviz/Orbit
-      Distance: 2.024759531021118
+      Distance: 1.5453948974609375
       Enable Stereo Rendering:
         Stereo Eye Separation: 0.05999999865889549
         Stereo Focal Distance: 1
@@ -417,9 +393,9 @@ Visualization Manager:
         Value: false
       Field of View: 0.7799999713897705
       Focal Point:
-        X: 0.36431100964546204
+        X: 0.4194673001766205
-        Y: -0.015723222866654396
+        Y: 0.033841054886579514
-        Z: 0.38234013319015503
+        Z: 0.4190172851085663
       Focal Shape Fixed Size: false
       Focal Shape Size: 0.05000000074505806
       Invert Z Axis: false
@@ -427,7 +403,7 @@ Visualization Manager:
       Near Clip Distance: 0.009999999776482582
       Pitch: 0.4999997615814209
       Target Frame: <Fixed Frame>
-      Yaw: 5.183177471160889
+      Yaw: 5.173177242279053
     Saved:
       - Class: rviz/Orbit
         Distance: 1.2000000476837158
@@ -450,18 +426,14 @@ Visualization Manager:
         Target Frame: <Fixed Frame>
         Yaw: 1.0399999618530273
 Window Geometry:
-  ColorImage:
-    collapsed: true
-  DepthImage:
-    collapsed: false
   Displays:
-    collapsed: true
-  Height: 787
-  Hide Left Dock: true
-  Hide Right Dock: true
-  InfraImage:
     collapsed: false
-  QMainWindow State: 000000ff00000000fd00000004000000000000015900000506fc0200000016fb0000001200530065006c0065006300740069006f006e00000001e10000009b0000005c00fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000a3fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afc0000003d000005060000000000fffffffaffffffff0100000002fb000000100044006900730070006c0061007900730000000000000001f30000015600fffffffb000000140043006f006c006f00720049006d0061006700650000000000ffffffff0000009b00fffffffb000000140043006f006c006f00720049006d00610067006501000001e4000001bf0000000000000000fb00000014004400650070007400680049006d006100670065000000029a000001090000001600fffffffb000000140043006f006c006f00720049006d00610067006501000001d2000001920000000000000000fb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261fb0000000a004400650070007400680000000166000000b50000000000000000fb000000280020002d0020005400720061006a006500630074006f0072007900200053006c00690064006500720000000000ffffffff0000000000000000fb0000001c004d006f00740069006f006e0050006c0061006e006e0069006e006700000001f4000001af0000000000000000fb00000044004d006f00740069006f006e0050006c0061006e006e0069006e00670020002d0020005400720061006a006500630074006f0072007900200053006c00690064006500720000000000ffffffff0000000000000000fb0000001c00460069006c0074006500720065006400200043006c006f00750064000000029a000001050000000000000000fb000000140043006f006c006f00720049006d006100670065010000019f000002040000000000000000fb0000000a0049006d00610067006500000001bc0000008a0000000000000000fb0000000a0049006d006100670065000000021c000000ee0000000000000000fb0000000c00430061006d00650072006100000002db000000c80000000000000000fb000000140043006f006c006f00720049006d006100670065010000019f000002040000000000000000fb0000000a0049006d006100670065010000019c000002070000000000000000000000010000017b00000366fc0200000003fb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a00560069006500770073000000003d00000366000000a400fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000780000000a9fc0100000002fb000000140049006e0066007200610049006d00610067006502000000000000006b00000280000001e0fb0000000a00560069006500770073030000004e00000080000002e100000197000000030000056300000039fc0100000002fb0000000800540069006d00650000000000000005630000041800fffffffb0000000800540069006d00650100000000000004500000000000000000000003ff000002b900000004000000040000000800000008fc00000002000000020000000000000002000000010000000a0054006f006f006c00730100000000ffffffff0000000000000000
+  Height: 871
+  Hide Left Dock: false
+  Hide Right Dock: true
+  Image:
+    collapsed: false
+  QMainWindow State: 000000ff00000000fd0000000400000000000001db0000030dfc0200000017fb0000001200530065006c0065006300740069006f006e00000001e10000009b0000005c00fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000a3fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afc0000003d00000161000000c900fffffffa000000000100000002fb000000100044006900730070006c0061007900730100000000000001f30000015600fffffffb000000140043006f006c006f00720049006d0061006700650000000000ffffffff0000000000000000fb000000140043006f006c006f00720049006d00610067006501000001e4000001bf0000000000000000fb00000014004400650070007400680049006d006100670065000000029a000001090000000000000000fb000000140043006f006c006f00720049006d00610067006501000001d2000001920000000000000000fb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261fb0000000a004400650070007400680000000166000000b50000000000000000fb000000280020002d0020005400720061006a006500630074006f0072007900200053006c00690064006500720000000000ffffffff0000000000000000fb0000001c004d006f00740069006f006e0050006c0061006e006e0069006e006700000001f4000001af0000000000000000fb00000044004d006f00740069006f006e0050006c0061006e006e0069006e00670020002d0020005400720061006a006500630074006f0072007900200053006c00690064006500720000000000ffffffff0000000000000000fb0000001c00460069006c0074006500720065006400200043006c006f00750064000000029a000001050000000000000000fb000000140043006f006c006f00720049006d006100670065010000019f000002040000000000000000fb0000000a0049006d00610067006501000001a4000001a60000001600fffffffb0000000a0049006d006100670065000000021c000000ee0000000000000000fb0000000c00430061006d00650072006100000002db000000c80000000000000000fb000000140043006f006c006f00720049006d006100670065010000019f000002040000000000000000fb0000000a0049006d006100670065010000019c000002070000000000000000fb0000000a0049006d0061006700650100000192000001b80000000000000000000000010000017b00000366fc0200000003fb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a00560069006500770073000000003d00000366000000a400fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000780000000a9fc0100000002fb000000140049006e0066007200610049006d00610067006502000000000000006b00000280000001e0fb0000000a00560069006500770073030000004e00000080000002e100000197000000030000056300000039fc0100000002fb0000000800540069006d00650000000000000005630000041800fffffffb0000000800540069006d006501000000000000045000000000000000000000035e0000030d00000004000000040000000800000008fc00000002000000020000000000000002000000010000000a0054006f006f006c00730100000000ffffffff0000000000000000
   Selection:
     collapsed: false
   Time:
@@ -470,6 +442,6 @@ Window Geometry:
     collapsed: false
   Views:
     collapsed: true
-  Width: 1023
+  Width: 1343
-  X: 1275
+  X: 1217
-  Y: 61
+  Y: 57

cfg/active_grasp.yaml

@@ -10,8 +10,9 @@ hw:
 grasp_controller:
   base_frame_id: panda_link0
   ee_grasp_offset: [0.0, 0.0, -0.383, 0.924, 0.0, 0.0, 0.065]  # offset to panda_link8
-  control_rate: 30
+  control_rate: 100
   linear_vel: 0.05
+  move_arm_time: 6
   camera:
     frame_id: camera_depth_optical_frame
     info_topic: /camera/depth/camera_info
@@ -28,7 +29,7 @@ vgn:
   qual_threshold: 0.9
 
 policy:
-  rate: 4
+  rate: 2
   window_size: 12
 
 nbv_grasp:

scripts/run.py

@@ -6,6 +6,7 @@ import pandas as pd
 from pathlib import Path
 import rospy
 from tqdm import tqdm
+import torch
 
 from active_grasp.controller import *
 from active_grasp.policy import make, registry
@@ -14,6 +15,8 @@ from robot_helpers.ros import tf
 
 
 def main():
+    torch.cuda.empty_cache()
+
     rospy.init_node("grasp_controller")
     tf.init()
 

src/active_grasp/__init__.py

@@ -1,11 +1,12 @@
 from .policy import register
 from .baselines import *
 from .nbv import NextBestView
-from.active_perception_policy import ActivePerceptionMultiViewPolicy
+from.active_perception_policy import *
 
 register("initial-view", InitialView)
 register("top-view", TopView)
 register("top-trajectory", TopTrajectory)
 register("fixed-trajectory", FixedTrajectory)
 register("nbv", NextBestView)
-register("ap-multi-view", ActivePerceptionMultiViewPolicy)
+register("ap-multi-view", ActivePerceptionMultiViewPolicy)
+register("ap-single-view", ActivePerceptionSingleViewPolicy)

src/active_grasp/active_perception_policy.py

@@ -2,7 +2,7 @@ import itertools
 from numba import jit
 import numpy as np
 import rospy
-from .policy import MultiViewPolicy
+from .policy import MultiViewPolicy, SingleViewPolicy
 from .timer import Timer
 from .active_perception_demo import APInferenceEngine
 from robot_helpers.spatial import Transform
@@ -160,3 +160,143 @@ class ActivePerceptionMultiViewPolicy(MultiViewPolicy):
                 return True
         return False
 
+
+class ActivePerceptionSingleViewPolicy(SingleViewPolicy):
+    def __init__(self):
+        super().__init__()
+        self.max_views = rospy.get_param("ap_grasp/max_views")
+        self.ap_config_path = rospy.get_param("ap_grasp/ap_config_path")
+        self.ap_inference_engine = APInferenceEngine(self.ap_config_path)
+        self.pcdvis = RealTime3DVisualizer()
+        self.updated = False
+
+
+    def update(self, img, seg, target_id, x, q):
+        # Visualize scene cloud
+        self.vis_scene_cloud(img, x)
+
+        # When policy hasn't produced an available grasp
+        while(self.updated == False):
+            # Inference with our model
+            target_points, scene_points = self.depth_image_to_ap_input(img, seg, target_id)
+            ap_input = {'target_pts': target_points,
+                        'scene_pts': scene_points}
+            ap_output = self.ap_inference_engine.inference(ap_input)
+            delta_rot_6d = ap_output['estimated_delta_rot_6d']
+
+            current_cam_pose = torch.from_numpy(x.as_matrix()).float().to("cuda:0")
+            est_delta_rot_mat = self.rotation_6d_to_matrix_tensor_batch(delta_rot_6d)[0]
+            look_at_center = torch.from_numpy(self.bbox.center).float().to("cuda:0")
+            nbv_tensor = self.get_transformed_mat(current_cam_pose, 
+                                                    est_delta_rot_mat, 
+                                                    look_at_center)
+            nbv_numpy = nbv_tensor.cpu().numpy()
+            nbv_transform = Transform.from_matrix(nbv_numpy)
+            x_d = nbv_transform
+
+            # Check if this pose available
+            if(self.solve_cam_ik(self.q0, x_d)):
+                self.vis_cam_pose(x_d)
+                self.x_d = x_d
+                self.updated = True
+                print("Found an NBV!")
+                return
+        # Policy has produced an available grasp
+        if(self.updated == True):
+            self.generate_grasp(q)
+            self.done = True
+            
+    def deactivate(self):
+        self.vis.clear_ig_views()
+        self.updated = False
+
+    def vis_cam_pose(self, x):
+        # Integrate
+        self.views.append(x)
+        self.vis.path(self.base_frame, self.intrinsic, self.views)
+    
+    def vis_scene_cloud(self, img, x):
+        self.tsdf.integrate(img, self.intrinsic, x.inv() * self.T_base_task)
+        scene_cloud = self.tsdf.get_scene_cloud()
+        self.vis.scene_cloud(self.task_frame, np.asarray(scene_cloud.points))
+
+    def generate_grasp(self, q):
+        tsdf_grid = self.tsdf.get_grid()
+        out = self.vgn.predict(tsdf_grid)
+        self.vis.quality(self.task_frame, self.tsdf.voxel_size, out.qual, 0.9)
+        grasps, qualities = self.filter_grasps(out, q)
+
+        if len(grasps) > 0:
+            self.best_grasp, quality = self.select_best_grasp(grasps, qualities)
+            self.vis.grasp(self.base_frame, self.best_grasp, quality)
+        else:
+            self.best_grasp = None
+            self.vis.clear_grasp()
+
+
+    def select_best_grasp(self, grasps, qualities):
+        i = np.argmax(qualities)
+        return grasps[i], qualities[i]
+
+    def rotation_6d_to_matrix_tensor_batch(self, d6: torch.Tensor) -> torch.Tensor:
+        a1, a2 = d6[..., :3], d6[..., 3:]
+        b1 = F.normalize(a1, dim=-1)
+        b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
+        b2 = F.normalize(b2, dim=-1)
+        b3 = torch.cross(b1, b2, dim=-1)
+        return torch.stack((b1, b2, b3), dim=-2)
+    
+
+    def get_transformed_mat(self, src_mat, delta_rot, target_center_w):
+        src_rot = src_mat[:3, :3] 
+        dst_rot = src_rot @ delta_rot.T
+        dst_mat = torch.eye(4).to(dst_rot.device)
+        dst_mat[:3, :3] = dst_rot
+        distance = torch.norm(target_center_w - src_mat[:3, 3])
+        z_axis_camera = dst_rot[:3, 2].reshape(-1)
+        new_camera_position_w = target_center_w - distance * z_axis_camera
+        dst_mat[:3, 3] = new_camera_position_w
+        return dst_mat
+
+    def depth_image_to_ap_input(self, depth_img, seg_img, target_id):
+        target_points = []
+        scene_points = []
+
+        K = self.intrinsic.K
+        depth_shape = depth_img.shape
+        seg_shape = seg_img.shape
+        if(depth_shape == seg_shape):
+            img_shape = depth_shape
+        else:
+            print("Depth image shape and segmentation image shape are not the same")
+            return None
+        
+        # Convert depth image to 3D points
+        u_indices , v_indices = np.meshgrid(np.arange(img_shape[1]), np.arange(img_shape[0]))
+        x_factors = (u_indices - K[0, 2]) / K[0, 0]
+        y_factors = (v_indices - K[1, 2]) / K[1, 1]
+        z_mat = depth_img
+        x_mat = x_factors * z_mat
+        y_mat = y_factors * z_mat
+        for i in range(img_shape[0]):
+            for j in range(img_shape[1]):
+                seg_id = seg_img[i, j]
+                x = x_mat[i][j]
+                y = y_mat[i][j]
+                z = z_mat[i][j]
+                if(int(seg_id) == int(target_id)):
+                    # This pixel belongs to the target object to be grasped
+                    target_points.append([x,y,z])
+                else:
+                    # This pixel belongs to the scene
+                    scene_points.append([x,y,z])
+        
+        target_points = np.asarray(target_points)
+        target_points = target_points.reshape(1, target_points.shape[0], 3)
+        self.pcdvis.update_points(target_points)
+        target_points = torch.from_numpy(target_points).float().to("cuda:0")
+        scene_points = np.asarray(scene_points)
+        scene_points = scene_points.reshape(1, scene_points.shape[0], 3)
+        scene_points = torch.from_numpy(scene_points).float().to("cuda:0")
+        
+        return target_points, scene_points

src/active_grasp/baselines.py

@@ -6,7 +6,7 @@ from .policy import SingleViewPolicy, MultiViewPolicy, compute_error
 class InitialView(SingleViewPolicy):
     def update(self, img, seg, target_id, x, q):
         self.x_d = x
-        super().update(img, x, q)
+        super().update(img, seg, target_id, x, q)
 
 
 class TopView(SingleViewPolicy):

src/active_grasp/controller.py

@@ -36,6 +36,7 @@ class GraspController:
         self.min_z_dist = rospy.get_param("~camera/min_z_dist")
         self.control_rate = rospy.get_param("~control_rate")
         self.linear_vel = rospy.get_param("~linear_vel")
+        self.move_arm_time = rospy.get_param("~move_arm_time")
         self.policy_rate = rospy.get_param("policy/rate")
 
     def init_service_proxies(self):
@@ -111,8 +112,13 @@ class GraspController:
             depth_img, seg_image, pose, q = self.get_state()
             target_seg_id = self.get_target_id(TargetIDRequest()).id
             self.policy.update(depth_img, seg_image, target_seg_id, pose, q)
-            r.sleep()
+            # Wait for the robot to move to its desired camera pose
-        rospy.sleep(0.2)  # Wait for a zero command to be sent to the robot.
+            ticks = self.move_arm_time*self.policy_rate
+            for i in range(ticks):
+                r.sleep()
+
+        # Wait for a zero command to be sent to the robot.
+        rospy.sleep(0.2)
         self.policy.deactivate()
         timer.shutdown()
         return self.policy.best_grasp