diff --git a/graspnetAPI/grasp.py b/graspnetAPI/grasp.py index c385443..f72c251 100644 --- a/graspnetAPI/grasp.py +++ b/graspnetAPI/grasp.py @@ -1,14 +1,15 @@ -__author__ = 'mhgou' +__author__ = 'mhgou, hsfang' import numpy as np import open3d as o3d import copy import cv2 -from .utils.utils import plot_gripper_pro_max, batch_rgbdxyz_2_rgbxy_depth, get_batch_key_points, batch_key_points_2_tuple, framexy_depth_2_xyz, batch_framexy_depth_2_xyz, center_depth, key_point_2_rotation, batch_center_depth, batch_framexy_depth_2_xyz, batch_key_point_2_rotation +from .utils.utils import plot_gripper_pro_max, plot_fric_reps, batch_rgbdxyz_2_rgbxy_depth, get_batch_key_points, batch_key_points_2_tuple, framexy_depth_2_xyz, batch_framexy_depth_2_xyz, center_depth, key_point_2_rotation, batch_center_depth, batch_framexy_depth_2_xyz, batch_key_point_2_rotation GRASP_ARRAY_LEN = 17 RECT_GRASP_ARRAY_LEN = 7 +FRIC_REP_ARRAY_LEN = 13+24*4+2 EPS = 1e-8 class Grasp(): @@ -31,7 +32,7 @@ def __init__(self, *args): raise TypeError('if only one arg is given, it must be np.ndarray.') elif len(args) == 7: score, width, height, depth, rotation_matrix, translation, object_id = args - self.grasp_array = np.concatenate([np.array((score, width, height, depth)),rotation_matrix.reshape(-1), translation, np.array((object_id)).reshape(-1)]).astype(np.float64) + self.grasp_array = np.concatenate([np.array((score, width, height, depth)),rotation_matrix.reshape(-1), translation, np.array([object_id]).reshape(-1)]).astype(np.float64) else: raise ValueError('only 1 or 7 arguments are accepted') @@ -1079,3 +1080,466 @@ def random_sample(self, numGrasp = 20): shuffled_rect_grasp_group = RectGraspGroup() shuffled_rect_grasp_group.rect_grasp_group_array = copy.deepcopy(shuffled_rect_grasp_group_array[:numGrasp]) return shuffled_rect_grasp_group + +class FricRep(): + def __init__(self, *args): + ''' + **Input:** + + - args can be a numpy array or tuple of the translation, view_rot, depth, fric_rep, height, object_id + + - the format of numpy array is [translation(3), view_rot(3x3), depth, fric_rep(48x2), height, object_id] + + - the length of the numpy array is 111. + ''' + if len(args) == 0: + self.fric_rep_array = np.array([0, 0, 0] + [1, 0, 0, 0, 1, 0, 0, 0, 1] + [0.02] + [0.02, 0.1]*48 + [0, 0.004, -1], dtype = np.float64) + elif len(args) == 1: + if type(args[0]) == np.ndarray: + self.fric_rep_array = copy.deepcopy(args[0]) + else: + raise TypeError('if only one arg is given, it must be np.ndarray.') + elif len(args) == 6: + translation, view_rot, depth, fric_rep, height, object_id = args + self.fric_rep_array = np.concatenate([translation, view_rot.reshape(-1), np.array([depth]),fric_rep.reshape(-1), np.array([height, object_id]).reshape(-1)]).astype(np.float64) + else: + raise ValueError('only 1 or 6 arguments are accepted') + + def __repr__(self): + return 'translation:{}\nview_rot:{}\ndepth:{}\nfric_rep:{}\nheight:{}, object id:{}'.format(self.translation, self.view_rot, self.depth, self.fric_rep, self.height, self.object_id) + + @property + def translation(self): + ''' + **Output:** + + - np.array of shape (3,) of the translation. + ''' + return self.fric_rep_array[0:3] + + @translation.setter + def translation(self, *args): + ''' + **Input:** + + - len(args) == 1: tuple of x, y, z + + - len(args) == 3: float of x, y, z + ''' + if len(args) == 1: + self.fric_rep_array[0:3] = np.array(args[0],dtype = np.float64) + elif len(args) == 3: + self.fric_rep_array[0:3] = np.array(args,dtype = np.float64) + + @property + def view_rot(self): + ''' + **Output:** + + - np.array of shape (9, ) of the view_rot. + ''' + return self.fric_rep_array[3:12].reshape((3,3)) + + @view_rot.setter + def view_rot(self, *args): + ''' + **Input:** + + - len(args) == 1: tuple of 9 + + - len(args) == 9: float of 9 + ''' + if len(args) == 1: + self.fric_rep_array[3:12] = np.array(args[0],dtype = np.float64).reshape(9) + elif len(args) == 3: + self.fric_rep_array[3:12] = np.array(args,dtype = np.float64) + + @property + def depth(self): + ''' + **Output:** + + - float of the depth. + ''' + return float(self.fric_rep_array[12]) + + @depth.setter + def depth(self, depth): + ''' + **input:** + + - float of the depth. + ''' + self.fric_rep_array[12] = depth + + @property + def fric_rep(self): + ''' + **Output:** + + - np.array of shape (48, 2) of the fric representation. + ''' + return self.fric_rep_array[13:13+48*2].reshape((48,2)) + + @fric_rep.setter + def fric_rep(self, *args): + ''' + **Input:** + + - len(args) == 1: tuple of fric representation + + - len(args) == 48*2: float of fric representation + ''' + if len(args) == 1: + self.fric_rep_array[13:13+48*2] = np.array(args[0],dtype = np.float64).reshape(48*2) + elif len(args) == 48*2: + self.fric_rep_array[13:13+48*2] = np.array(args,dtype = np.float64) + + @property + def height(self): + ''' + **Output:** + + - float of the height. + ''' + return float(self.fric_rep_array[13+48*2]) + + @height.setter + def height(self, height): + ''' + **input:** + + - float of the height. + ''' + self.fric_rep_array[13+48*2] = height + + @property + def object_id(self): + ''' + **Output:** + + - int of the object id that this grasp grasps + ''' + return self.fric_rep_array[13+48*2+1] + + @object_id.setter + def object_id(self, object_id): + ''' + **Input:** + + - int of the object_id. + ''' + self.fric_rep_array[13+48*2+1] = object_id + + def transform(self, T): + ''' + **Input:** + + - T: np.array of shape (4, 4) + + **Output:** + + - FricRep instance after transformation, the original FricRep will also be changed. + ''' + rotation = T[:3,:3] + translation = T[:3,3] + self.translation = np.dot(rotation, self.translation.reshape((3,1))).reshape(-1) + translation + self.view_rot = np.dot(rotation, self.view_rot) + return self + + def to_open3d_geometry(self, color=None): + ''' + **Input:** + + - color: optional, tuple of shape (3) denotes (r, g, b), e.g., (1,0,0) for red + + **Ouput:** + + - list of open3d.geometry.Geometry of the gripper. + ''' + return plot_fric_reps(self.translation, self.view_rot, self.depth, self.fric_rep) + +class FricRepGroup(): + def __init__(self, *args): + ''' + **Input:** + + - args can be (1) nothing (2) numpy array of grasp group array (3) str of the npy file. + ''' + if len(args) == 0: + self.fric_rep_group_array = np.zeros((0, FRIC_REP_ARRAY_LEN), dtype=np.float64) + elif len(args) == 1: + if isinstance(args[0], np.ndarray): + self.fric_rep_group_array = args[0] + elif isinstance(args[0], str): + self.fric_rep_group_array = np.load(args[0]) + else: + raise ValueError('args must be nothing, numpy array or string.') + else: + raise ValueError('args must be nothing, numpy array or string.') + + def __len__(self): + ''' + **Output:** + + - int of the length. + ''' + return len(self.fric_rep_group_array) + + def __repr__(self): + repr = '----------\FricRep Group, Number={}:\n'.format(self.__len__()) + if self.__len__() <= 6: + for fric_rep_array in self.fric_rep_group_array: + repr += FricRep(fric_rep_array).__repr__() + '\n' + else: + for i in range(3): + repr += FricRep(self.fric_rep_group_array[i]).__repr__() + '\n' + repr += '......\n' + for i in range(3): + repr += FricRep(self.fric_rep_group_array[-(3-i)]).__repr__() + '\n' + return repr + '----------' + + def __getitem__(self, index): + ''' + **Input:** + + - index: int, slice, list or np.ndarray. + + **Output:** + + - if index is int, return FricRep instance. + + - if index is slice, np.ndarray or list, return FricRepGroup instance. + ''' + if type(index) == int: + return FricRep(self.fric_rep_group_array[index]) + elif type(index) == slice: + fricrepgroup = FricRepGroup() + fricrepgroup.fric_rep_group_array = copy.deepcopy(self.fric_rep_group_array[index]) + return fricrepgroup + elif type(index) == np.ndarray: + return FricRepGroup(self.fric_rep_group_array[index]) + elif type(index) == list: + return FricRepGroup(self.fric_rep_group_array[index]) + else: + raise TypeError('unknown type "{}" for calling __getitem__ for FricRepGroup'.format(type(index))) + + + @property + def translations(self): + ''' + **Output:** + + - np.array of shape (-1,3) of the translations. + ''' + return self.fric_rep_group_array[:,0:3] + + @translations.setter + def translations(self, translations): + ''' + **Input:** + + - len(args) == 1: tuple of x, y, z + + - len(args) == 3: float of x, y, z + ''' + assert translations.shape == (len(self), 3) + self.fric_rep_group_array[:,0:3] = copy.deepcopy(translations) + + @property + def view_rots(self): + ''' + **Output:** + + - np.array of shape (-1, 9) of the view_rots. + ''' + return self.fric_rep_group_array[:,3:12].reshape((-1,3,3)) + + @view_rots.setter + def view_rots(self, view_rots): + ''' + **Input:** + + - view_rotation matrices of shape (-1,3,3) + ''' + assert view_rots.shape == (len(self), 3, 3) + self.fric_rep_group_array[:,3:12] = copy.deepcopy(view_rots.reshape((-1, 9))) + + @property + def depths(self): + ''' + **Output:** + + - float of the depths. + ''' + return self.fric_rep_group_array[:,12] + + @depths.setter + def depths(self, depth): + ''' + **input:** + + - float of the depths. + ''' + assert depths.size == len(self) + self.fric_rep_group_array[:,12] = copy.deepcopy(depths) + + @property + def fric_reps(self): + ''' + **Output:** + + - np.array of shape (-1, 48, 2) of the fric representations. + ''' + return self.fric_rep_group_array[:,13:13+48*2].reshape((-1,48,2)) + + @fric_reps.setter + def fric_reps(self, fric_reps): + ''' + **Input:** + + - fric_reps: matrices of fric representations + ''' + assert view_rots.shape == (len(self), 48, 2) + self.fric_rep_group_array[:,13:13+48*2] = copy.deepcopy(fric_reps.reshape((-1, 48*2))) + + @property + def heights(self): + ''' + **Output:** + + - float of the heights. + ''' + return self.fric_rep_group_array[:,13+48*2] + + @heights.setter + def heights(self, heights): + ''' + **input:** + + - float of the heights. + ''' + assert heights.size == len(self) + self.fric_rep_group_array[:,13+48*2] = copy.deepcopy(heights) + + @property + def object_ids(self): + ''' + **Output:** + + - int of the object id that this grasp grasps + ''' + return self.fric_rep_group_array[:,13+48*2+1] + + @object_ids.setter + def object_ids(self, object_ids): + ''' + **Input:** + + - int of the object_ids. + ''' + assert object_ids.size == len(self) + self.fric_rep_group_array[:,13+48*2+1] = copy.deepcopy(object_ids) + + def transform(self, T): + ''' + **Input:** + + - T: np.array of shape (4, 4) + + **Output:** + + - GraspGroup instance after transformation, the original GraspGroup will also be changed. + ''' + rotation = T[:3,:3] + translation = T[:3,3] + self.translations = np.dot(rotation, self.translations.T).T + translation # (-1, 3) + self.view_rots = np.matmul(rotation, self.view_rots).reshape((-1, 3, 3)) # (-1, 9) + return self + + def add(self, element): + ''' + **Input:** + + - element: FricRep instance or FricRepGroup instance. + ''' + if isinstance(element, FricRep): + self.fric_rep_group_array = np.concatenate((self.fric_rep_group_array, element.fric_rep_array.reshape((-1, FRIC_REP_ARRAY_LEN)))) + elif isinstance(element, FricRepGroup): + self.fric_rep_group_array = np.concatenate((self.fric_rep_group_array, element.fric_rep_group_array)) + else: + raise TypeError('Unknown type:{}'.format(element)) + return self + + def remove(self, index): + ''' + **Input:** + + - index: list of the index of fricrep + ''' + self.fric_rep_group_array = np.delete(self.fric_rep_group_array, index, axis = 0) + return self + + def from_npy(self, npy_file_path): + ''' + **Input:** + + - npy_file_path: string of the file path. + ''' + self.fric_rep_group_array = np.load(npy_file_path) + return self + + def save_npy(self, npy_file_path): + ''' + **Input:** + + - npy_file_path: string of the file path. + ''' + np.save(npy_file_path, self.fric_rep_group_array) + + def to_open3d_geometry_list(self): + ''' + **Output:** + + - list of open3d.geometry.Geometry of the grippers. + ''' + geometry = [] + for i in range(len(self.fric_rep_group_array)): + f = FricRep(self.fric_rep_group_array[i]) + geometry.append(f.to_open3d_geometry()) + return geometry + + def sort_by_non_empty(self, reverse = False): + ''' + **Input:** + + - reverse: bool of order, if False, from high to low, if True, from low to high. + + **Output:** + + - no output but sort the grasp group. + ''' + non_empty = np.sum(self.fric_rep_group_array[:,14:13+48*2:2]>0) + index = np.argsort(score) + if not reverse: + index = index[::-1] + self.fric_rep_group_array = self.fric_rep_group_array[index] + return self + + def random_sample(self, numFricRep = 20): + ''' + **Input:** + + - numFricRep: int of the number of sampled fric representation. + + **Output:** + + - FricRepGroup instance of sample fric reps. + ''' + if numFricRep > self.__len__(): + raise ValueError('Number of sampled fricrep should be no more than the total number of fricreps in the group') + shuffled_fric_rep_group_array = copy.deepcopy(self.fric_rep_group_array) + np.random.shuffle(shuffled_fric_rep_group_array) + shuffled_fricrep_group = FricRepGroup() + shuffled_fricrep_group.fric_rep_group_array = copy.deepcopy(shuffled_fric_rep_group_array[:numFricRep]) + return shuffled_fricrep_group \ No newline at end of file diff --git a/graspnetAPI/graspnet.py b/graspnetAPI/graspnet.py index ae8d7ed..9faa026 100755 --- a/graspnetAPI/graspnet.py +++ b/graspnetAPI/graspnet.py @@ -33,13 +33,17 @@ # loadScenePointCloud - Load point cloud constructed by the depth and color image. # loadWorkSpace - Load the workspace bounding box. # loadGraspLabels - Load grasp labels with the specified object ids. +# loadObjFricReps - Load grasp friction based representations with the specified object ids, # loadObjModels - Load object 3d mesh model with the specified object ids. # loadObjTrimesh - Load object 3d mesh in Trimesh format. # loadCollisionLabels - Load collision labels with the specified scene ids. # loadGrasp - Load grasp labels with the specified scene and annotation id. +# loadSceneFricReps - Load grasp friction based representations with the specified scene and annotation id. # loadData - Load data path with the specified data ids. # showObjGrasp - Save visualization of the grasp pose of specified object ids. # showSceneGrasp - Save visualization of the grasp pose of specified scene ids. +# showObjFricReps - Save visualization of the grasp friction based representations of specified object ids. +# showSceneFricReps - Save visualization of the grasp friction based representations of specified scene ids. # show6DPose - Save visualization of the 6d pose of specified scene ids, project obj models onto pointcloud # Throughout the API "ann"=annotation, "obj"=object, and "img"=image. @@ -55,7 +59,7 @@ import cv2 import trimesh -from .grasp import Grasp, GraspGroup, RectGrasp, RectGraspGroup, RECT_GRASP_ARRAY_LEN +from .grasp import Grasp, GraspGroup, FricRep, FricRepGroup, RectGrasp, RectGraspGroup, RECT_GRASP_ARRAY_LEN from .utils.utils import transform_points, parse_posevector from .utils.xmlhandler import xmlReader @@ -85,19 +89,19 @@ def __init__(self, root, camera='kinect', split='train'): self.split = split self.collisionLabels = {} - if split == 'all': - self.sceneIds = list(range(TOTAL_SCENE_NUM)) - elif split == 'train': - self.sceneIds = list(range(100)) - elif split == 'test': - self.sceneIds = list(range(100, 190)) - elif split == 'test_seen': - self.sceneIds = list(range(100, 130)) - elif split == 'test_similar': - self.sceneIds = list(range(130, 160)) - elif split == 'test_novel': - self.sceneIds = list(range(160, 190)) - + # if split == 'all': + # self.sceneIds = list(range(TOTAL_SCENE_NUM)) + # elif split == 'train': + # self.sceneIds = list(range(100)) + # elif split == 'test': + # self.sceneIds = list(range(100, 190)) + # elif split == 'test_seen': + # self.sceneIds = list(range(100, 130)) + # elif split == 'test_similar': + # self.sceneIds = list(range(130, 160)) + # elif split == 'test_novel': + # self.sceneIds = list(range(160, 190)) + self.sceneIds = list([180]) self.rgbPath = [] self.depthPath = [] self.segLabelPath = [] @@ -278,6 +282,27 @@ def loadGraspLabels(self, objIds=None): file = np.load(os.path.join(self.root, 'grasp_label', '{}_labels.npz'.format(str(i).zfill(3)))) graspLabels[i] = (file['points'].astype(np.float32), file['offsets'].astype(np.float32), file['scores'].astype(np.float32)) return graspLabels + + def loadObjFricReps(self, objIds=None): + ''' + **Input:** + + - objIds: int or list of int of the object ids. + + **Output:** + + - a dict of friction representations of each object. + ''' + # load object-level grasp labels of the given obj ids + objIds = self.objIds if objIds is None else objIds + assert _isArrayLike(objIds) or isinstance(objIds, int), 'objIds must be an integer or a list/numpy array of integers' + objIds = objIds if _isArrayLike(objIds) else [objIds] + ObjFricReps = {} + for i in tqdm(objIds, desc='Loading object friction representations...'): + file = np.load(os.path.join(self.root, 'fric_rep', '{}_labels_small.npz'.format(str(i).zfill(3)))) + ObjFricReps[i] = (file['points'].astype(np.float32), file['data'].astype(np.float32)) + return ObjFricReps + def loadObjModels(self, objIds=None): ''' @@ -347,6 +372,28 @@ def loadCollisionLabels(self, sceneIds=None): collisionLabels['scene_'+str(sid).zfill(4)] = collisionLabel return collisionLabels + def loadFricCollisionLabels(self, sceneIds=None): + ''' + **Input:** + + - sceneIds: int or list of int of the scene ids. + + **Output:** + + - dict of the collision labels. + ''' + sceneIds = self.sceneIds if sceneIds is None else sceneIds + assert _isArrayLike(sceneIds) or isinstance(sceneIds, int), 'sceneIds must be an integer or a list/numpy array of integers' + sceneIds = sceneIds if _isArrayLike(sceneIds) else [sceneIds] + collisionLabels = {} + for sid in tqdm(sceneIds, desc='Loading fric_rep collision labels...'): + labels = np.load(os.path.join(self.root, 'fric_collision_label','scene_'+str(sid).zfill(4), 'collision_labels.npz')) + collisionLabel = [] + for j in range(len(labels)): + collisionLabel.append(labels['arr_{}'.format(j)]) + collisionLabels['scene_'+str(sid).zfill(4)] = collisionLabel + return collisionLabels + def loadRGB(self, sceneId, camera, annId): ''' **Input:** @@ -662,6 +709,132 @@ def loadGrasp(self, sceneId, annId=0, format = '6d', camera='kinect', grasp_labe rect_grasps = RectGraspGroup(os.path.join(self.root,'scenes','scene_%04d' % sceneId,camera,'rect','%04d.npy' % annId)) return rect_grasps + def loadSceneFricReps(self, sceneId, annId=0, camera='realsense', fric_reps = None, fric_collision_labels = None, fric_coef_thresh=0.5): + ''' + TODO@fang-haoshu + **Input:** + + - sceneId: int of scene id. + + - annId: int of annotation id. + + - camera: string of camera type, 'kinect' or 'realsense'. + + - fric_reps: dict of object friction based representations. Call self.loadObjFricReps if not given. + + - collision_labels: dict of collision labels. Call self.loadFricCollisionLabels if not given. + + - fric_coef_thresh: float of the frcition coefficient threshold of the grasp. + + **ATTENTION** + + the LOWER the friction coefficient is (except 0), the better the contact is. + + **Output:** + + - return a FricRepGroup instance. + ''' + from .utils.xmlhandler import xmlReader + from .utils.utils import get_obj_pose_list, generate_views, transform_points + from .utils.rotation import batch_viewpoint_params_to_matrix + + camera_poses = np.load(os.path.join(self.root,'scenes','scene_%04d' %(sceneId,),camera, 'camera_poses.npy')) + camera_pose = camera_poses[annId] + scene_reader = xmlReader(os.path.join(self.root,'scenes','scene_%04d' %(sceneId,),camera,'annotations','%04d.xml' %(annId,))) + pose_vectors = scene_reader.getposevectorlist() + + obj_list,pose_list = get_obj_pose_list(camera_pose,pose_vectors) + if fric_reps is None: + print('warning: fric_reps are not given, calling self.loadObjFricReps to retrieve them') + fric_reps = self.loadObjFricReps(objIds = obj_list) + if fric_collision_labels is None: + print('warning: collision_labels are not given, calling self.loadFricCollisionLabels to retrieve them') + fric_collision_labels = self.loadFricCollisionLabels(sceneId) + + num_views, num_angles, num_depths = 300, 48, 5 + template_views = generate_views(num_views) + template_views = template_views[np.newaxis, :, np.newaxis, np.newaxis, :] + template_views = np.tile(template_views, [1, 1, 1, num_depths, 1]) + template_view_inds = np.arange(300) + template_view_inds = template_view_inds[np.newaxis, :, np.newaxis, np.newaxis] + template_view_inds = np.tile(template_view_inds, [1, 1, 1, num_depths]) + collision_dump = fric_collision_labels['scene_'+str(sceneId).zfill(4)] + + # grasp = dict() + fric_rep_group = FricRepGroup() + for i, (obj_idx, trans) in enumerate(zip(obj_list, pose_list)): + sampled_points, datas = fric_reps[obj_idx] + collision = collision_dump[i] + print('00000\n') + point_inds = np.arange(sampled_points.shape[0]) + + angle_inds = np.arange(48) + depth_inds = np.arange(5) + + num_points = len(point_inds) + target_points = sampled_points[:, np.newaxis, np.newaxis, :] + target_points = np.tile(target_points, [1, num_views, num_depths, 1]).reshape(-1,3) + point_inds = point_inds[:, np.newaxis, np.newaxis] + point_inds = np.tile(point_inds, [1, num_views, num_depths]).reshape(-1,1) + views = np.tile(template_views, [num_points, 1, 1, 1, 1]).reshape(-1,3) + view_inds = np.tile(template_view_inds, [num_points, 1, 1, 1]).reshape(-1,1) + depths = np.array([0.005, 0.01, 0.02, 0.03, 0.04]).repeat(num_points*num_views) + depth_inds = np.array(depth_inds).repeat(num_points*num_views).reshape(-1,1) + circular_reps = np.swapaxes(datas,2,3).reshape(-1,48,2) + print('11111\n') + collision = np.swapaxes(collision,2,3).reshape(-1,48) + print('22222\n') + grasp_score_min = np.minimum(circular_reps[:,:24,1], circular_reps[:,24:,1]) + grasp_score_max = np.maximum(circular_reps[:,:24,1], circular_reps[:,24:,1]) + sum_mu1, sum_mu2 = np.sum(circular_reps[:,:24,1], axis=1), np.sum(circular_reps[:,24:,1], axis=1) + + mask1 = np.any((grasp_score_min>0) & (grasp_score_max fric_coef_thresh*24) | (sum_mu1 <= 1) | (sum_mu2 > fric_coef_thresh*24) | (sum_mu2 <= 1)) + mask = mask1 & mask2 + target_points = target_points[mask] + point_inds = point_inds[mask] + view_inds = view_inds[mask] + depth_inds = depth_inds[mask] + views = views[mask] + depths = depths[mask] + circular_reps = circular_reps[mask] + collision = collision[mask] + print('33333\n') + circular_reps[collision.astype("bool"),1] = 2 + print('44444\n') + grasp_score_min = np.minimum(circular_reps[:,:24,1], circular_reps[:,24:,1]) + grasp_score_max = np.maximum(circular_reps[:,:24,1], circular_reps[:,24:,1]) + sum_mu1, sum_mu2 = np.sum(circular_reps[:,:24,1], axis=1), np.sum(circular_reps[:,24:,1], axis=1) + + mask1 = np.any((grasp_score_min>0) & (grasp_score_max fric_coef_thresh*48) | (sum_mu1 <= 1) | (sum_mu2 > fric_coef_thresh*48) | (sum_mu2 <= 1)) + mask = mask1 & mask2 + target_points = target_points[mask] + target_points = transform_points(target_points, trans) + target_points = transform_points(target_points, np.linalg.inv(camera_pose)) + point_inds = point_inds[mask] + view_inds = view_inds[mask] + depth_inds = depth_inds[mask] + views = views[mask] + depths = depths[mask] + circular_reps = circular_reps[mask] + print('55555\n') + Rs = batch_viewpoint_params_to_matrix(-views, np.zeros(len(views))) + Rs = np.matmul(trans[np.newaxis, :3, :3], Rs) + Rs = np.matmul(np.linalg.inv(camera_pose)[np.newaxis,:3,:3], Rs) + + num_fricrep = depths.shape[0] + circular_reps = circular_reps.reshape(-1,48*2) + heights = GRASP_HEIGHT * np.ones((num_fricrep,1)) + depths = depths.reshape(-1,1) + rotations = Rs.reshape((-1,9)) + object_ids = point_inds*100000 + view_inds*10 + depth_inds#obj_idx * np.ones((num_fricrep,1), dtype=np.int32) + + obj_fric_rep_array = np.hstack([target_points, rotations, depths, circular_reps, heights, object_ids]).astype(np.float32) + + fric_rep_group.fric_rep_group_array = np.concatenate((fric_rep_group.fric_rep_group_array, obj_fric_rep_array)) + return fric_rep_group + def loadData(self, ids=None, *extargs): ''' **Input:** @@ -768,6 +941,64 @@ def showSceneGrasp(self, sceneId, camera = 'kinect', annId = 0, format = '6d', n cv2.waitKey(0) cv2.destroyAllWindows() + def showObjFricReps(self, objIds=[], numFricRep=10, th=0.5, saveFolder='save_fig', show=False): + ''' + **Input:** + + - objIds: int of list of objects ids. + + - numFricRep: how many fric-reps to show in the image. + + - th: threshold of the coefficient of friction. + + - saveFolder: string of the path to save the rendered image. + + - show: bool of whether to show the image. + + **Output:** + + - No output but save the rendered image and maybe show it. + ''' + from .utils.vis import visObjFricReps + objIds = objIds if _isArrayLike(objIds) else [objIds] + if len(objIds) == 0: + print('You need to specify object ids.') + return 0 + + if not os.path.exists(saveFolder): + os.mkdir(saveFolder) + for obj_id in objIds: + visObjFricReps(self.root, obj_id, num_fric_rep=numFricRep, th=th, save_folder=saveFolder, show=show) + + def showSceneFricReps(self, sceneId, camera = 'kinect', annId = 0, numFricRep = 10, fric_reps = None, show_object = True, coef_fric_thresh = 0.5): + ''' + **Input:** + + - sceneId: int of the scene index. + + - camera: string of the camera type, 'realsense' or 'kinect'. + + - annId: int of the annotation index. + + - numFricRep: int of the displayed fric-reps number, reps will be randomly sampled. + + - fric_reps: fric representations + + - coef_fric_thresh: float of the friction coefficient of rep. + ''' + geometries = [] + SceneFricRep = self.loadSceneFricReps(sceneId = sceneId, annId = annId, camera = camera, fric_reps = fric_reps, fric_coef_thresh = coef_fric_thresh) + SceneFricRep = SceneFricRep.random_sample(numFricRep = numFricRep) + scenePCD = self.loadScenePointCloud(sceneId = sceneId, camera = camera, annId = annId, align = False) + geometries.append(scenePCD) + geometries += SceneFricRep.to_open3d_geometry_list() + print(int(SceneFricRep.object_ids[0]/100000), int((SceneFricRep.object_ids[0]%100000)/10), int((SceneFricRep.object_ids[0]%10))) + print(SceneFricRep.object_ids[0]) + if show_object: + objectPCD = self.loadSceneModel(sceneId = sceneId, camera = camera, annId = annId, align = False) + geometries += objectPCD + o3d.visualization.draw_geometries(geometries) + def show6DPose(self, sceneIds, saveFolder='save_fig', show=False, perObj=False): ''' **Input:** diff --git a/graspnetAPI/utils/utils.py b/graspnetAPI/utils/utils.py index fae9f77..3d6836f 100755 --- a/graspnetAPI/utils/utils.py +++ b/graspnetAPI/utils/utils.py @@ -354,6 +354,15 @@ def get_model_grasps(datapath): collision = label['collision'] return points, offsets, scores, collision +def get_model_fricreps(datapath): + ''' Author: chenxi-wang, h.s-fang + Load grasp labels from .npz files. + ''' + label = np.load(datapath) + points = label['points'] + data = label['data'] + return points, data + def parse_posevector(posevector): ''' Author: chenxi-wang Decode posevector to object id and transformation matrix. @@ -452,7 +461,7 @@ def plot_axis(R,center,length,grid_size = 0.01): def plot_gripper_pro_max(center, R, width, depth, score=1, color=None): ''' - Author: chenxi-wang + Author: chenxi-wang, h.s-fang **Input:** @@ -462,8 +471,12 @@ def plot_gripper_pro_max(center, R, width, depth, score=1, color=None): - width: float, gripper width + - depth: appraoch depth of gripper + - score: float, grasp quality score + - color: gripper color, by default red for high score and blue for low score + **Output:** - open3d.geometry.TriangleMesh @@ -521,6 +534,125 @@ def plot_gripper_pro_max(center, R, width, depth, score=1, color=None): gripper.vertex_colors = o3d.utility.Vector3dVector(colors) return gripper +def plot_fric_reps(center, view_rot, depth, fric_reps): + ''' + Author: h.s-fang, chenxi-wang + + **Input:** + + - center: numpy array of (3,), target point as gripper center + + - view_rot: approach view rotation of (3,3) on the object surface + + - depth: appraoch depth of gripper + + - fric_reps: array of (48, 2): 48x(distance, mu), distances to the object surface and the frictions + + **Output:** + + - open3d.geometry.TriangleMesh + ''' + x, y, z = center + height=0.004 + finger_width = 0.004 + tail_length = 0.04 + depth_base = 0.02 + + all_d1, all_mu1, all_d2, all_mu2 = fric_reps[:24,0], fric_reps[:24,1], fric_reps[24:,0], fric_reps[24:,1] + + + color1_r = (2-np.clip(all_mu1,0,2))/2 + 0.05 # red for high score + color1_g = 0.05*np.ones(all_mu1.shape) + color1_b = 0.05*np.ones(all_mu1.shape) + + color2_r = 0.05*np.ones(all_mu2.shape) + color2_g = (2-np.clip(all_mu2,0,2))/2 + 0.05 # green for high score + color2_b = 0.05*np.ones(all_mu2.shape) + + angle_inds = np.arange(24) + angles = ((angle_inds-12)/24)*np.pi + + + vertices = [] + triangles = [] + colors = [] + count = 0 + for i in range(24): + R1 = np.array([[1, 0, 0], + [0, np.cos(angles[i]), -np.sin(angles[i])], + [0, np.sin(angles[i]), np.cos(angles[i])]]) + rotation = view_rot.dot(R1) + if all_mu1[i] > 0: + right = create_mesh_box(depth+depth_base+finger_width, finger_width, height) + right_points = np.array(right.vertices) + right_triangles = np.array(right.triangles) + 8*count + count += 1 + right_points[:,0] -= depth_base + finger_width + right_points[:,1] += all_d1[i] + right_points[:,2] -= height/2 + colors_right = np.array([ [color1_r[i],color1_g[i],color1_b[i]] for _ in range(len(right.vertices))]) + vertices.append(np.dot(rotation, right_points.T).T + center) + triangles.append(right_triangles) + colors.append(colors_right) + + bottom1 = create_mesh_box(finger_width, all_d1[i], height) + bottom1_points = np.array(bottom1.vertices) + bottom1_triangles = np.array(bottom1.triangles) + 8*count + count += 1 + bottom1_points[:,0] -= finger_width + depth_base + # bottom1_points[:,1] -= all_d1[i] + bottom1_points[:,2] -= height/2 + colors_bottom1 = np.array([ [color1_r[i],color1_g[i],color1_b[i]] for _ in range(len(bottom1.vertices))]) + vertices.append(np.dot(rotation, bottom1_points.T).T + center) + triangles.append(bottom1_triangles) + colors.append(colors_bottom1) + + + if all_mu2[i] > 0: + left = create_mesh_box(depth+depth_base+finger_width, finger_width, height) + left_points = np.array(left.vertices) + left_triangles = np.array(left.triangles) + 8*count + count += 1 + left_points[:,0] -= depth_base + finger_width + left_points[:,1] -= all_d2[i] + finger_width + left_points[:,2] -= height/2 + colors_left = np.array([ [color2_r[i],color2_g[i],color2_b[i]] for _ in range(len(left.vertices))]) + vertices.append(np.dot(rotation, left_points.T).T + center) + triangles.append(left_triangles) + colors.append(colors_left) + + bottom2 = create_mesh_box(finger_width, all_d2[i], height) + bottom2_points = np.array(bottom2.vertices) + bottom2_triangles = np.array(bottom2.triangles) + 8*count + count += 1 + bottom2_points[:,0] -= finger_width + depth_base + bottom2_points[:,1] -= all_d2[i] + bottom2_points[:,2] -= height/2 + colors_bottom2 = np.array([ [color2_r[i],color2_g[i],color2_b[i]] for _ in range(len(bottom2.vertices))]) + vertices.append(np.dot(rotation, bottom2_points.T).T + center) + triangles.append(bottom2_triangles) + colors.append(colors_bottom2) + + tail = create_mesh_box(tail_length, finger_width, height) + tail_points = np.array(tail.vertices) + tail_triangles = np.array(tail.triangles) + 8*count + tail_points[:,0] -= tail_length + finger_width + depth_base + tail_points[:,1] -= finger_width / 2 + tail_points[:,2] -= height/2 + colors_tail = np.array([ [0.3,0.3,0.3] for _ in range(len(tail.vertices))]) + vertices.append(np.dot(rotation, tail_points.T).T + center) + triangles.append(tail_triangles) + colors.append(colors_tail) + + vertices_np = np.concatenate(vertices, axis=0) + triangles_np = np.concatenate(triangles, axis=0) + colors_np = np.concatenate(colors, axis=0) + + rep = o3d.geometry.TriangleMesh() + rep.vertices = o3d.utility.Vector3dVector(vertices_np) + rep.triangles = o3d.utility.Vector3iVector(triangles_np) + rep.vertex_colors = o3d.utility.Vector3dVector(colors_np) + return rep def find_scene_by_model_id(dataset_root, model_id_list): picked_scene_names = [] diff --git a/graspnetAPI/utils/vis.py b/graspnetAPI/utils/vis.py index 62887f1..303c6b5 100644 --- a/graspnetAPI/utils/vis.py +++ b/graspnetAPI/utils/vis.py @@ -3,7 +3,7 @@ import numpy as np import open3d as o3d from transforms3d.euler import euler2mat, quat2mat -from .utils import generate_scene_model, generate_scene_pointcloud, generate_views, get_model_grasps, plot_gripper_pro_max, transform_points +from .utils import generate_scene_model, generate_scene_pointcloud, generate_views, get_model_grasps, get_model_fricreps, plot_gripper_pro_max, plot_fric_reps, transform_points from .rotation import viewpoint_params_to_matrix, batch_viewpoint_params_to_matrix def create_table_cloud(width, height, depth, dx=0, dy=0, dz=0, grid_size=0.01): @@ -264,7 +264,7 @@ def visObjGrasp(dataset_root, obj_idx, num_grasp=10, th=0.5, max_width=0.08, sav plyfile = os.path.join(dataset_root, 'models', '%03d'%obj_idx, 'nontextured.ply') model = o3d.io.read_point_cloud(plyfile) - num_views, num_angles, num_depths = 300, 12, 4 + num_views = 300 views = generate_views(num_views) vis = o3d.visualization.Visualizer() @@ -323,6 +323,87 @@ def visObjGrasp(dataset_root, obj_idx, num_grasp=10, th=0.5, max_width=0.08, sav if show: o3d.visualization.draw_geometries([model, *grippers]) +def visObjFricReps(dataset_root, obj_idx, num_fric_rep=3, th=0.8, save_folder='save_fig', show=False): + ''' + Author: chenxi-wang, h.s-fang + + **Input:** + + - dataset_root: str, graspnet dataset root + + - obj_idx: int, index of object model + + - num_fric_rep: int, number of sampled representations + + - th: float, threshold of averaged friction coefficient of a view + + - save_folder: str, folder to save screen captures + + - show: bool, show visualization in open3d window if set to True + ''' + plyfile = os.path.join(dataset_root, 'models', '%03d'%obj_idx, 'nontextured.ply') + model = o3d.io.read_point_cloud(plyfile) + + num_views = 300 + views = generate_views(num_views) + + vis = o3d.visualization.Visualizer() + vis.create_window(width = 1280, height = 720) + ctr = vis.get_view_control() + param = get_camera_parameters(camera='kinect') + + cam_pos = np.load(os.path.join(dataset_root, 'scenes', 'scene_0000', 'kinect', 'cam0_wrt_table.npy')) + param.extrinsic = np.linalg.inv(cam_pos).tolist() + + sampled_points, data = get_model_fricreps('%s/fric_rep/%03d_labels_small.npz'%(dataset_root, obj_idx)) + + cnt = 0 + point_inds = [735] + # np.random.shuffle(point_inds) + reps = [] + + for point_ind in point_inds: + target_point = sampled_points[point_ind] + offset = data[point_ind] + view_inds = [252] + # np.random.shuffle(view_inds) + flag = False + for v in view_inds: + if flag: break + depth_inds = [0] + # np.random.shuffle(depth_inds) + for d in depth_inds: + if flag: break + # grasp_score_min = np.minimum(offset[v, :24, d, 1], offset[v, 24:, d, 1]) + # grasp_score_max = np.maximum(offset[v, :24, d, 1], offset[v, 24:, d, 1]) + # if not(( (grasp_score_min>0) & (grasp_score_max th*24 or sum_mu1 <= 1 or sum_mu2 > th*24 or sum_mu2 <= 1: + # continue + depth = 0.005 if d==0 else 0.01*d + view = views[v] + t = target_point + view_rot = viewpoint_params_to_matrix(-view, 0) + fric_rep = plot_fric_reps(t, view_rot, depth, offset[v,:,d,:]) + print(point_ind, v, d) + reps.append(fric_rep) + flag = True + if flag: + cnt += 1 + if cnt == num_fric_rep: + break + + vis.add_geometry(model) + for fric_rep in reps: + vis.add_geometry(fric_rep) + ctr.convert_from_pinhole_camera_parameters(param) + vis.poll_events() + filename = os.path.join(save_folder, 'object_{}_fric_rep.png'.format(obj_idx)) + vis.capture_screen_image(filename, do_render=True) + if show: + o3d.visualization.draw_geometries([model, *reps]) + def vis_rec_grasp(rec_grasp_tuples,numGrasp,image_path,save_path,show=False): ''' author: Minghao Gou