from graspnetAPI import GraspNet, GraspNetEval
import open3d as o3d
import cv2

# GraspNetAPI example for loading grasp for a scene.
# change the graspnet_root path

####################################################################
graspnet_root = '/home/bam/graspnetAPI/graspnet' # ROOT PATH FOR GRASPNET
####################################################################
sceneId = 2
annId = 2
camera = 'realsense'

# initialize a GraspNet instance  
g = GraspNetEval(graspnet_root, camera=camera, split='train')

# load grasps of scene 1 with annotation id = 3, camera = kinect and fric_coef_thresh = 0.2
_6d_grasp = g.loadGrasp(sceneId = sceneId, annId = annId, format = '6d', camera = camera, fric_coef_thresh = 0.2)
print('6d grasp:\n{}'.format(_6d_grasp))

# visualize the grasps using open3d
geometries = []
geometries.append(g.loadScenePointCloud(sceneId = sceneId, annId = annId, camera = camera))
selected_grasps = _6d_grasp.random_sample(numGrasp = 20)
geometries += selected_grasps.to_open3d_geometry_list()
o3d.visualization.draw_geometries(geometries)

ann_id_list = [annId]
grasp_group_list = [selected_grasps]
# grasp_list_list, score_list_list, collision_list_list = g.eval_scene_all_grasps(sceneId, ann_id_list, grasp_group_list, vis=True, max_width = 0.1, use_cache = True)

g._eval_scene(sceneId, ann_id_list, grasp_group_list, vis=True, max_width = 0.1, use_cache = False)