def get_data(self, index, return_raw_cloud=False):
color = np.array(Image.open(self.colorpath[index]), dtype=np.float32) / 255.0
depth = np.array(Image.open(self.depthpath[index]))
seg = np.array(Image.open(self.labelpath[index]))
meta = scio.loadmat(self.metapath[index])
scene = self.scenename[index]
try:
intrinsic = meta['intrinsic_matrix']
factor_depth = meta['factor_depth']
except Exception as e:
print(repr(e))
print(scene)
camera = CameraInfo(1280.0, 720.0, intrinsic[0][0], intrinsic[1][1], intrinsic[0][2], intrinsic[1][2], factor_depth)
# generate cloud
cloud = create_point_cloud_from_depth_image(depth, camera, organized=True)
# get valid points
depth_mask = (depth > 0)
seg_mask = (seg > 0)
if self.remove_outlier:
camera_poses = np.load(os.path.join(self.root, 'scenes', scene, self.camera, 'camera_poses.npy'))
align_mat = np.load(os.path.join(self.root, 'scenes', scene, self.camera, 'cam0_wrt_table.npy'))
trans = np.dot(align_mat, camera_poses[self.frameid[index]])
workspace_mask = get_workspace_mask(cloud, seg, trans=trans, organized=True, outlier=0.02)
mask = (depth_mask & workspace_mask)
else:
mask = depth_mask
cloud_masked = cloud[mask]
color_masked = color[mask]
seg_masked = seg[mask]
if return_raw_cloud:
return cloud_masked, color_masked
# sample points
if len(cloud_masked) >= self.num_points:
idxs = np.random.choice(len(cloud_masked), self.num_points, replace=False)
else:
idxs1 = np.arange(len(cloud_masked))
idxs2 = np.random.choice(len(cloud_masked), self.num_points-len(cloud_masked), replace=True)
idxs = np.concatenate([idxs1, idxs2], axis=0)
cloud_sampled = cloud_masked[idxs]
color_sampled = color_masked[idxs]
ret_dict = {}
ret_dict['point_clouds'] = cloud_sampled.astype(np.float32)
ret_dict['cloud_colors'] = color_sampled.astype(np.float32)
return ret_dict
def get_and_process_data(data_dir):
# load data
color = np.array(Image.open(os.path.join(data_dir, 'color.png')),
dtype=np.float32) / 255.0
depth = np.array(Image.open(os.path.join(data_dir, 'depth.png')))
workspace_mask = np.array(
Image.open(os.path.join(data_dir, 'workspace_mask.png')))
meta = scio.loadmat(os.path.join(data_dir, 'meta.mat'))
intrinsic = meta['intrinsic_matrix']
factor_depth = meta['factor_depth']
# generate cloud
camera = CameraInfo(1280.0, 720.0, intrinsic[0][0], intrinsic[1][1],
intrinsic[0][2], intrinsic[1][2], factor_depth)
cloud = create_point_cloud_from_depth_image(depth, camera, organized=True)
# get valid points
mask = (workspace_mask & (depth > 0))
cloud_masked = cloud[mask]
color_masked = color[mask]
# sample points
if len(cloud_masked) >= cfgs.num_point:
idxs = np.random.choice(len(cloud_masked),
cfgs.num_point,
replace=False)
else:
idxs1 = np.arange(len(cloud_masked))
idxs2 = np.random.choice(len(cloud_masked),
cfgs.num_point - len(cloud_masked),
replace=True)
idxs = np.concatenate([idxs1, idxs2], axis=0)
cloud_sampled = cloud_masked[idxs]
color_sampled = color_masked[idxs]
# convert data
cloud = o3d.geometry.PointCloud()
cloud.points = o3d.utility.Vector3dVector(cloud_masked.astype(np.float32))
cloud.colors = o3d.utility.Vector3dVector(color_masked.astype(np.float32))
end_points = dict()
cloud_sampled = torch.from_numpy(cloud_sampled[np.newaxis].astype(
np.float32))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
cloud_sampled = cloud_sampled.to(device)
end_points['point_clouds'] = cloud_sampled
end_points['cloud_colors'] = color_sampled
return end_points, cloud
def _get_one_annid_data(self, scene_index, view_index,return_raw_cloud=False):
# import ipdb; ipdb.set_trace()
ret_dict = dict(
scene_indexs=self.real_ids[scene_index],
annid_offset=self.annoffset[scene_index][view_index],
)
data_path = self.pointpath[scene_index][view_index]
if os.path.exists(data_path):
ret_dict['point_clouds'] = np.load(data_path)
return ret_dict
# color = np.array(Image.open(self.colorpath[scene_index][view_index]), dtype=np.float32) / 255.0
depth = np.array(Image.open(self.depthpath[scene_index][view_index]))
seg = np.array(Image.open(self.labelpath[scene_index][view_index]))
meta = scio.loadmat(self.metapath[scene_index][view_index])
scene = self.scenename[scene_index][view_index]
try:
intrinsic = meta['intrinsic_matrix']
factor_depth = meta['factor_depth']
except Exception as e:
print(repr(e))
print(scene)
camera = CameraInfo(1280.0, 720.0, intrinsic[0][0], intrinsic[1][1], intrinsic[0][2], intrinsic[1][2], factor_depth)
# generate cloud
cloud = create_point_cloud_from_depth_image(depth, camera, organized=True)
# get valid points
depth_mask = (depth > 0)
seg_mask = (seg > 0)
if self.remove_outlier:
camera_poses = np.load(os.path.join(self.root, 'scenes', scene, self.camera, 'camera_poses.npy'))
align_mat = np.load(os.path.join(self.root, 'scenes', scene, self.camera, 'cam0_wrt_table.npy'))
trans = np.dot(align_mat, camera_poses[self.frameid[scene_index][view_index]])
workspace_mask = get_workspace_mask(cloud, seg, trans=trans, organized=True, outlier=0.02)
mask = (depth_mask & workspace_mask)
else:
mask = depth_mask
cloud_masked = cloud[mask]
# color_masked = color[mask]
seg_masked = seg[mask]
if return_raw_cloud:
return cloud_masked #, color_masked
# sample points
if len(cloud_masked) >= self.num_points:
idxs = np.random.choice(len(cloud_masked), self.num_points, replace=False)
else:
idxs1 = np.arange(len(cloud_masked))
idxs2 = np.random.choice(len(cloud_masked), self.num_points-len(cloud_masked), replace=True)
idxs = np.concatenate([idxs1, idxs2], axis=0)
cloud_sampled = cloud_masked[idxs].astype(np.float32)
# color_sampled = color_masked[idxs]
data_dir = os.path.split(data_path)[0]
if not os.path.exists(data_dir):
os.mkdir(data_dir)
# import ipdb; ipdb.set_trace()
np.save(data_path, cloud_sampled)
ret_dict['point_clouds'] = cloud_sampled
# ret_dict = dict(
# scene_indexs=self.real_ids[scene_index],
# annid_offset=self.annoffset[scene_index][view_index],
# point_clouds=cloud_sampled.astype(np.float32),
# )
return ret_dict
def get_data_label(self, index):
color = np.array(Image.open(self.colorpath[index]),
dtype=np.float32) / 255.0
depth = np.array(Image.open(self.depthpath[index]))
seg = np.array(Image.open(self.labelpath[index]))
meta = scio.loadmat(self.metapath[index])
scene = self.scenename[index]
try:
obj_idxs = meta['cls_indexes'].flatten().astype(np.int32)
poses = meta['poses']
intrinsic = meta['intrinsic_matrix']
factor_depth = meta['factor_depth']
except Exception as e:
print(repr(e))
print(scene)
camera = CameraInfo(1280.0, 720.0, intrinsic[0][0], intrinsic[1][1],
intrinsic[0][2], intrinsic[1][2], factor_depth)
# generate cloud
cloud = create_point_cloud_from_depth_image(depth,
camera,
organized=True)
# get valid points
depth_mask = (depth > 0)
seg_mask = (seg > 0)
if self.remove_outlier:
camera_poses = np.load(
os.path.join(self.root, 'scenes', scene, self.camera,
'camera_poses.npy'))
align_mat = np.load(
os.path.join(self.root, 'scenes', scene, self.camera,
'cam0_wrt_table.npy'))
trans = np.dot(align_mat, camera_poses[self.frameid[index]])
workspace_mask = get_workspace_mask(cloud,
seg,
trans=trans,
organized=True,
outlier=0.02)
mask = (depth_mask & workspace_mask)
else:
mask = depth_mask
cloud_masked = cloud[mask]
color_masked = color[mask]
seg_masked = seg[mask]
# sample points
if len(cloud_masked) >= self.num_points:
idxs = np.random.choice(len(cloud_masked),
self.num_points,
replace=False)
else:
idxs1 = np.arange(len(cloud_masked))
idxs2 = np.random.choice(len(cloud_masked),
self.num_points - len(cloud_masked),
replace=True)
idxs = np.concatenate([idxs1, idxs2], axis=0)
cloud_sampled = cloud_masked[idxs]
color_sampled = color_masked[idxs]
seg_sampled = seg_masked[idxs]
objectness_label = seg_sampled.copy()
objectness_label[objectness_label > 1] = 1
object_poses_list = []
grasp_points_list = []
grasp_offsets_list = []
grasp_scores_list = []
grasp_tolerance_list = []
for i, obj_idx in enumerate(obj_idxs):
if obj_idx not in self.valid_obj_idxs:
continue
if (seg_sampled == obj_idx).sum() < 50:
continue
object_poses_list.append(poses[:, :, i])
points, offsets, scores, tolerance = self.grasp_labels[obj_idx]
collision = self.collision_labels[scene][i] #(Np, V, A, D)
# remove invisible grasp points
if self.remove_invisible:
visible_mask = remove_invisible_grasp_points(
cloud_sampled[seg_sampled == obj_idx],
points,
poses[:, :, i],
th=0.01)
points = points[visible_mask]
offsets = offsets[visible_mask]
scores = scores[visible_mask]
tolerance = tolerance[visible_mask]
collision = collision[visible_mask]
idxs = np.random.choice(len(points),
min(max(int(len(points) / 4), 300),
len(points)),
replace=False)
grasp_points_list.append(points[idxs])
grasp_offsets_list.append(offsets[idxs])
collision = collision[idxs].copy()
scores = scores[idxs].copy()
scores[collision] = 0
grasp_scores_list.append(scores)
tolerance = tolerance[idxs].copy()
tolerance[collision] = 0
grasp_tolerance_list.append(tolerance)
if self.augment:
cloud_sampled, object_poses_list = self.augment_data(
cloud_sampled, object_poses_list)
ret_dict = {}
ret_dict['point_clouds'] = cloud_sampled.astype(np.float32)
ret_dict['cloud_colors'] = color_sampled.astype(np.float32)
ret_dict['objectness_label'] = objectness_label.astype(np.int64)
ret_dict['object_poses_list'] = object_poses_list
ret_dict['grasp_points_list'] = grasp_points_list
ret_dict['grasp_offsets_list'] = grasp_offsets_list
ret_dict['grasp_labels_list'] = grasp_scores_list
ret_dict['grasp_tolerance_list'] = grasp_tolerance_list
return ret_dict