def execute_policy(self, rgbd_image_state, grasping_policy, grasp_pose_publisher, pose_frame): """ Executes a grasping policy on an RgbdImageState Parameters ---------- rgbd_image_state: :obj:`RgbdImageState` RgbdImageState from perception module to encapsulate depth and color image along with camera intrinsics grasping_policy: :obj:`GraspingPolicy` grasping policy to use grasp_pose_publisher: :obj:`Publisher` ROS Publisher to publish planned grasp's ROS Pose only for visualization pose_frame: :obj:`str` frame of reference to publish pose alone in """ # execute the policy's action rospy.loginfo('Planning Grasp') grasp_planning_start_time = time.time() grasp = grasping_policy(rgbd_image_state) # create GQCNNGrasp return msg and populate it gqcnn_grasp = GQCNNGrasp() gqcnn_grasp.grasp_success_prob = grasp.q_value gqcnn_grasp.pose = grasp.grasp.pose().pose_msg # create and publish the pose alone for visualization ease of grasp pose in rviz pose_stamped = PoseStamped() pose_stamped.pose = grasp.grasp.pose().pose_msg header = Header() header.stamp = rospy.Time.now() header.frame_id = pose_frame pose_stamped.header = header grasp_pose_publisher.publish(pose_stamped) # return GQCNNGrasp msg rospy.loginfo('Total grasp planning time: ' + str(time.time() - grasp_planning_start_time) + ' secs.') if self.cfg['vis']['vis_final_grasp']: vis.imshow(rgbd_image_state.rgbd_im) vis.grasp(grasp.grasp, scale=1.5, show_center=False, show_axis=True) vis.show() return gqcnn_grasp
def visualize(self): """ Visualize predictions """ logging.info('Visualizing ' + self.datapoint_type) # iterate through shuffled file indices for i in self.indices: im_filename = self.im_filenames[i] pose_filename = self.pose_filenames[i] label_filename = self.label_filenames[i] logging.info('Loading Image File: ' + im_filename + ' Pose File: ' + pose_filename + ' Label File: ' + label_filename) # load tensors from files metric_tensor = np.load(os.path.join(self.data_dir, label_filename))['arr_0'] label_tensor = 1 * (metric_tensor > self.metric_thresh) image_tensor = np.load(os.path.join(self.data_dir, im_filename))['arr_0'] hand_poses_tensor = np.load( os.path.join(self.data_dir, pose_filename))['arr_0'] pose_tensor = self._read_pose_data(hand_poses_tensor, self.input_data_mode) # score with neural network pred_p_success_tensor = self._gqcnn.predict( image_tensor, pose_tensor) # compute results classification_result = ClassificationResult( [pred_p_success_tensor], [label_tensor]) logging.info('Error rate on files: %.3f' % (classification_result.error_rate)) logging.info('Precision on files: %.3f' % (classification_result.precision)) logging.info('Recall on files: %.3f' % (classification_result.recall)) mispred_ind = classification_result.mispredicted_indices() correct_ind = classification_result.correct_indices() # IPython.embed() if self.datapoint_type == 'true_positive' or self.datapoint_type == 'true_negative': vis_ind = correct_ind else: vis_ind = mispred_ind num_visualized = 0 # visualize for ind in vis_ind: # limit the number of sampled datapoints displayed per object if num_visualized >= self.samples_per_object: break num_visualized += 1 # don't visualize the datapoints that we don't want if self.datapoint_type == 'true_positive': if classification_result.labels[ind] == 0: continue elif self.datapoint_type == 'true_negative': if classification_result.labels[ind] == 1: continue elif self.datapoint_type == 'false_positive': if classification_result.labels[ind] == 0: continue elif self.datapoint_type == 'false_negative': if classification_result.labels[ind] == 1: continue logging.info('Datapoint %d of files for %s' % (ind, im_filename)) logging.info('Depth: %.3f' % (hand_poses_tensor[ind, 2])) data = image_tensor[ind, ...] if self.display_image_type == RenderMode.SEGMASK: image = BinaryImage(data) elif self.display_image_type == RenderMode.GRAYSCALE: image = GrayscaleImage(data) elif self.display_image_type == RenderMode.COLOR: image = ColorImage(data) elif self.display_image_type == RenderMode.DEPTH: image = DepthImage(data) elif self.display_image_type == RenderMode.RGBD: image = RgbdImage(data) elif self.display_image_type == RenderMode.GD: image = GdImage(data) vis2d.figure() if self.display_image_type == RenderMode.RGBD: vis2d.subplot(1, 2, 1) vis2d.imshow(image.color) grasp = Grasp2D(Point(image.center, 'img'), 0, hand_poses_tensor[ind, 2], self.gripper_width_m) grasp.camera_intr = grasp.camera_intr.resize(1.0 / 3.0) vis2d.grasp(grasp) vis2d.subplot(1, 2, 2) vis2d.imshow(image.depth) vis2d.grasp(grasp) elif self.display_image_type == RenderMode.GD: vis2d.subplot(1, 2, 1) vis2d.imshow(image.gray) grasp = Grasp2D(Point(image.center, 'img'), 0, hand_poses_tensor[ind, 2], self.gripper_width_m) grasp.camera_intr = grasp.camera_intr.resize(1.0 / 3.0) vis2d.grasp(grasp) vis2d.subplot(1, 2, 2) vis2d.imshow(image.depth) vis2d.grasp(grasp) else: vis2d.imshow(image) grasp = Grasp2D(Point(image.center, 'img'), 0, hand_poses_tensor[ind, 2], self.gripper_width_m) grasp.camera_intr = grasp.camera_intr.resize(1.0 / 3.0) vis2d.grasp(grasp) vis2d.title('Datapoint %d: Pred: %.3f Label: %.3f' % (ind, classification_result.pred_probs[ind, 1], classification_result.labels[ind])) vis2d.show() # cleanup self._cleanup()
def action(self, state): """ Plans the grasp with the highest probability of success on the given RGB-D image. Attributes ---------- state : :obj:`RgbdImageState` image to plan grasps on Returns ------- :obj:`ParallelJawGrasp` grasp to execute """ # take the greedy action with prob 1 - epsilon if np.random.rand() > self.epsilon: logging.debug('Taking greedy action') return CrossEntropyAntipodalGraspingPolicy.action(self, state) # otherwise take a random action logging.debug('Taking random action') # check valid input if not isinstance(state, RgbdImageState): raise ValueError('Must provide an RGB-D image state.') # parse state rgbd_im = state.rgbd_im camera_intr = state.camera_intr segmask = state.segmask # sample random antipodal grasps grasps = self._grasp_sampler.sample( rgbd_im, camera_intr, self._num_seed_samples, segmask=segmask, visualize=self.config['vis']['grasp_sampling'], seed=self._seed) num_grasps = len(grasps) if num_grasps == 0: logging.warning('No valid grasps could be found') return None # choose a grasp uniformly at random grasp_ind = np.random.choice(num_grasps, size=1)[0] grasp = grasps[grasp_ind] depth = grasp.depth # create transformed image image_tensor, pose_tensor = self.grasps_to_tensors([grasp], state) image = DepthImage(image_tensor[0, ...]) # predict prob success output_arr = self.gqcnn.predict(image_tensor, pose_tensor) q_value = output_arr[0, -1] # visualize planned grasp if self.config['vis']['grasp_plan']: scale_factor = float(self.gqcnn.im_width) / float(self._crop_width) scaled_camera_intr = camera_intr.resize(scale_factor) vis_grasp = Grasp2D(Point(image.center), 0.0, depth, width=self._gripper_width, camera_intr=scaled_camera_intr) vis.figure() vis.imshow(image) vis.grasp(vis_grasp, scale=1.5, show_center=False, show_axis=True) vis.title('Best Grasp: d=%.3f, q=%.3f' % (depth, q_value)) vis.show() # return action return ParallelJawGrasp(grasp, q_value, image)
def action(self, state): """ Plans the grasp with the highest probability of success on the given RGB-D image. Attributes ---------- state : :obj:`RgbdImageState` image to plan grasps on Returns ------- :obj:`ParallelJawGrasp` grasp to execute """ # check valid input if not isinstance(state, RgbdImageState): raise ValueError('Must provide an RGB-D image state.') # parse state rgbd_im = state.rgbd_im camera_intr = state.camera_intr segmask = state.segmask # sample grasps grasps = self._grasp_sampler.sample( rgbd_im, camera_intr, self._num_seed_samples, segmask=segmask, visualize=self.config['vis']['grasp_sampling'], seed=self._seed) num_grasps = len(grasps) if num_grasps == 0: logging.warning('No valid grasps could be found') return None # form tensors image_tensor, pose_tensor = self.grasps_to_tensors(grasps, state) if self.config['vis']['tf_images']: # read vis params k = self.config['vis']['k'] d = utils.sqrt_ceil(k) # display grasp transformed images vis.figure(size=(FIGSIZE, FIGSIZE)) for i, image_tf in enumerate(image_tensor[:k, ...]): depth = pose_tensor[i][0] vis.subplot(d, d, i + 1) vis.imshow(DepthImage(image_tf)) vis.title('Image %d: d=%.3f' % (i, depth)) # display grasp transformed images vis.figure(size=(FIGSIZE, FIGSIZE)) for i in range(d): image_tf = image_tensor[i, ...] depth = pose_tensor[i][0] grasp = grasps[i] vis.subplot(d, 2, 2 * i + 1) vis.imshow(rgbd_im.depth) vis.grasp(grasp, scale=1.5, show_center=False, show_axis=True) vis.title('Grasp %d: d=%.3f' % (i, depth)) vis.subplot(d, 2, 2 * i + 2) vis.imshow(DepthImage(image_tf)) vis.title('TF image %d: d=%.3f' % (i, depth)) vis.show() # iteratively refit and sample for j in range(self._num_iters): logging.debug('CEM iter %d' % (j)) # predict grasps predict_start = time() output_arr = self.gqcnn.predict(image_tensor, pose_tensor) q_values = output_arr[:, -1] logging.debug('Prediction took %.3f sec' % (time() - predict_start)) # sort grasps q_values_and_indices = zip(q_values, np.arange(num_grasps)) q_values_and_indices.sort(key=lambda x: x[0], reverse=True) if self.config['vis']['grasp_candidates']: # display each grasp on the original image, colored by predicted success vis.figure(size=(FIGSIZE, FIGSIZE)) vis.imshow(rgbd_im.depth) for grasp, q in zip(grasps, q_values): vis.grasp(grasp, scale=1.5, show_center=False, show_axis=True, color=plt.cm.RdYlBu(q)) vis.title('Sampled grasps iter %d' % (j)) vis.show() if self.config['vis']['grasp_ranking']: # read vis params k = self.config['vis']['k'] d = utils.sqrt_ceil(k) # form camera intr for the thumbnail (to compute gripper width) scale_factor = float(self.gqcnn.im_width) / float( self._crop_width) scaled_camera_intr = camera_intr.resize(scale_factor) vis.figure(size=(FIGSIZE, FIGSIZE)) for i, p in enumerate(q_values_and_indices[:k]): # read stats for grasp q_value = p[0] ind = p[1] depth = pose_tensor[ind][0] image = DepthImage(image_tensor[ind, ...]) grasp = Grasp2D(Point(image.center), 0.0, depth, width=self._gripper_width, camera_intr=scaled_camera_intr) # plot vis.subplot(d, d, i + 1) vis.imshow(image) vis.grasp(grasp, scale=1.5) vis.title('K=%d: d=%.3f, q=%.3f' % (i, depth, q_value)) vis.show() # fit elite set num_refit = max(int(np.ceil(self._gmm_refit_p * num_grasps)), 1) elite_q_values = [i[0] for i in q_values_and_indices[:num_refit]] elite_grasp_indices = [ i[1] for i in q_values_and_indices[:num_refit] ] elite_grasps = [grasps[i] for i in elite_grasp_indices] elite_grasp_arr = np.array([g.feature_vec for g in elite_grasps]) if self.config['vis']['elite_grasps']: # display each grasp on the original image, colored by predicted success vis.figure(size=(FIGSIZE, FIGSIZE)) vis.imshow(rgbd_im.depth) for grasp, q in zip(elite_grasps, elite_q_values): vis.grasp(grasp, scale=1.5, show_center=False, show_axis=True, color=plt.cm.RdYlBu(q)) vis.title('Elite grasps iter %d' % (j)) vis.show() # normalize elite set elite_grasp_mean = np.mean(elite_grasp_arr, axis=0) elite_grasp_std = np.std(elite_grasp_arr, axis=0) elite_grasp_std[elite_grasp_std == 0] = 1.0 elite_grasp_arr = (elite_grasp_arr - elite_grasp_mean) / elite_grasp_std # fit a GMM to the top samples num_components = max( int(np.ceil(self._gmm_component_frac * num_refit)), 1) uniform_weights = (1.0 / num_components) * np.ones(num_components) gmm = GaussianMixture(n_components=num_components, weights_init=uniform_weights, reg_covar=self._gmm_reg_covar) train_start = time() gmm.fit(elite_grasp_arr) train_duration = time() - train_start logging.debug('GMM fitting with %d components took %.3f sec' % (num_components, train_duration)) # sample the next grasps sample_start = time() grasp_vecs, _ = gmm.sample(n_samples=self._num_gmm_samples) grasp_vecs = elite_grasp_std * grasp_vecs + elite_grasp_mean sample_duration = time() - sample_start logging.debug('GMM sampling took %.3f sec' % (sample_duration)) # convert features to grasps grasps = [] for grasp_vec in grasp_vecs: grasps.append( Grasp2D.from_feature_vec(grasp_vec, width=self._gripper_width, camera_intr=camera_intr)) num_grasps = len(grasps) if num_grasps == 0: logging.warning('No valid grasps could be found') return None # form tensors image_tensor, pose_tensor = self.grasps_to_tensors(grasps, state) if self.config['vis']['tf_images']: # read vis params k = self.config['vis']['k'] d = utils.sqrt_ceil(k) # display grasp transformed images vis.figure(size=(FIGSIZE, FIGSIZE)) for i, image_tf in enumerate(image_tensor[:k, ...]): depth = pose_tensor[i][0] vis.subplot(d, d, i + 1) vis.imshow(DepthImage(image_tf)) vis.title('Image %d: d=%.3f' % (i, depth)) vis.show() # predict final set of grasps predict_start = time() output_arr = self.gqcnn.predict(image_tensor, pose_tensor) q_values = output_arr[:, -1] logging.debug('Final prediction took %.3f sec' % (time() - predict_start)) if self.config['vis']['grasp_candidates']: # display each grasp on the original image, colored by predicted success vis.figure(size=(FIGSIZE, FIGSIZE)) vis.imshow(rgbd_im.depth) for grasp, q in zip(grasps, q_values): vis.grasp(grasp, scale=1.5, show_center=False, show_axis=True, color=plt.cm.RdYlBu(q)) vis.title('Final sampled grasps') vis.show() # select grasp index = self.select(grasps, q_values) grasp = grasps[index] q_value = q_values[index] image = DepthImage(image_tensor[index, ...]) pose = pose_tensor[index, ...] depth = pose[0] if self.config['vis']['grasp_plan']: scale_factor = float(self.gqcnn.im_width) / float(self._crop_width) scaled_camera_intr = camera_intr.resize(scale_factor) grasp = Grasp2D(Point(image.center), 0.0, pose[0], width=self._gripper_width, camera_intr=scaled_camera_intr) vis.figure() vis.imshow(image) vis.grasp(grasp, scale=1.5, show_center=False, show_axis=True) vis.title('Best Grasp: d=%.3f, q=%.3f' % (depth, q_value)) vis.show() # return action return ParallelJawGrasp(grasp, q_value, image)
def action(self, state): """ Plans the grasp with the highest probability of success on the given RGB-D image. Attributes ---------- state : :obj:`RgbdImageState` image to plan grasps on Returns ------- :obj:`ParallelJawGrasp` grasp to execute """ # check valid input if not isinstance(state, RgbdImageState): raise ValueError('Must provide an RGB-D image state.') # parse state rgbd_im = state.rgbd_im camera_intr = state.camera_intr segmask = state.segmask # sample grasps grasps = self._grasp_sampler.sample( rgbd_im, camera_intr, self._num_grasp_samples, segmask=segmask, visualize=self.config['vis']['grasp_sampling'], seed=None) num_grasps = len(grasps) # form tensors image_tensor, pose_tensor = self.grasps_to_tensors(grasps, state) if self.config['vis']['tf_images']: # read vis params k = self.config['vis']['k'] d = utils.sqrt_ceil(k) # display grasp transformed images vis.figure(size=(FIGSIZE, FIGSIZE)) for i, image_tf in enumerate(image_tensor[:k, ...]): depth = pose_tensor[i][0] vis.subplot(d, d, i + 1) vis.imshow(DepthImage(image_tf)) vis.title('Image %d: d=%.3f' % (i, depth)) vis.show() # predict grasps predict_start = time() output_arr = self.gqcnn.predict(image_tensor, pose_tensor) q_values = output_arr[:, -1] logging.debug('Prediction took %.3f sec' % (time() - predict_start)) if self.config['vis']['grasp_candidates']: # display each grasp on the original image, colored by predicted success vis.figure(size=(FIGSIZE, FIGSIZE)) vis.imshow(rgbd_im.depth) for grasp, q in zip(grasps, q_values): vis.grasp(grasp, scale=1.5, show_center=False, show_axis=True, color=plt.cm.RdYlBu(q)) vis.title('Sampled grasps') vis.show() if self.config['vis']['grasp_ranking']: # read vis params k = self.config['vis']['k'] d = utils.sqrt_ceil(k) # form camera intr for the thumbnail (to compute gripper width) scale_factor = float(self.gqcnn.im_width) / float(self._crop_width) scaled_camera_intr = camera_intr.resize(scale_factor) # sort grasps q_values_and_indices = zip(q_values, np.arange(num_grasps)) q_values_and_indices.sort(key=lambda x: x[0], reverse=True) vis.figure(size=(FIGSIZE, FIGSIZE)) for i, p in enumerate(q_values_and_indices[:k]): # read stats for grasp q_value = p[0] ind = p[1] depth = pose_tensor[ind][0] image = DepthImage(image_tensor[ind, ...]) grasp = Grasp2D(Point(image.center), 0.0, depth, width=self._gripper_width, camera_intr=scaled_camera_intr) # plot vis.subplot(d, d, i + 1) vis.imshow(image) vis.grasp(grasp, scale=1.5) vis.title('K=%d: d=%.3f, q=%.3f' % (i, depth, q_value)) vis.show() # select grasp index = self.select(grasps, q_values) grasp = grasps[index] q_value = q_values[index] image = DepthImage(image_tensor[index, ...]) pose = pose_tensor[index, ...] depth = pose[0] if self.config['vis']['grasp_plan']: scale_factor = float(self.gqcnn.im_width) / float(self._crop_width) scaled_camera_intr = camera_intr.resize(scale_factor) grasp = Grasp2D(Point(image.center), 0.0, pose[0], width=self._gripper_width, camera_intr=scaled_camera_intr) vis.figure() vis.imshow(image) vis.grasp(grasp, scale=1.5, show_center=False, show_axis=True) vis.title('Best Grasp: d=%.3f, q=%.3f' % (depth, q_value)) vis.show() # return action return ParallelJawGrasp(grasp, q_value, image)
# setup sensor sensor = RgbdSensorFactory.sensor(sensor_type, config['sensor']) sensor.start() camera_intr = sensor.ir_intrinsics # read images color_im, depth_im, _ = sensor.frames() color_im = color_im.inpaint(rescale_factor=inpaint_rescale_factor) depth_im = depth_im.inpaint(rescale_factor=inpaint_rescale_factor) rgbd_im = RgbdImage.from_color_and_depth(color_im, depth_im) state = RgbdImageState(rgbd_im, camera_intr) # init policy policy = CrossEntropyAntipodalGraspingPolicy(policy_config) policy_start = time.time() action = policy(state) logging.info('Planning took %.3f sec' % (time.time() - policy_start)) # vis final grasp if policy_config['vis']['final_grasp']: vis.figure(size=(10, 10)) vis.subplot(1, 2, 1) vis.imshow(rgbd_im.color) vis.grasp(action.grasp, scale=1.5, show_center=False, show_axis=True) vis.title('Planned grasp on color (Q=%.3f)' % (action.q_value)) vis.subplot(1, 2, 2) vis.imshow(rgbd_im.depth) vis.grasp(action.grasp, scale=1.5, show_center=False, show_axis=True) vis.title('Planned grasp on depth (Q=%.3f)' % (action.q_value)) vis.show()
def generate_gqcnn_dataset(dataset_path, database, target_object_keys, env_rv_params, gripper_name, config): """ Generates a GQ-CNN TensorDataset for training models with new grippers, quality metrics, objects, and cameras. Parameters ---------- dataset_path : str path to save the dataset to database : :obj:`Hdf5Database` Dex-Net database containing the 3D meshes, grasps, and grasp metrics target_object_keys : :obj:`OrderedDict` dictionary mapping dataset names to target objects (either 'all' or a list of specific object keys) env_rv_params : :obj:`OrderedDict` parameters of the camera and object random variables used in sampling (see meshpy_berkeley.UniformPlanarWorksurfaceImageRandomVariable for more info) gripper_name : str name of the gripper to use config : :obj:`autolab_core.YamlConfig` other parameters for dataset generation Notes ----- Required parameters of config are specified in Other Parameters Other Parameters ---------------- images_per_stable_pose : int number of object and camera poses to sample for each stable pose stable_pose_min_p : float minimum probability of occurrence for a stable pose to be used in data generation (used to prune bad stable poses gqcnn/crop_width : int width, in pixels, of crop region around each grasp center, before resize (changes the size of the region seen by the GQ-CNN) gqcnn/crop_height : int height, in pixels, of crop region around each grasp center, before resize (changes the size of the region seen by the GQ-CNN) gqcnn/final_width : int width, in pixels, of final transformed grasp image for input to the GQ-CNN (defaults to 32) gqcnn/final_height : int height, in pixels, of final transformed grasp image for input to the GQ-CNN (defaults to 32) table_alignment/max_approach_table_angle : float max angle between the grasp axis and the table normal when the grasp approach is maximally aligned with the table normal table_alignment/max_approach_offset : float max deviation from perpendicular approach direction to use in grasp collision checking table_alignment/num_approach_offset_samples : int number of approach samples to use in collision checking collision_checking/table_offset : float max allowable interpenetration between the gripper and table to be considered collision free collision_checking/table_mesh_filename : str path to a table mesh for collision checking (default data/meshes/table.obj) collision_checking/approach_dist : float distance, in meters, between the approach pose and final grasp pose along the grasp axis collision_checking/delta_approach : float amount, in meters, to discretize the straight-line path from the gripper approach pose to the final grasp pose tensors/datapoints_per_file : int number of datapoints to store in each unique tensor file on disk tensors/fields : :obj:`dict` dictionary mapping field names to dictionaries specifying the data type, height, width, and number of channels for each tensor debug : bool True (or 1) if the random seed should be set to enforce deterministic behavior, False (0) otherwise vis/candidate_grasps : bool True (or 1) if the collision free candidate grasps should be displayed in 3D (for debugging) vis/rendered_images : bool True (or 1) if the rendered images for each stable pose should be displayed (for debugging) vis/grasp_images : bool True (or 1) if the transformed grasp images should be displayed (for debugging) """ # read data gen params output_dir = dataset_path gripper = RobotGripper.load(gripper_name) image_samples_per_stable_pose = config['images_per_stable_pose'] stable_pose_min_p = config['stable_pose_min_p'] # read gqcnn params gqcnn_params = config['gqcnn'] im_crop_height = gqcnn_params['crop_height'] im_crop_width = gqcnn_params['crop_width'] im_final_height = gqcnn_params['final_height'] im_final_width = gqcnn_params['final_width'] cx_crop = float(im_crop_width) / 2 cy_crop = float(im_crop_height) / 2 # open database dataset_names = target_object_keys.keys() datasets = [database.dataset(dn) for dn in dataset_names] # set target objects for dataset in datasets: if target_object_keys[dataset.name] == 'all': target_object_keys[dataset.name] = dataset.object_keys # setup grasp params table_alignment_params = config['table_alignment'] min_grasp_approach_offset = -np.deg2rad( table_alignment_params['max_approach_offset']) max_grasp_approach_offset = np.deg2rad( table_alignment_params['max_approach_offset']) max_grasp_approach_table_angle = np.deg2rad( table_alignment_params['max_approach_table_angle']) num_grasp_approach_samples = table_alignment_params[ 'num_approach_offset_samples'] phi_offsets = [] if max_grasp_approach_offset == min_grasp_approach_offset: phi_inc = 1 elif num_grasp_approach_samples == 1: phi_inc = max_grasp_approach_offset - min_grasp_approach_offset + 1 else: phi_inc = (max_grasp_approach_offset - min_grasp_approach_offset) / ( num_grasp_approach_samples - 1) phi = min_grasp_approach_offset while phi <= max_grasp_approach_offset: phi_offsets.append(phi) phi += phi_inc # setup collision checking coll_check_params = config['collision_checking'] approach_dist = coll_check_params['approach_dist'] delta_approach = coll_check_params['delta_approach'] table_offset = coll_check_params['table_offset'] table_mesh_filename = coll_check_params['table_mesh_filename'] if not os.path.isabs(table_mesh_filename): table_mesh_filename = os.path.join( os.path.dirname(os.path.realpath(__file__)), '..', table_mesh_filename) table_mesh = ObjFile(table_mesh_filename).read() # set tensor dataset config tensor_config = config['tensors'] tensor_config['fields']['depth_ims_tf_table']['height'] = im_final_height tensor_config['fields']['depth_ims_tf_table']['width'] = im_final_width tensor_config['fields']['obj_masks']['height'] = im_final_height tensor_config['fields']['obj_masks']['width'] = im_final_width # add available metrics (assuming same are computed for all objects) metric_names = [] dataset = datasets[0] obj_keys = dataset.object_keys if len(obj_keys) == 0: raise ValueError('No valid objects in dataset %s' % (dataset.name)) obj = dataset[obj_keys[0]] grasps = dataset.grasps(obj.key, gripper=gripper.name) grasp_metrics = dataset.grasp_metrics(obj.key, grasps, gripper=gripper.name) metric_names = grasp_metrics[grasp_metrics.keys()[0]].keys() for metric_name in metric_names: tensor_config['fields'][metric_name] = {} tensor_config['fields'][metric_name]['dtype'] = 'float32' # init tensor dataset tensor_dataset = TensorDataset(output_dir, tensor_config) tensor_datapoint = tensor_dataset.datapoint_template # setup log file experiment_log_filename = os.path.join(output_dir, 'dataset_generation.log') formatter = logging.Formatter('%(asctime)s %(levelname)s: %(message)s') hdlr = logging.FileHandler(experiment_log_filename) hdlr.setFormatter(formatter) logging.getLogger().addHandler(hdlr) root_logger = logging.getLogger() # copy config out_config_filename = os.path.join(output_dir, 'dataset_generation.json') ordered_dict_config = collections.OrderedDict() for key in config.keys(): ordered_dict_config[key] = config[key] with open(out_config_filename, 'w') as outfile: json.dump(ordered_dict_config, outfile) # 1. Precompute the set of valid grasps for each stable pose: # i) Perpendicular to the table # ii) Collision-free along the approach direction # load grasps if they already exist grasp_cache_filename = os.path.join(output_dir, CACHE_FILENAME) if os.path.exists(grasp_cache_filename): logging.info('Loading grasp candidates from file') candidate_grasps_dict = pkl.load(open(grasp_cache_filename, 'rb')) # otherwise re-compute by reading from the database and enforcing constraints else: # create grasps dict candidate_grasps_dict = {} # loop through datasets and objects for dataset in datasets: logging.info('Reading dataset %s' % (dataset.name)) for obj in dataset: if obj.key not in target_object_keys[dataset.name]: continue # init candidate grasp storage candidate_grasps_dict[obj.key] = {} # setup collision checker collision_checker = GraspCollisionChecker(gripper) collision_checker.set_graspable_object(obj) # read in the stable poses of the mesh stable_poses = dataset.stable_poses(obj.key) for i, stable_pose in enumerate(stable_poses): # render images if stable pose is valid if stable_pose.p > stable_pose_min_p: candidate_grasps_dict[obj.key][stable_pose.id] = [] # setup table in collision checker T_obj_stp = stable_pose.T_obj_table.as_frames( 'obj', 'stp') T_obj_table = obj.mesh.get_T_surface_obj( T_obj_stp, delta=table_offset).as_frames('obj', 'table') T_table_obj = T_obj_table.inverse() collision_checker.set_table(table_mesh_filename, T_table_obj) # read grasp and metrics grasps = dataset.grasps(obj.key, gripper=gripper.name) logging.info( 'Aligning %d grasps for object %s in stable %s' % (len(grasps), obj.key, stable_pose.id)) # align grasps with the table aligned_grasps = [ grasp.perpendicular_table(stable_pose) for grasp in grasps ] # check grasp validity logging.info( 'Checking collisions for %d grasps for object %s in stable %s' % (len(grasps), obj.key, stable_pose.id)) for aligned_grasp in aligned_grasps: # check angle with table plane and skip unaligned grasps _, grasp_approach_table_angle, _ = aligned_grasp.grasp_angles_from_stp_z( stable_pose) perpendicular_table = ( np.abs(grasp_approach_table_angle) < max_grasp_approach_table_angle) if not perpendicular_table: continue # check whether any valid approach directions are collision free collision_free = False for phi_offset in phi_offsets: rotated_grasp = aligned_grasp.grasp_y_axis_offset( phi_offset) collides = collision_checker.collides_along_approach( rotated_grasp, approach_dist, delta_approach) if not collides: collision_free = True break # store if aligned to table candidate_grasps_dict[obj.key][ stable_pose.id].append( GraspInfo(aligned_grasp, collision_free)) # visualize if specified if collision_free and config['vis'][ 'candidate_grasps']: logging.info('Grasp %d' % (aligned_grasp.id)) vis.figure() vis.gripper_on_object(gripper, aligned_grasp, obj, stable_pose.T_obj_world) vis.show() # save to file logging.info('Saving to file') pkl.dump(candidate_grasps_dict, open(grasp_cache_filename, 'wb')) # 2. Render a dataset of images and associate the gripper pose with image coordinates for each grasp in the Dex-Net database # setup variables obj_category_map = {} pose_category_map = {} cur_pose_label = 0 cur_obj_label = 0 cur_image_label = 0 # render images for each stable pose of each object in the dataset render_modes = [RenderMode.SEGMASK, RenderMode.DEPTH_SCENE] for dataset in datasets: logging.info('Generating data for dataset %s' % (dataset.name)) # iterate through all object keys object_keys = dataset.object_keys for obj_key in object_keys: obj = dataset[obj_key] if obj.key not in target_object_keys[dataset.name]: continue # read in the stable poses of the mesh stable_poses = dataset.stable_poses(obj.key) for i, stable_pose in enumerate(stable_poses): # render images if stable pose is valid if stable_pose.p > stable_pose_min_p: # log progress logging.info('Rendering images for object %s in %s' % (obj.key, stable_pose.id)) # add to category maps if obj.key not in obj_category_map.keys(): obj_category_map[obj.key] = cur_obj_label pose_category_map['%s_%s' % (obj.key, stable_pose.id)] = cur_pose_label # read in candidate grasps and metrics candidate_grasp_info = candidate_grasps_dict[obj.key][ stable_pose.id] candidate_grasps = [g.grasp for g in candidate_grasp_info] grasp_metrics = dataset.grasp_metrics(obj.key, candidate_grasps, gripper=gripper.name) # compute object pose relative to the table T_obj_stp = stable_pose.T_obj_table.as_frames('obj', 'stp') T_obj_stp = obj.mesh.get_T_surface_obj(T_obj_stp) # sample images from random variable T_table_obj = RigidTransform(from_frame='table', to_frame='obj') scene_objs = { 'table': SceneObject(table_mesh, T_table_obj) } urv = UniformPlanarWorksurfaceImageRandomVariable( obj.mesh, render_modes, 'camera', env_rv_params, stable_pose=stable_pose, scene_objs=scene_objs) render_start = time.time() render_samples = urv.rvs( size=image_samples_per_stable_pose) render_stop = time.time() logging.info('Rendering images took %.3f sec' % (render_stop - render_start)) # visualize if config['vis']['rendered_images']: d = int(np.ceil( np.sqrt(image_samples_per_stable_pose))) # binary vis2d.figure() for j, render_sample in enumerate(render_samples): vis2d.subplot(d, d, j + 1) vis2d.imshow(render_sample.renders[ RenderMode.SEGMASK].image) # depth table vis2d.figure() for j, render_sample in enumerate(render_samples): vis2d.subplot(d, d, j + 1) vis2d.imshow(render_sample.renders[ RenderMode.DEPTH_SCENE].image) vis2d.show() # tally total amount of data num_grasps = len(candidate_grasps) num_images = image_samples_per_stable_pose num_save = num_images * num_grasps logging.info('Saving %d datapoints' % (num_save)) # for each candidate grasp on the object compute the projection # of the grasp into image space for render_sample in render_samples: # read images binary_im = render_sample.renders[ RenderMode.SEGMASK].image depth_im_table = render_sample.renders[ RenderMode.DEPTH_SCENE].image # read camera params T_stp_camera = render_sample.camera.object_to_camera_pose shifted_camera_intr = render_sample.camera.camera_intr # read pixel offsets cx = depth_im_table.center[1] cy = depth_im_table.center[0] # compute intrinsics for virtual camera of the final # cropped and rescaled images camera_intr_scale = float(im_final_height) / float( im_crop_height) cropped_camera_intr = shifted_camera_intr.crop( im_crop_height, im_crop_width, cy, cx) final_camera_intr = cropped_camera_intr.resize( camera_intr_scale) # create a thumbnail for each grasp for grasp_info in candidate_grasp_info: # read info grasp = grasp_info.grasp collision_free = grasp_info.collision_free # get the gripper pose T_obj_camera = T_stp_camera * T_obj_stp.as_frames( 'obj', T_stp_camera.from_frame) grasp_2d = grasp.project_camera( T_obj_camera, shifted_camera_intr) # center images on the grasp, rotate to image x axis dx = cx - grasp_2d.center.x dy = cy - grasp_2d.center.y translation = np.array([dy, dx]) binary_im_tf = binary_im.transform( translation, grasp_2d.angle) depth_im_tf_table = depth_im_table.transform( translation, grasp_2d.angle) # crop to image size binary_im_tf = binary_im_tf.crop( im_crop_height, im_crop_width) depth_im_tf_table = depth_im_tf_table.crop( im_crop_height, im_crop_width) # resize to image size binary_im_tf = binary_im_tf.resize( (im_final_height, im_final_width), interp='nearest') depth_im_tf_table = depth_im_tf_table.resize( (im_final_height, im_final_width)) # visualize the transformed images if config['vis']['grasp_images']: grasp_center = Point( depth_im_tf_table.center, frame=final_camera_intr.frame) tf_grasp_2d = Grasp2D( grasp_center, 0, grasp_2d.depth, width=gripper.max_width, camera_intr=final_camera_intr) # plot 2D grasp image vis2d.figure() vis2d.subplot(2, 2, 1) vis2d.imshow(binary_im) vis2d.grasp(grasp_2d) vis2d.subplot(2, 2, 2) vis2d.imshow(depth_im_table) vis2d.grasp(grasp_2d) vis2d.subplot(2, 2, 3) vis2d.imshow(binary_im_tf) vis2d.grasp(tf_grasp_2d) vis2d.subplot(2, 2, 4) vis2d.imshow(depth_im_tf_table) vis2d.grasp(tf_grasp_2d) vis2d.title('Coll Free? %d' % (grasp_info.collision_free)) vis2d.show() # plot 3D visualization vis.figure() T_obj_world = vis.mesh_stable_pose( obj.mesh, stable_pose.T_obj_world, style='surface', dim=0.5) vis.gripper(gripper, grasp, T_obj_world, color=(0.3, 0.3, 0.3)) vis.show() # form hand pose array hand_pose = np.r_[grasp_2d.center.y, grasp_2d.center.x, grasp_2d.depth, grasp_2d.angle, grasp_2d.center.y - shifted_camera_intr.cy, grasp_2d.center.x - shifted_camera_intr.cx, grasp_2d.width_px] # store to data buffers tensor_datapoint[ 'depth_ims_tf_table'] = depth_im_tf_table.raw_data tensor_datapoint[ 'obj_masks'] = binary_im_tf.raw_data tensor_datapoint['hand_poses'] = hand_pose tensor_datapoint['collision_free'] = collision_free tensor_datapoint['obj_labels'] = cur_obj_label tensor_datapoint['pose_labels'] = cur_pose_label tensor_datapoint['image_labels'] = cur_image_label for metric_name, metric_val in grasp_metrics[ grasp.id].iteritems(): coll_free_metric = ( 1 * collision_free) * metric_val tensor_datapoint[ metric_name] = coll_free_metric tensor_dataset.add(tensor_datapoint) # update image label cur_image_label += 1 # update pose label cur_pose_label += 1 # force clean up gc.collect() # update object label cur_obj_label += 1 # force clean up gc.collect() # save last file tensor_dataset.flush() # save category mappings obj_cat_filename = os.path.join(output_dir, 'object_category_map.json') json.dump(obj_category_map, open(obj_cat_filename, 'w')) pose_cat_filename = os.path.join(output_dir, 'pose_category_map.json') json.dump(pose_category_map, open(pose_cat_filename, 'w'))
def visualize_tensor_dataset(dataset, config): """ Visualizes a Tensor dataset. Parameters ---------- dataset : :obj:`TensorDataset` dataset to visualize config : :obj:`autolab_core.YamlConfig` parameters for visualization Notes ----- Required parameters of config are specified in Other Parameters Other Parameters ---------------- field_name : str name of the field in the TensorDataset to visualize (defaults to depth_ims_tf_table, which is a single view point cloud of the object on a table) field_type : str type of image that the field name correspondes to (defaults to depth, can also be `segmask` if using the field `object_masks`) print_fields : :obj:`list` of str names of additiona fields to print to the command line filter : :obj:`dict` mapping str to :obj:`dict` contraints that all displayed datapoints must satisfy (supports any univariate field name as a key and numeric thresholds) gripper_width_px : float width of the gripper to plot in pixels font_size : int size of font on the rendered images """ # shuffle the tensor indices indices = dataset.datapoint_indices np.random.shuffle(indices) # read config field_name = config['field_name'] field_type = config['field_type'] font_size = config['font_size'] print_fields = config['print_fields'] gripper_width_px = config['gripper_width_px'] num = 0 for i, ind in enumerate(indices): datapoint = dataset[ind] data = datapoint[field_name] if field_type == RenderMode.SEGMASK: image = BinaryImage(data) elif field_type == RenderMode.DEPTH: image = DepthImage(data) else: raise ValueError('Field type %s not supported!' % (field_type)) skip_datapoint = False for f, filter_cfg in config['filter'].iteritems(): data = datapoint[f] if 'greater_than' in filter_cfg.keys( ) and data < filter_cfg['greater_than']: skip_datapoint = True break elif 'less_than' in filter_cfg.keys( ) and data > filter_cfg['less_than']: skip_datapoint = True break if skip_datapoint: continue logging.info('DATAPOINT %d' % (num)) for f in print_fields: data = datapoint[f] logging.info('Field %s:' % (f)) print data grasp_2d = Grasp2D(Point(image.center), 0, datapoint['hand_poses'][2]) vis2d.figure() if field_type == RenderMode.RGBD: vis2d.subplot(1, 2, 1) vis2d.imshow(image.color) vis2d.grasp(grasp_2d, width=gripper_width_px) vis2d.subplot(1, 2, 2) vis2d.imshow(image.depth) vis2d.grasp(grasp_2d, width=gripper_width_px) elif field_type == RenderMode.GD: vis2d.subplot(1, 2, 1) vis2d.imshow(image.gray) vis2d.grasp(grasp_2d, width=gripper_width_px) vis2d.subplot(1, 2, 2) vis2d.imshow(image.depth) vis2d.grasp(grasp_2d, width=gripper_width_px) else: vis2d.imshow(image) vis2d.grasp(grasp_2d, width=gripper_width_px) vis2d.title('Datapoint %d: %s' % (ind, field_type)) vis2d.show() num += 1