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
