def generate_write(self, dir_path, mesh_file):
# prepare mesh and +generate Grasps
mesh_processor = mp.MeshProcessor(os.path.join(dir_path, mesh_file),
self.mesh_out)
mesh_processor.generate_graspable(self.egad_config)
obj = GraspableObject3D(mesh_processor.sdf, mesh_processor.mesh)
self.collision_checker = GraspCollisionChecker(self.gripper)
self.collision_checker.set_graspable_object(obj)
spd_list = list()
# read in the stable poses of the mesh
stable_poses = mesh_processor.stable_poses
# find aligned grasps for each stable pose
for i, stable_pose in enumerate(stable_poses):
# render images if stable pose is valid
if stable_pose.p > self.stable_pose_min_p:
# add to total mesh+stableposes list
T_table_obj = self.setup_Table(stable_pose, obj)
grasp_data = self.get_grasps(stable_pose, obj)
spd_list.append(
StablePoseData(mesh_file.replace(".obj", "_proc.obj"),
T_table_obj.matrix, grasp_data))
# write pickle file
print("found %d grasps, writing to file" % len(spd_list))
return spd_list
def render_data(self):
logging.basicConfig(level=logging.WARNING)
phi_offsets = self._generate_phi_offsets()
old_coll = []
new_coll = []
cnt = 0
for dataset in self.datasets:
logging.info('Generating data for dataset %s' % dataset.name)
object_keys = dataset.object_keys
for obj_key in object_keys:
if self.cur_obj_label % 10 == 0:
logging.info("Object number: %d" % self.cur_obj_label)
if ONLY_TEST and self.cur_obj_label not in self.test_obj_labels:
self.cur_obj_label += 1
continue
obj = dataset[obj_key]
grasps = dataset.grasps(obj.key, gripper=self.gripper.name)
# Load grasp metrics
stable_poses = dataset.stable_poses(obj.key)
collision_checker = GraspCollisionChecker(self.gripper)
collision_checker.set_graspable_object(obj)
# Iterate through stable poses
for i, stable_pose in enumerate(stable_poses):
if not stable_pose.p > 0.0:
continue
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=0.005).as_frames('obj', 'table')
T_table_obj = T_obj_table.inverse()
aligned_grasps = [
grasp.perpendicular_table(stable_pose)
for grasp in grasps
]
collision_checker.set_table(self._table_mesh_filename,
T_table_obj)
for grasp in aligned_grasps:
_alpha, _beta, _ = grasp.grasp_angles_from_stp_z(
stable_pose)
if np.abs(_beta) > np.deg2rad(5):
continue
old_coll.append(
self.is_grasp_collision_free(
grasp, collision_checker, phi_offsets))
new_coll.append(
self.is_grasp_collision_free(
grasp, collision_checker, [0]))
if len(old_coll) == 1000:
np.save(
"{}old_coll_{:05d}.npy".format(
self.output_dir, cnt), old_coll)
np.save(
"{}new_coll_{:05d}.npy".format(
self.output_dir, cnt), new_coll)
old_coll = []
new_coll = []
cnt += 1
self.cur_obj_label += 1
def render_data(self):
logging.basicConfig(level=logging.WARNING)
for dataset in self.datasets:
logging.info('Generating data for dataset %s' % dataset.name)
object_keys = dataset.object_keys
for obj_key in object_keys:
self.tensor_datapoint['obj_labels'] = self.cur_obj_label
if self.cur_obj_label % 10 == 0:
logging.info("Object number: %d" % self.cur_obj_label)
self.obj = dataset[obj_key]
grasps = dataset.grasps(self.obj.key, gripper=self.gripper.name)
# Load grasp metrics
grasp_metrics = dataset.grasp_metrics(self.obj.key,
grasps,
gripper=self.gripper.name)
# setup collision checker
collision_checker = GraspCollisionChecker(self.gripper)
collision_checker.set_graspable_object(self.obj)
# read in the stable poses of the mesh
stable_poses = dataset.stable_poses(self.obj.key)
# Iterate through stable poses
for i, stable_pose in enumerate(stable_poses):
if not stable_pose.p > self._stable_pose_min_p:
continue
self.tensor_datapoint['pose_labels'] = self.cur_pose_label
# setup table in collision checker
T_obj_stp = stable_pose.T_obj_table.as_frames('obj', 'stp')
T_obj_table = self.obj.mesh.get_T_surface_obj(T_obj_stp,
delta=self._table_offset).as_frames('obj', 'table')
T_table_obj = T_obj_table.inverse()
T_obj_stp = self.obj.mesh.get_T_surface_obj(T_obj_stp)
collision_checker.set_table(self._table_mesh_filename, T_table_obj)
# sample images from random variable
T_table_obj = RigidTransform(from_frame='table', to_frame='obj')
scene_objs = {'table': SceneObject(self.table_mesh, T_table_obj)}
# Set up image renderer
samples = self.render_images(scene_objs,
stable_pose,
self.config['images_per_stable_pose'],
camera_config=self._camera_configs())
for sample in samples:
self.save_samples(sample, grasps, T_obj_stp, collision_checker, grasp_metrics, stable_pose)
self.cur_pose_label += 1
gc.collect()
# next stable pose
self.cur_obj_label += 1
# next object
# Save dataset
self.tensor_dataset.flush()
def filter_grasps_generic(graspable, grasps, gripper, progress_reporter=lambda x: None):
progress_reporter(0)
collision_checker = GraspCollisionChecker(gripper)
collision_checker.set_graspable_object(graspable)
collision_free_grasps = []
colliding_grasps = []
for k, grasp in enumerate(grasps):
progress_reporter(float(k) / len(grasps))
collision_free = False
for rot_idx in range(0, consts.GENERAL_COLLISION_CHECKING_NUM_OFFSETS):
rotated_grasp = grasp.grasp_y_axis_offset(rot_idx * consts.GENERAL_COLLISION_CHECKING_PHI)
collides = collision_checker.collides_along_approach(rotated_grasp,
consts.APPROACH_DIST,
consts.DELTA_APPROACH)
if not collides:
collision_free = True
collision_free_grasps.append(grasp)
break
if not collision_free:
colliding_grasps.append(grasp)
return collision_free_grasps, colliding_grasps
def filter_grasps_stbp(graspable, grasps, gripper, stable_poses, progress_reporter=lambda x: None):
progress_reporter(0)
collision_checker = GraspCollisionChecker(gripper)
collision_checker.set_graspable_object(graspable)
stbp_grasps_indices = []
stbp_grasps_aligned = []
for k, stable_pose in enumerate(stable_poses):
# set up collision checker with table
T_obj_stp = RigidTransform(rotation=stable_pose.r, from_frame='obj', to_frame='stp')
T_obj_table = graspable.mesh.get_T_surface_obj(T_obj_stp,
delta=consts.COLLISION_CONFIG['table_offset']).as_frames('obj', 'table')
T_table_obj = T_obj_table.inverse()
collision_checker.set_table(consts.COLLISION_CONFIG['table_mesh_filename'], T_table_obj)
aligned_grasps = [grasp.perpendicular_table(stable_pose) for grasp in grasps]
this_stbp_grasps_indices = []
this_stbp_grasps_aligned = []
for idx, aligned_grasp in enumerate(aligned_grasps):
progress_reporter(float(idx) / (len(grasps) * len(stable_poses)) + float(k) / len(stable_poses))
_, grasp_approach_table_angle, _ = aligned_grasp.grasp_angles_from_stp_z(stable_pose)
perpendicular_table = (np.abs(grasp_approach_table_angle) < consts.MAX_GRASP_APPROACH_TABLE_ANGLE)
if not perpendicular_table:
continue
# check whether any valid approach directions are collision free
for phi_offset in consts.PHI_OFFSETS:
rotated_grasp = aligned_grasp.grasp_y_axis_offset(phi_offset)
collides = collision_checker.collides_along_approach(rotated_grasp, consts.APPROACH_DIST, consts.DELTA_APPROACH)
if not collides:
this_stbp_grasps_indices.append(idx)
this_stbp_grasps_aligned.append(aligned_grasp)
break
stbp_grasps_indices.append(this_stbp_grasps_indices)
stbp_grasps_aligned.append(this_stbp_grasps_aligned)
return stbp_grasps_indices, stbp_grasps_aligned
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 save_dexnet_objects(output_path, database, target_object_keys, config,
pointers, num):
slice_dataset = False
files = []
if not os.path.exists(output_path):
os.mkdir(output_path)
for each_file in os.listdir(output_path):
os.remove(output_path + '/' + each_file)
gripper = RobotGripper.load(config['gripper'])
# 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']
stable_pose_min_p = config['stable_pose_min_p']
table_mesh_filename = coll_check_params['table_mesh_filename']
if not os.path.isabs(table_mesh_filename):
table_mesh_filename = os.path.join(DATA_DIR, 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()
dataset_names = target_object_keys.keys()
datasets = [database.dataset(dn) for dn in dataset_names]
if slice_dataset:
datasets = [dataset.subset(0, 100) for dataset in datasets]
start = 0
for dataset in datasets:
target_object_keys[dataset.name] = []
end = start + len(dataset.object_keys)
for cnt, _id in enumerate(pointers.obj_ids):
if _id >= end or _id < start:
continue
target_object_keys[dataset.name].append(dataset.object_keys[_id -
start])
files.append(
tuple([
dataset.object_keys[_id - start], pointers.tensor[cnt],
pointers.array[cnt], pointers.depths[cnt]
]))
start += end
print(target_object_keys)
print("target object keys:", len(target_object_keys['3dnet']),
len(target_object_keys['kit']))
files = np.array(files,
dtype=[('Object_id', (np.str_, 40)), ('Tensor', int),
('Array', int), ('Depth', float)])
# Precompute set of valid grasps
candidate_grasps_dict = {}
counter = 0
for dataset in datasets:
for obj in dataset:
if obj.key not in target_object_keys[dataset.name]:
continue
print("Object in subset")
# Initiate 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)
try:
stable_pose = stable_poses[pointers.pose_num[counter]]
except IndexError:
print(
"Problems with reading pose. Tensor %d, Array %d, Pose %d"
% (pointers.tensor[counter], pointers.array[counter],
pointers.pose_num[counter]))
print("Stable poses:", stable_poses)
print("Pointers pose:", pointers.pose_num[counter])
counter += 1
print("Continue.")
continue
print("Read in stable pose")
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)
# align grasps with the table
aligned_grasps = [
grasp.perpendicular_table(stable_pose) for grasp in grasps
]
i = 0
found = False
if len(aligned_grasps) < pointers.grasp_num[counter]:
raise IndexError
print("pointers grasp num", pointers.grasp_num[counter])
print("tensor", pointers.tensor[counter])
print("array", pointers.array[counter])
# Check grasp validity
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
if i == pointers.grasp_num[counter]:
found, contact_points = aligned_grasp.close_fingers(obj)
print("Contact points", contact_points)
if not found:
print("Could not find contact points. continue.")
break
else:
print("Original metric: ", pointers.metrics[counter])
print(
"Metrics mapped point: ",
dataset.grasp_metrics(obj.key, [aligned_grasp],
gripper=gripper.name))
candidate_grasps_dict[obj.key][stable_pose.id].append(
GraspInfo(aligned_grasp, collision_free, [
contact_points[0].point,
contact_points[1].point
]))
# logging.info('Grasp %d' % (aligned_grasp.id))
# vis.figure()
# vis.gripper_on_object(gripper, aligned_grasp, obj, stable_pose.T_obj_world, plot_table=False)
# vis.show()
break
i += 1
if found:
# Save files
print("Saving file: ", obj.key)
savefile = ObjFile(DATA_DIR + "/meshes/dexnet/" + obj.key +
".obj")
savefile.write(obj.mesh)
# print("Vertices:", obj.mesh.vertices.shape)
# print("Triangles:", obj.mesh.triangles.shape)
mesh = o3d.geometry.TriangleMesh(
o3d.utility.Vector3dVector(obj.mesh.vertices),
o3d.utility.Vector3iVector(obj.mesh.triangles))
o3d.io.write_triangle_mesh(
DATA_DIR + '/PerfectPredictions/3D_meshes/' +
"%05d_%03d.ply" %
(pointers.tensor[counter], pointers.array[counter]), mesh)
# Save stable poses
save_stp = StablePoseFile(DATA_DIR + "/meshes/dexnet/" +
obj.key + ".stp")
save_stp.write(stable_poses)
# Save candidate grasp info
pkl.dump(
candidate_grasps_dict[obj.key],
open(DATA_DIR + "/meshes/dexnet/" + obj.key + ".pkl",
'wb'))
# Save grasp 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)
write_metrics = json.dumps(grasp_metrics)
f = open(DATA_DIR + "/meshes/dexnet/" + obj.key + ".json", "w")
f.write(write_metrics)
f.close()
counter += 1
if num is not None and counter >= num:
break
with open(DATA_DIR + '/meshes/dexnet/files.csv', 'w') as csv_file:
csv_writer = csv.writer(csv_file, delimiter=',')
for point in files:
csv_writer.writerow(point)
def render_data(self):
logging.basicConfig(level=logging.WARNING)
for dataset in self.datasets:
logging.info('Generating data for dataset %s' % dataset.name)
object_keys = dataset.object_keys
for obj_key in object_keys:
logging.info("Object number: %d" % self.cur_obj_label)
self.obj = dataset[obj_key]
grasps = dataset.grasps(self.obj.key,
gripper=self.gripper.name)
positive_grasps = self.get_positive_grasps(dataset, grasps)
# Load grasp metrics
grasp_metrics = dataset.grasp_metrics(
self.obj.key, grasps, gripper=self.gripper.name)
# setup collision checker
collision_checker = GraspCollisionChecker(self.gripper)
collision_checker.set_graspable_object(self.obj)
# read in the stable poses of the mesh
stable_poses = dataset.stable_poses(self.obj.key)
# Iterate through stable poses
for i, stable_pose in enumerate(stable_poses):
if not stable_pose.p > self._stable_pose_min_p:
continue
# setup table in collision checker
T_obj_stp = stable_pose.T_obj_table.as_frames('obj', 'stp')
T_obj_table = self.obj.mesh.get_T_surface_obj(
T_obj_stp,
delta=self._table_offset).as_frames('obj', 'table')
T_table_obj = T_obj_table.inverse()
T_obj_stp = self.obj.mesh.get_T_surface_obj(T_obj_stp)
collision_checker.set_table(self._table_mesh_filename,
T_table_obj)
# sample images from random variable
T_table_obj = RigidTransform(from_frame='table',
to_frame='obj')
scene_objs = {
'table': SceneObject(self.table_mesh, T_table_obj)
}
# Set up image renderer
for _ in range(self.config['images_per_stable_pose']):
elev_angle = np.random.choice(
self.camera_distributions)
camera_config = self._camera_configs()
camera_config['min_elev'] = elev_angle * 5.0
camera_config['max_elev'] = (elev_angle + 1) * 5.0
sample = self.render_images(
scene_objs,
stable_pose,
1,
camera_config=camera_config)
self.save_samples(sample, grasps, T_obj_stp,
collision_checker, grasp_metrics)
# Get camera elevation angle for current sample
elev_angle = sample.camera.elev * 180 / np.pi
elev_bar = int(elev_angle // 5)
# Sample number of positive images from distribution
number_positive_images = np.random.choice(
self.grasp_upsample_distribuitions[elev_bar])
# Render only positive images
new_config = self._camera_configs()
new_config['min_elev'] = elev_angle
new_config['max_elev'] = elev_angle
positive_render_samples = self.render_images(
scene_objs,
stable_pose,
number_positive_images,
camera_config=new_config)
if type(positive_render_samples) == list:
for pos_sample in positive_render_samples:
self.save_samples(pos_sample,
positive_grasps,
T_obj_stp,
collision_checker,
grasp_metrics,
only_positive=True)
else:
self.save_samples(positive_render_samples,
positive_grasps,
T_obj_stp,
collision_checker,
grasp_metrics,
only_positive=True)
self.cur_pose_label += 1
gc.collect()
# next stable pose
self.cur_obj_label += 1
# next object
# Save dataset
self.tensor_dataset.flush()
cur_image_label = 0
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:
print("Object number: ", cur_obj_label)
obj = dataset[obj_key]
grasps = dataset.grasps(obj.key, gripper=gripper.name)
# 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)
# Iterate through stable poses
for i, stable_pose in enumerate(stable_poses):
if not stable_pose.p > stable_pose_min_p:
continue
# 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
# setup table in collision checker
class DexnetGetGraspdata:
def __init__(self, mesh_out, sp_min_p, egad_path, dexnet_path):
## VARIABLES
# min stable pose probability
self.stable_pose_min_p = sp_min_p
# mesh output dir
self.mesh_out = mesh_out
## SETUP CONFIG
self.egad_config = None
self.table_offset = None
self.table_mesh_filename = None
self.approach_dist = None
self.delta_approach = None
self.max_grasp_approach_table_angle = None
self.phi_offsets = None
self.setup_config(egad_path, dexnet_path)
## DEXNET SETUP
# setup collision checker and sampler
os.chdir(str(Path(dexnet.__file__).resolve().parent.parent.parent))
self.gripper = RobotGripper.load(
'yumi_metal_spline', gripper_dir=self.egad_config['gripper_dir'])
self.sampler = AntipodalGraspSampler(self.gripper, self.egad_config)
# setup Grasp Quality Function
self.metric_config = GraspQualityConfigFactory().create_config(
self.egad_config['metrics']['ferrari_canny'])
# has to be reinitiated to avoid errors
self.collision_checker = None
self.grasp_data_list = None
# setup config
def setup_config(self, egad_path, dexnet_path):
# Use local config file
self.egad_config = YamlConfig(
os.path.join(egad_path, "scripts/cfg/dexnet_api_settings.yaml"))
# setup collision params for alignment
coll_vis_config = YamlConfig(
os.path.join(dexnet_path, "cfg/tools/generate_gqcnn_dataset.yaml"))
coll_check_params = coll_vis_config['collision_checking']
self.table_offset = coll_check_params['table_offset']
self.table_mesh_filename = coll_check_params['table_mesh_filename']
self.approach_dist = coll_check_params['approach_dist']
self.delta_approach = coll_check_params['delta_approach']
# paramet for alignement of grasps
table_alignment_params = coll_vis_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'])
num_grasp_approach_samples = table_alignment_params[
'num_approach_offset_samples']
self.max_grasp_approach_table_angle = np.deg2rad(
table_alignment_params['max_approach_table_angle'])
self.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:
self.phi_offsets.append(phi)
phi += phi_inc
# get list of grasps for given stable_pose and obj
def get_grasps(self, stable_pose, obj):
# align grasps with the table
grasps = self.sampler.generate_grasps(
obj, max_iter=self.egad_config['max_grasp_sampling_iters'])
quality_fn = GraspQualityFunctionFactory.create_quality_function(
obj, self.metric_config)
# https://github.com/BerkeleyAutomation/dex-net/blob/0f63f706668815e96bdeea16e14d2f2b266f9262/src/dexnet/grasping/quality.py
aligned_grasps = [
grasp.perpendicular_table(stable_pose) for grasp in grasps
]
self.grasp_data_list = list()
# check grasp validity
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) <
self.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 self.phi_offsets:
rotated_grasp = aligned_grasp.grasp_y_axis_offset(phi_offset)
collides = self.collision_checker.collides_along_approach(
rotated_grasp, self.approach_dist, self.delta_approach)
if not collides:
collision_free = True
break
if (collision_free and aligned_grasp.close_fingers(obj)[0]):
quality = convert_quality(quality_fn(aligned_grasp).quality)
self.grasp_data_list.append(
GraspData(aligned_grasp.T_grasp_obj.matrix,
aligned_grasp.close_fingers(obj)[1][0].point,
aligned_grasp.close_fingers(obj)[1][1].point,
quality))
return self.grasp_data_list
# setup table data and in collision checker, return table to obj trans matrix
def setup_Table(self, stable_pose, obj):
# get table trans matrix
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=self.table_offset).as_frames('obj', 'table')
T_table_obj = T_obj_table.inverse()
# setup table in collision checker
self.collision_checker.set_table(self.table_mesh_filename, T_table_obj)
return T_table_obj
# Generate Stable Poses, Grasps and write to pickle
def generate_write(self, dir_path, mesh_file):
# prepare mesh and +generate Grasps
mesh_processor = mp.MeshProcessor(os.path.join(dir_path, mesh_file),
self.mesh_out)
mesh_processor.generate_graspable(self.egad_config)
obj = GraspableObject3D(mesh_processor.sdf, mesh_processor.mesh)
self.collision_checker = GraspCollisionChecker(self.gripper)
self.collision_checker.set_graspable_object(obj)
spd_list = list()
# read in the stable poses of the mesh
stable_poses = mesh_processor.stable_poses
# find aligned grasps for each stable pose
for i, stable_pose in enumerate(stable_poses):
# render images if stable pose is valid
if stable_pose.p > self.stable_pose_min_p:
# add to total mesh+stableposes list
T_table_obj = self.setup_Table(stable_pose, obj)
grasp_data = self.get_grasps(stable_pose, obj)
spd_list.append(
StablePoseData(mesh_file.replace(".obj", "_proc.obj"),
T_table_obj.matrix, grasp_data))
# write pickle file
print("found %d grasps, writing to file" % len(spd_list))
return spd_list