import numpy as np import pickle import json import os from task_gen_dependencies.transforms import transform_coordinates_3d from grasp_nms import nms_grasp class OmniObject: def __init__(self, name='obj', cam_info=None, type='Active', mask=None, pose=np.eye(4), size=np.array([0.001, 0.001, 0.001]) ): self.name = name self.cam_info = cam_info self.type = type self.obj_pose = pose self.obj_length = size self.xyz = np.array([0,0,0]) self.direction = np.array([0,0,0.05]) self.direction_proposals = None self.elements = {} if name=='gripper': self.elements['active'] = { 'push': [ { "part": "finger edge", "task": "forward push", "xyz": np.array([0,0,0]), "direction": np.array([0,0,0.08]), }, { "part": "finger edge", "task": "side push", "xyz": np.array([0,0,0]), "direction": np.array([1,0,0.3]), }, { "part": "finger edge", "task": "side push", "xyz": np.array([0,0,0]), "direction": np.array([-1,0,0.3]), }, ], 'press': { "part": "finger edge", "task": "press", "xyz": np.array([0,0,0]), "direction": np.array([0,0,1]), }, 'touch': { "part": "finger edge", "task": "touch", "xyz": np.array([0,0,0]), "direction": np.array([0,0,1]), }, 'pull': { "part": "finger edge", "task": "pull", "xyz": np.array([0,0,0]), "direction": np.array([0,0,0.08]), }, 'rotate': { "part": "finger edge", "task": "rotate", "xyz": np.array([0,0,0]), "direction": np.array([0,0,0.08]), } } @classmethod def from_obj_dir(cls, obj_dir, obj_info=None): if "interaction" in obj_info: obj_info = obj_info interaction_info=obj_info["interaction"] part_joint_limits_info = obj_info.get("part_joint_limits", None) else: obj_info_file = '%s/object_parameters.json'%obj_dir interaction_label_file = '%s/interaction.json'%obj_dir assert os.path.exists(obj_info_file), 'object_parameters.json not found in %s'%obj_dir assert os.path.exists(interaction_label_file), 'interaction.json not found in %s'%obj_dir obj_info = json.load(open(obj_info_file)) interaction_data = json.load(open(interaction_label_file)) interaction_info = interaction_data['interaction'] part_joint_limits_info = interaction_data.get("part_joint_limits", None) obj = cls( name=obj_info['object_id'], size=obj_info['size'] ) mesh_file = '%s/Aligned.obj'%obj_dir if os.path.exists(mesh_file): obj_info['mesh_file'] = mesh_file obj.part_joint_limits = part_joint_limits_info ''' Load interaction labels ''' obj.part_ids = [] for type in ['active', 'passive']: if type not in interaction_info: continue for action in interaction_info[type]: action_info = interaction_info[type][action] # import ipdb;ipdb.set_trace() if action=='grasp' and type=='passive': for grasp_part in action_info: grasp_files = action_info[grasp_part] grasp_data={ "grasp_pose": [], "width": [] } if isinstance(grasp_files, str): grasp_files = [grasp_files] for grasp_file in grasp_files: grasp_file = '%s/%s'%(obj_dir, grasp_file) if not os.path.exists(grasp_file): print('Grasp file %s not found'%grasp_file) continue _data = pickle.load(open(grasp_file, 'rb')) _data['grasp_pose'] = np.array(_data['grasp_pose']) _data['width'] = np.array(_data['width']) if _data['grasp_pose'].shape[0]==0: print('Empty grasp pose in %s'%grasp_file) continue grasp_data['grasp_pose'].append(_data['grasp_pose']) grasp_data['width'].append(_data['width']) if len(grasp_data['grasp_pose'])==0: print('Empty grasp pose in %s'%grasp_file) continue grasp_data['grasp_pose'] = np.concatenate(grasp_data['grasp_pose']) grasp_data['width'] = np.concatenate(grasp_data['width']) N_grasp = grasp_data['grasp_pose'].shape[0] max_grasps = 100 use_nms = True if N_grasp>60 else False if use_nms: # import ipdb;ipdb.set_trace() gripper_pose = grasp_data['grasp_pose'] N = gripper_pose.shape[0] width = grasp_data['width'][:N, np.newaxis] rotation = gripper_pose[:, :3, :3].reshape(N, -1) translation = gripper_pose[:, :3, 3] height = 0.02 * np.ones_like(width) depth = np.zeros_like(width) score = np.ones_like(width) * 0.1 obj_id = -1 * np.ones_like(width) grasp_group_array = np.concatenate([score, width, height, depth, rotation, translation, obj_id], axis=-1) if True: nms_info = interaction_info.get('grasp_nms_threshold', {'translation': 0.015, 'rotation': 25.0}) translation_thresh = nms_info['translation'] rotation_thresh = nms_info['rotation'] / 180.0 * np.pi grasp_group_array = nms_grasp(grasp_group_array, translation_thresh, rotation_thresh) rotation = grasp_group_array[:, 4:4+9] translation = grasp_group_array[:, 4+9:4+9+3] width = grasp_group_array[:, 1] grasp_data['grasp_pose'] = np.tile(np.eye(4), (grasp_group_array.shape[0], 1, 1)) grasp_data['grasp_pose'][:, :3, :3] = rotation.reshape(-1,3,3) grasp_data['grasp_pose'][:, :3, 3] = translation grasp_data['width'] = width #print(f"Grasp num of {obj.name} after NMS: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}") # downsample if grasp_data['grasp_pose'].shape[0]>max_grasps: grasp_num = grasp_data['grasp_pose'].shape[0] index = np.random.choice(grasp_num, max_grasps, replace=False) grasp_data['grasp_pose'] = grasp_data['grasp_pose'][index] grasp_data['width'] = grasp_data['width'][index] #print(f"Grasp num of {obj.name} after downsample: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}") action_info[grasp_part] = grasp_data interaction_info[type][action] = action_info else: for primitive in action_info: for primitive_info in action_info[primitive]: if 'part_id' in primitive_info: obj.part_ids.append(primitive_info['part_id']) obj.elements = interaction_info obj.info = obj_info # import ipdb;ipdb.set_trace() obj.is_articulated = True if part_joint_limits_info is not None else False return obj def set_element(self, element): action = element['action'] self.elements[action] = element def set_mask(self, mask, roi=None): self.mask = mask self.roi = roi def set_pose(self, pose, length): self.obj_pose = pose self.obj_length = length def set_part(self, xyz=None, direction=None, direction_proposals=None, relative=True): if xyz is not None: if not isinstance(xyz, np.ndarray): xyz = np.array(xyz) if relative: xyz = xyz * self.obj_length / 2.0 self.xyz = xyz if direction is not None: if not isinstance(direction, np.ndarray): direction = np.array(direction) direction = direction / np.linalg.norm(direction) * 0.08 self.direction = direction if direction_proposals is not None: self.direction_proposals = direction_proposals def format_object(self, relative=True): xyz, direction = self.xyz, self.direction obj_type = self.type.lower() if obj_type=='active': xyz_start = xyz xyz_end = xyz_start + direction elif obj_type=='passive' or obj_type=='plane': xyz_end = xyz xyz_start = xyz_end - direction arrow_in_obj = np.array([xyz_start, xyz_end]).transpose(1,0) arrow_in_world = transform_coordinates_3d(arrow_in_obj, self.obj_pose).transpose(1,0) xyz_start_world, xyz_end_world = arrow_in_world direction_world = xyz_end_world - xyz_start_world direction_world = direction_world / np.linalg.norm(direction_world) part2obj = np.eye(4) part2obj[:3, 3] = xyz_start self.obj2part = np.linalg.inv(part2obj) object_world = { 'pose': self.obj_pose, 'length': self.obj_length, 'xyz_start': xyz_start, 'xyz_end': xyz_end, 'xyz_start_world': xyz_start_world, 'xyz_end_world': xyz_end_world, 'direction': direction_world, 'obj2part': self.obj2part } return object_world