diff --git a/grasp_tracking/demo.py b/grasp_tracking/demo.py new file mode 100644 index 0000000..6051c55 --- /dev/null +++ b/grasp_tracking/demo.py @@ -0,0 +1,105 @@ +import os +import argparse +import numpy as np +import open3d as o3d +from PIL import Image +from graspnetAPI import GraspGroup + +from tracker import AnyGraspTracker + +parser = argparse.ArgumentParser() +parser.add_argument('--checkpoint_path', required=True, help='Model checkpoint path') +parser.add_argument('--filter', type=str, default='oneeuro', help='filter to smooth grasp parameters(rotation, width, depth). [oneeuro/kalman/none]') + +parser.add_argument('--debug', action='store_true', help='Enable debug mode') +cfgs = parser.parse_args() + +class CameraInfo: + def __init__(self, width, height, fx, fy, cx, cy, scale): + self.width = width + self.height = height + self.fx = fx + self.fy = fy + self.cx = cx + self.cy = cy + self.scale = scale + +def create_point_cloud_from_depth_image(depth, camera, organized=True): + assert(depth.shape[0] == camera.height and depth.shape[1] == camera.width) + xmap = np.arange(camera.width) + ymap = np.arange(camera.height) + xmap, ymap = np.meshgrid(xmap, ymap) + points_z = depth / camera.scale + points_x = (xmap - camera.cx) * points_z / camera.fx + points_y = (ymap - camera.cy) * points_z / camera.fy + points = np.stack([points_x, points_y, points_z], axis=-1) + if not organized: + points = points.reshape([-1, 3]) + return points + +def get_data(data_dir, index): + # load image + colors = np.array(Image.open(os.path.join(data_dir, 'color_%03d.png'%index)), dtype=np.float32) / 255.0 + depths = np.load(os.path.join(data_dir, 'depth_%03d.npy'%index)) + + # set camera intrinsics + width, height = depths.shape[1], depths.shape[0] + fx, fy = 927.17, 927.37 + cx, cy = 651.32, 349.62 + scale = 1000.0 + camera = CameraInfo(width, height, fx, fy, cx, cy, scale) + + # get point cloud + points = create_point_cloud_from_depth_image(depths, camera) + mask = (points[:,:,2] > 0) & (points[:,:,2] < 1.5) + points = points[mask] + colors = colors[mask] + + return points, colors + +def demo(data_dir_list, indices): + # intialization + anygrasp_tracker = AnyGraspTracker(cfgs) + anygrasp_tracker.load_net() + + grasp_ids = [0] + vis = o3d.visualization.Visualizer() + vis.create_window(height=720, width=1280) + for i in range(len(indices)): + # get prediction + points, colors = get_data(data_dir_list, indices[i]) + target_gg, curr_gg, target_grasp_ids, corres_preds = anygrasp_tracker.update(points, colors, grasp_ids) + + if i == 0: + # select grasps on objects to track for the 1st frame + grasp_mask_x = ((curr_gg.translations[:,0]>-0.18) & (curr_gg.translations[:,0]<0.18)) + grasp_mask_y = ((curr_gg.translations[:,1]>-0.12) & (curr_gg.translations[:,1]<0.12)) + grasp_mask_z = ((curr_gg.translations[:,2]>0.35) & (curr_gg.translations[:,2]<0.55)) + grasp_ids = np.where(grasp_mask_x & grasp_mask_y & grasp_mask_z)[0][:30:6] + target_gg = curr_gg[grasp_ids] + else: + grasp_ids = target_grasp_ids + print(i, target_grasp_ids) + + # visualization + if cfgs.debug: + trans_mat = np.array([[1,0,0,0],[0,1,0,0],[0,0,-1,0],[0,0,0,1]]) + cloud = o3d.geometry.PointCloud() + cloud.points = o3d.utility.Vector3dVector(points) + cloud.colors = o3d.utility.Vector3dVector(colors) + cloud.transform(trans_mat) + grippers = target_gg.to_open3d_geometry_list() + for gripper in grippers: + gripper.transform(trans_mat) + vis.add_geometry(cloud) + for gripper in grippers: + vis.add_geometry(gripper) + vis.poll_events() + vis.remove_geometry(cloud) + for gripper in grippers: + vis.remove_geometry(gripper) + +if __name__ == "__main__": + data_dir = "example_data" + data_dir_list = [x for x in range(30)] + demo(data_dir, data_dir_list) \ No newline at end of file diff --git a/grasp_tracking/demo.sh b/grasp_tracking/demo.sh new file mode 100644 index 0000000..a8748ea --- /dev/null +++ b/grasp_tracking/demo.sh @@ -0,0 +1 @@ +python demo.py --checkpoint_path log/checkpoint_track_graspnetv1_4depth.tar --filter oneeuro \ No newline at end of file diff --git a/grasp_tracking/example_data/color_000.png b/grasp_tracking/example_data/color_000.png new file mode 100644 index 0000000..cc658b0 Binary files /dev/null and b/grasp_tracking/example_data/color_000.png differ diff --git a/grasp_tracking/example_data/color_001.png b/grasp_tracking/example_data/color_001.png new file mode 100644 index 0000000..23795f2 Binary files /dev/null and b/grasp_tracking/example_data/color_001.png differ diff --git a/grasp_tracking/example_data/color_002.png b/grasp_tracking/example_data/color_002.png new file mode 100644 index 0000000..18781d4 Binary files /dev/null and b/grasp_tracking/example_data/color_002.png differ diff --git a/grasp_tracking/example_data/color_003.png b/grasp_tracking/example_data/color_003.png new file mode 100644 index 0000000..69a2368 Binary files /dev/null and b/grasp_tracking/example_data/color_003.png differ diff --git a/grasp_tracking/example_data/color_004.png b/grasp_tracking/example_data/color_004.png new file mode 100644 index 0000000..bd827b3 Binary files /dev/null and b/grasp_tracking/example_data/color_004.png differ diff --git a/grasp_tracking/example_data/color_005.png b/grasp_tracking/example_data/color_005.png new file mode 100644 index 0000000..8d7c7c3 Binary files /dev/null and b/grasp_tracking/example_data/color_005.png differ diff --git a/grasp_tracking/example_data/color_006.png b/grasp_tracking/example_data/color_006.png new file mode 100644 index 0000000..6a9bb83 Binary files /dev/null and b/grasp_tracking/example_data/color_006.png differ diff --git a/grasp_tracking/example_data/color_007.png b/grasp_tracking/example_data/color_007.png new file mode 100644 index 0000000..18901b2 Binary files /dev/null and b/grasp_tracking/example_data/color_007.png differ diff --git a/grasp_tracking/example_data/color_008.png b/grasp_tracking/example_data/color_008.png new file mode 100644 index 0000000..993d02d Binary files /dev/null and b/grasp_tracking/example_data/color_008.png differ diff --git a/grasp_tracking/example_data/color_009.png b/grasp_tracking/example_data/color_009.png new file mode 100644 index 0000000..d05e1c0 Binary files /dev/null and b/grasp_tracking/example_data/color_009.png differ diff --git a/grasp_tracking/example_data/color_010.png b/grasp_tracking/example_data/color_010.png new file mode 100644 index 0000000..584df89 Binary files /dev/null and b/grasp_tracking/example_data/color_010.png differ diff --git a/grasp_tracking/example_data/color_011.png b/grasp_tracking/example_data/color_011.png new file mode 100644 index 0000000..6e810d2 Binary files /dev/null and b/grasp_tracking/example_data/color_011.png differ diff --git a/grasp_tracking/example_data/color_012.png b/grasp_tracking/example_data/color_012.png new file mode 100644 index 0000000..73e0a34 Binary files /dev/null and b/grasp_tracking/example_data/color_012.png differ diff --git a/grasp_tracking/example_data/color_013.png b/grasp_tracking/example_data/color_013.png new file mode 100644 index 0000000..fcfd28e Binary files /dev/null and b/grasp_tracking/example_data/color_013.png differ diff --git a/grasp_tracking/example_data/color_014.png b/grasp_tracking/example_data/color_014.png new file mode 100644 index 0000000..c6c1d08 Binary files /dev/null and b/grasp_tracking/example_data/color_014.png differ diff --git a/grasp_tracking/example_data/color_015.png b/grasp_tracking/example_data/color_015.png new file mode 100644 index 0000000..3a27a2f Binary files /dev/null and b/grasp_tracking/example_data/color_015.png differ diff --git a/grasp_tracking/example_data/color_016.png b/grasp_tracking/example_data/color_016.png new file mode 100644 index 0000000..58ca1f6 Binary files /dev/null and b/grasp_tracking/example_data/color_016.png differ diff --git a/grasp_tracking/example_data/color_017.png b/grasp_tracking/example_data/color_017.png new file mode 100644 index 0000000..7251e17 Binary files /dev/null and b/grasp_tracking/example_data/color_017.png differ diff --git a/grasp_tracking/example_data/color_018.png b/grasp_tracking/example_data/color_018.png new file mode 100644 index 0000000..00e04c4 Binary files /dev/null and b/grasp_tracking/example_data/color_018.png differ diff --git a/grasp_tracking/example_data/color_019.png b/grasp_tracking/example_data/color_019.png new file mode 100644 index 0000000..a4415ff Binary files /dev/null and b/grasp_tracking/example_data/color_019.png differ diff --git a/grasp_tracking/example_data/color_020.png b/grasp_tracking/example_data/color_020.png new file mode 100644 index 0000000..00a63a5 Binary files /dev/null and b/grasp_tracking/example_data/color_020.png differ diff --git a/grasp_tracking/example_data/color_021.png b/grasp_tracking/example_data/color_021.png new file mode 100644 index 0000000..8a0c6fe Binary files /dev/null and b/grasp_tracking/example_data/color_021.png differ diff --git a/grasp_tracking/example_data/color_022.png b/grasp_tracking/example_data/color_022.png new file mode 100644 index 0000000..be7fedb Binary files /dev/null and b/grasp_tracking/example_data/color_022.png differ diff --git a/grasp_tracking/example_data/color_023.png b/grasp_tracking/example_data/color_023.png new file mode 100644 index 0000000..736f41a Binary files /dev/null and b/grasp_tracking/example_data/color_023.png differ diff --git a/grasp_tracking/example_data/color_024.png b/grasp_tracking/example_data/color_024.png new file mode 100644 index 0000000..7e8e933 Binary files /dev/null and b/grasp_tracking/example_data/color_024.png differ diff --git a/grasp_tracking/example_data/color_025.png b/grasp_tracking/example_data/color_025.png new file mode 100644 index 0000000..5697119 Binary files /dev/null and b/grasp_tracking/example_data/color_025.png differ diff --git a/grasp_tracking/example_data/color_026.png b/grasp_tracking/example_data/color_026.png new file mode 100644 index 0000000..cd39df7 Binary files /dev/null and b/grasp_tracking/example_data/color_026.png differ diff --git a/grasp_tracking/example_data/color_027.png b/grasp_tracking/example_data/color_027.png new file mode 100644 index 0000000..dc6aca1 Binary files /dev/null and b/grasp_tracking/example_data/color_027.png differ diff --git a/grasp_tracking/example_data/color_028.png b/grasp_tracking/example_data/color_028.png new file mode 100644 index 0000000..20b6429 Binary files /dev/null and b/grasp_tracking/example_data/color_028.png differ diff --git a/grasp_tracking/example_data/color_029.png b/grasp_tracking/example_data/color_029.png new file mode 100644 index 0000000..22e328c Binary files /dev/null and b/grasp_tracking/example_data/color_029.png differ diff --git a/grasp_tracking/example_data/depth_000.npy b/grasp_tracking/example_data/depth_000.npy new file mode 100644 index 0000000..9b8be20 Binary files /dev/null and b/grasp_tracking/example_data/depth_000.npy differ diff --git a/grasp_tracking/example_data/depth_001.npy b/grasp_tracking/example_data/depth_001.npy new file mode 100644 index 0000000..f634c2d Binary files /dev/null and b/grasp_tracking/example_data/depth_001.npy differ diff --git a/grasp_tracking/example_data/depth_002.npy b/grasp_tracking/example_data/depth_002.npy new file mode 100644 index 0000000..d93c577 Binary files /dev/null and b/grasp_tracking/example_data/depth_002.npy differ diff --git a/grasp_tracking/example_data/depth_003.npy b/grasp_tracking/example_data/depth_003.npy new file mode 100644 index 0000000..89cc56c Binary files /dev/null and b/grasp_tracking/example_data/depth_003.npy differ diff --git a/grasp_tracking/example_data/depth_004.npy b/grasp_tracking/example_data/depth_004.npy new file mode 100644 index 0000000..0b966f2 Binary files /dev/null and b/grasp_tracking/example_data/depth_004.npy differ diff --git a/grasp_tracking/example_data/depth_005.npy b/grasp_tracking/example_data/depth_005.npy new file mode 100644 index 0000000..d29074e Binary files /dev/null and b/grasp_tracking/example_data/depth_005.npy differ diff --git a/grasp_tracking/example_data/depth_006.npy b/grasp_tracking/example_data/depth_006.npy new file mode 100644 index 0000000..f66b292 Binary files /dev/null and b/grasp_tracking/example_data/depth_006.npy differ diff --git a/grasp_tracking/example_data/depth_007.npy b/grasp_tracking/example_data/depth_007.npy new file mode 100644 index 0000000..1edbfd2 Binary files /dev/null and b/grasp_tracking/example_data/depth_007.npy differ diff --git a/grasp_tracking/example_data/depth_008.npy b/grasp_tracking/example_data/depth_008.npy new file mode 100644 index 0000000..decaa34 Binary files /dev/null and b/grasp_tracking/example_data/depth_008.npy differ diff --git a/grasp_tracking/example_data/depth_009.npy b/grasp_tracking/example_data/depth_009.npy new file mode 100644 index 0000000..3269a18 Binary files /dev/null and b/grasp_tracking/example_data/depth_009.npy differ diff --git a/grasp_tracking/example_data/depth_010.npy b/grasp_tracking/example_data/depth_010.npy new file mode 100644 index 0000000..5429eaa Binary files /dev/null and b/grasp_tracking/example_data/depth_010.npy differ diff --git a/grasp_tracking/example_data/depth_011.npy b/grasp_tracking/example_data/depth_011.npy new file mode 100644 index 0000000..cfd564b Binary files /dev/null and b/grasp_tracking/example_data/depth_011.npy differ diff --git a/grasp_tracking/example_data/depth_012.npy b/grasp_tracking/example_data/depth_012.npy new file mode 100644 index 0000000..8748c25 Binary files /dev/null and b/grasp_tracking/example_data/depth_012.npy differ diff --git a/grasp_tracking/example_data/depth_013.npy b/grasp_tracking/example_data/depth_013.npy new file mode 100644 index 0000000..0165407 Binary files /dev/null and b/grasp_tracking/example_data/depth_013.npy differ diff --git a/grasp_tracking/example_data/depth_014.npy b/grasp_tracking/example_data/depth_014.npy new file mode 100644 index 0000000..fa26258 Binary files /dev/null and b/grasp_tracking/example_data/depth_014.npy differ diff --git a/grasp_tracking/example_data/depth_015.npy b/grasp_tracking/example_data/depth_015.npy new file mode 100644 index 0000000..4b1fc84 Binary files /dev/null and b/grasp_tracking/example_data/depth_015.npy differ diff --git a/grasp_tracking/example_data/depth_016.npy b/grasp_tracking/example_data/depth_016.npy new file mode 100644 index 0000000..44162e9 Binary files /dev/null and b/grasp_tracking/example_data/depth_016.npy differ diff --git a/grasp_tracking/example_data/depth_017.npy b/grasp_tracking/example_data/depth_017.npy new file mode 100644 index 0000000..b89c069 Binary files /dev/null and b/grasp_tracking/example_data/depth_017.npy differ diff --git a/grasp_tracking/example_data/depth_018.npy b/grasp_tracking/example_data/depth_018.npy new file mode 100644 index 0000000..c28b03a Binary files /dev/null and b/grasp_tracking/example_data/depth_018.npy differ diff --git a/grasp_tracking/example_data/depth_019.npy b/grasp_tracking/example_data/depth_019.npy new file mode 100644 index 0000000..7029cb8 Binary files /dev/null and b/grasp_tracking/example_data/depth_019.npy differ diff --git a/grasp_tracking/example_data/depth_020.npy b/grasp_tracking/example_data/depth_020.npy new file mode 100644 index 0000000..a72447a Binary files /dev/null and b/grasp_tracking/example_data/depth_020.npy differ diff --git a/grasp_tracking/example_data/depth_021.npy b/grasp_tracking/example_data/depth_021.npy new file mode 100644 index 0000000..cd63635 Binary files /dev/null and b/grasp_tracking/example_data/depth_021.npy differ diff --git a/grasp_tracking/example_data/depth_022.npy b/grasp_tracking/example_data/depth_022.npy new file mode 100644 index 0000000..d12b277 Binary files /dev/null and b/grasp_tracking/example_data/depth_022.npy differ diff --git a/grasp_tracking/example_data/depth_023.npy b/grasp_tracking/example_data/depth_023.npy new file mode 100644 index 0000000..78acea7 Binary files /dev/null and b/grasp_tracking/example_data/depth_023.npy differ diff --git a/grasp_tracking/example_data/depth_024.npy b/grasp_tracking/example_data/depth_024.npy new file mode 100644 index 0000000..b295503 Binary files /dev/null and b/grasp_tracking/example_data/depth_024.npy differ diff --git a/grasp_tracking/example_data/depth_025.npy b/grasp_tracking/example_data/depth_025.npy new file mode 100644 index 0000000..2f2c291 Binary files /dev/null and b/grasp_tracking/example_data/depth_025.npy differ diff --git a/grasp_tracking/example_data/depth_026.npy b/grasp_tracking/example_data/depth_026.npy new file mode 100644 index 0000000..3938e26 Binary files /dev/null and b/grasp_tracking/example_data/depth_026.npy differ diff --git a/grasp_tracking/example_data/depth_027.npy b/grasp_tracking/example_data/depth_027.npy new file mode 100644 index 0000000..1df17cd Binary files /dev/null and b/grasp_tracking/example_data/depth_027.npy differ diff --git a/grasp_tracking/example_data/depth_028.npy b/grasp_tracking/example_data/depth_028.npy new file mode 100644 index 0000000..cbdaa65 Binary files /dev/null and b/grasp_tracking/example_data/depth_028.npy differ diff --git a/grasp_tracking/example_data/depth_029.npy b/grasp_tracking/example_data/depth_029.npy new file mode 100644 index 0000000..5330fdd Binary files /dev/null and b/grasp_tracking/example_data/depth_029.npy differ