added contactgraspnet class and testscript

Author: Ruthrash

robot_toolkit/robot_arm_algos/src/inference/contact_graspnet.py

@@ -1,12 +1,13 @@
 import numpy as np
 import cv2
+import timeit
 # Tensorflow
 
 import tensorflow.compat.v1 as tf
 tf.disable_eager_execution()
 physical_devices = tf.config.experimental.list_physical_devices('GPU')
 tf.config.experimental.set_memory_growth(physical_devices[0], True)
-
+import matplotlib.pyplot as plt
 import sys 
 sys.path.insert(0, '/root/git/contact_graspnet/contact_graspnet')
 
@@ -15,8 +16,9 @@
 from contact_grasp_estimator import GraspEstimator
 from visualization_utils import visualize_grasps, show_image
 
-from ..camera.camera import get_bbox_annotations, get_segmap_from_bbox
+from ..camera.camera import get_bbox_annotations, get_segmap_from_bbox_with_depth , get_segmap_from_bbox
 from ._grasp_predictor import GraspPredictor
+from ..logger import logger
 class ContactGraspNet(GraspPredictor):#(object)
     def __init__(self,
                 global_config,
@@ -47,11 +49,18 @@ def __init__(self,
 
         pass
 
-    def generate_grasps(self, rgbd_camera):
+    def generate_grasps(self, rgbd_camera, is_visualize_grasps = False, use_depth_for_seg = False):
         rgb_image, depth_image = rgbd_camera.get_current_rgbd_frames()
         bbox, _ = get_bbox_annotations(rgb_image)
-        segmap = get_segmap_from_bbox(rgb_image, bbox)
+        if use_depth_for_seg:
+            segmap = get_segmap_from_bbox_with_depth(rgb_image, depth_image, bbox)
+        else:            
+            segmap = get_segmap_from_bbox(rgb_image, bbox)        
+        # 
+        plt.imshow(segmap)
+        plt.show()    
         cam_k = rgbd_camera.camera_matrix
+        start = timeit.timeit()
         pc_full, pc_segments, pc_colors = self.grasp_estimator.extract_point_clouds(depth_image, 
                                                                                     cam_k, 
                                                                                     segmap=segmap,
@@ -63,10 +72,15 @@ def generate_grasps(self, rgbd_camera):
                                                                                             local_regions=self.local_regions, 
                                                                                             filter_grasps=self.filter_grasps, 
                                                                                             forward_passes=self.forward_passes)  
+        end = timeit.timeit()
+        logger.info(f"elapsed time {end-start}")
 
         show_image(rgb_image, segmap)
-        visualize_grasps(pc_full, pred_grasps_cam, scores, plot_opencv_cam=True, pc_colors=pc_colors)
-        
+        plt.imshow(segmap)
+        plt.show()
+        if is_visualize_grasps:
+            visualize_grasps(pc_full, pred_grasps_cam, scores, plot_opencv_cam=True, pc_colors=pc_colors)
+        return pred_grasps_cam, scores, contact_pts#, _
         # get_segmap 
         # rgb image 
         # depth image 

robot_toolkit/test_contact_graspnet.py

@@ -1,13 +1,35 @@
+import numpy as np
+import open3d as o3d
 from robot_arm_algos.src.camera.realsense_camera import RealSenseCamera
 from robot_arm_algos.src.inference.contact_graspnet import ContactGraspNet as cgn
 import robot_arm_algos.src.inference._contact_graspnet_config_utils as config_utils
 
+
+test_config = {
+    "use_depth" : False,
+    "checkpoint_dir" : "/root/git/scratchpad/scene_test_2048_bs3_hor_sigma_0025"
+}
+
 def main(): 
     rs_camera = RealSenseCamera()
-    checkpoint_dir = "/root/git/scratchpad/scene_test_2048_bs3_hor_sigma_0025"
-    global_config = config_utils.load_config(checkpoint_dir, batch_size=1, arg_configs=[])
-    cgn_ = cgn(global_config, checkpoint_dir )
-    cgn_.generate_grasps(rs_camera)
+    global_config = config_utils.load_config(test_config["checkpoint_dir"], batch_size=1, arg_configs=[])
+    cgn_ = cgn(global_config, test_config["checkpoint_dir"] )
+    grasps, scores, contact_points = cgn_.generate_grasps(rs_camera, use_depth_for_seg = test_config["use_depth"])
+    
+    # sorted_grasps = [x for _, x in sorted(zip(grasps[255.0], scores[255.0])) ]
+    sorted_grasps = grasps[255.0][np.argsort(scores[255.0])]
+    plot_grasps = sorted_grasps[:3]
+    plot_grasp_frames = []
+    for plot_grasp in plot_grasps:
+        frame_mesh = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.2).transform(plot_grasp)
+        plot_grasp_frames.append(frame_mesh)
+        
+    rgb_image, depth_image = rs_camera.get_current_rgbd_frames()
+    pcd = rs_camera.get_pointcloud_rgbd(rgb_image, depth_image)
+    plot_grasp_frames.append(pcd)
+    o3d.visualization.draw_geometries(plot_grasp_frames)
+
+    
 
 if __name__ == "__main__":
     main()