diff --git a/cfg/active_grasp.rviz b/cfg/active_grasp.rviz index c41ae6c..2b0cc47 100644 --- a/cfg/active_grasp.rviz +++ b/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 - Y: -0.015723222866654396 - Z: 0.38234013319015503 + X: 0.4194673001766205 + Y: 0.033841054886579514 + 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: - Yaw: 5.183177471160889 + Yaw: 5.173177242279053 Saved: - Class: rviz/Orbit Distance: 1.2000000476837158 @@ -450,18 +426,14 @@ Visualization Manager: Target 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 - X: 1275 - Y: 61 + Width: 1343 + X: 1217 + Y: 57 diff --git a/cfg/active_grasp.yaml b/cfg/active_grasp.yaml index 7f7edc1..dc4d91b 100644 --- a/cfg/active_grasp.yaml +++ b/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: diff --git a/scripts/run.py b/scripts/run.py index 65328c6..d072912 100755 --- a/scripts/run.py +++ b/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() diff --git a/src/active_grasp/__init__.py b/src/active_grasp/__init__.py index 6372414..5833f63 100644 --- a/src/active_grasp/__init__.py +++ b/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) \ No newline at end of file +register("ap-multi-view", ActivePerceptionMultiViewPolicy) +register("ap-single-view", ActivePerceptionSingleViewPolicy) \ No newline at end of file diff --git a/src/active_grasp/active_perception_policy.py b/src/active_grasp/active_perception_policy.py index e19590a..542f130 100644 --- a/src/active_grasp/active_perception_policy.py +++ b/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 diff --git a/src/active_grasp/baselines.py b/src/active_grasp/baselines.py index 5d76b88..1661064 100644 --- a/src/active_grasp/baselines.py +++ b/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): diff --git a/src/active_grasp/controller.py b/src/active_grasp/controller.py index 43c3ed4..e59e3f3 100644 --- a/src/active_grasp/controller.py +++ b/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() - rospy.sleep(0.2) # Wait for a zero command to be sent to the robot. + # Wait for the robot to move to its desired camera pose + 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