def __init__(self,
data_paths,
input_transform=None,
target_transform=None,
data_root='/',
explicit_rotation=-1,
ignore_label=255,
return_transformation=False,
augment_data=False,
config=None,
**kwargs):
if explicit_rotation > 0:
raise NotImplementedError
self.augment_data = augment_data
self.config = config
super().__init__(data_paths,
input_transform=input_transform,
target_transform=target_transform,
cache=cache,
data_root=data_root,
explicit_rotation=explicit_rotation,
ignore_mask=ignore_label,
return_transformation=return_transformation)
self.sparse_voxelizer = Voxelizer(
voxel_size=self.VOXEL_SIZE,
clip_bound=self.CLIP_BOUND,
use_augmentation=augment_data,
scale_augmentation_bound=self.SCALE_AUGMENTATION_BOUND,
rotation_augmentation_bound=self.ROTATION_AUGMENTATION_BOUND,
translation_augmentation_ratio_bound=self.
TRANSLATION_AUGMENTATION_RATIO_BOUND,
ignore_label=ignore_label)
# map labels not evaluated to ignore_label
if self.IGNORE_LABELS is not None:
label_map = {}
n_used = 0
for l in range(self.NUM_LABELS):
if l in self.IGNORE_LABELS:
label_map[l] = self.ignore_mask
elif l not in self.INSTANCE_LABELS:
label_map[l] = n_used
n_used += 1
for l in self.INSTANCE_LABELS:
label_map[l] = n_used
n_used += 1
label_map[self.ignore_mask] = self.ignore_mask
self.label_map = label_map
self.NUM_LABELS -= len(self.IGNORE_LABELS)
def __init__(self,
data_paths,
prevoxel_transform=None,
input_transform=None,
target_transform=None,
data_root='/',
ignore_label=255,
return_transformation=False,
augment_data=False,
config=None,
**kwargs):
self.augment_data = augment_data
self.config = config
VoxelizationDatasetBase.__init__(
self,
data_paths,
prevoxel_transform=prevoxel_transform,
input_transform=input_transform,
target_transform=target_transform,
cache=cache,
data_root=data_root,
ignore_mask=ignore_label,
return_transformation=return_transformation)
# Prevoxel transformations
self.voxelizer = Voxelizer(
voxel_size=self.VOXEL_SIZE,
clip_bound=self.CLIP_BOUND,
use_augmentation=augment_data,
scale_augmentation_bound=self.SCALE_AUGMENTATION_BOUND,
rotation_augmentation_bound=self.ROTATION_AUGMENTATION_BOUND,
translation_augmentation_ratio_bound=self.
TRANSLATION_AUGMENTATION_RATIO_BOUND,
ignore_label=ignore_label)
# map labels not evaluated to ignore_label
label_map = {}
n_used = 0
for l in range(self.NUM_LABELS):
if l in self.IGNORE_LABELS:
label_map[l] = self.ignore_mask
else:
label_map[l] = n_used
n_used += 1
label_map[self.ignore_mask] = self.ignore_mask
self.label_map = label_map
self.NUM_LABELS -= len(self.IGNORE_LABELS)
class VoxelizationDataset(VoxelizationDatasetBase):
"""This dataset loads RGB point clouds and their labels as a list of points
and voxelizes the pointcloud with sufficient data augmentation.
"""
# Voxelization arguments
VOXEL_SIZE = 0.05 # 5cm
# Coordinate Augmentation Arguments: Unlike feature augmentation, coordinate
# augmentation has to be done before voxelization
SCALE_AUGMENTATION_BOUND = (0.9, 1.1)
ROTATION_AUGMENTATION_BOUND = ((-np.pi / 6, np.pi / 6), (-np.pi, np.pi),
(-np.pi / 6, np.pi / 6))
TRANSLATION_AUGMENTATION_RATIO_BOUND = ((-0.2, 0.2), (-0.05, 0.05), (-0.2,
0.2))
ELASTIC_DISTORT_PARAMS = None
# MISC.
PREVOXELIZATION_VOXEL_SIZE = None
# Augment coords to feats
AUGMENT_COORDS_TO_FEATS = False
def __init__(self,
data_paths,
prevoxel_transform=None,
input_transform=None,
target_transform=None,
data_root='/',
ignore_label=255,
return_transformation=False,
augment_data=False,
config=None,
**kwargs):
self.augment_data = augment_data
self.config = config
VoxelizationDatasetBase.__init__(
self,
data_paths,
prevoxel_transform=prevoxel_transform,
input_transform=input_transform,
target_transform=target_transform,
cache=cache,
data_root=data_root,
ignore_mask=ignore_label,
return_transformation=return_transformation)
# Prevoxel transformations
self.voxelizer = Voxelizer(
voxel_size=self.VOXEL_SIZE,
clip_bound=self.CLIP_BOUND,
use_augmentation=augment_data,
scale_augmentation_bound=self.SCALE_AUGMENTATION_BOUND,
rotation_augmentation_bound=self.ROTATION_AUGMENTATION_BOUND,
translation_augmentation_ratio_bound=self.
TRANSLATION_AUGMENTATION_RATIO_BOUND,
ignore_label=ignore_label)
# map labels not evaluated to ignore_label
label_map = {}
n_used = 0
for l in range(self.NUM_LABELS):
if l in self.IGNORE_LABELS:
label_map[l] = self.ignore_mask
else:
label_map[l] = n_used
n_used += 1
label_map[self.ignore_mask] = self.ignore_mask
self.label_map = label_map
self.NUM_LABELS -= len(self.IGNORE_LABELS)
def _augment_coords_to_feats(self, coords, feats, labels=None):
norm_coords = coords - coords.mean(0)
# color must come first.
if isinstance(coords, np.ndarray):
feats = np.concatenate((feats, norm_coords), 1)
else:
feats = torch.cat((feats, norm_coords), 1)
return coords, feats, labels
def convert_mat2cfl(self, mat):
# Generally, xyz,rgb,label
return mat[:, :3], mat[:, 3:-1], mat[:, -1]
def __getitem__(self, index):
coords, feats, labels, center = self.load_ply(index)
# Downsample the pointcloud with finer voxel size before transformation for memory and speed
if self.PREVOXELIZATION_VOXEL_SIZE is not None:
inds = ME.utils.sparse_quantize(coords /
self.PREVOXELIZATION_VOXEL_SIZE,
return_index=True)
coords = coords[inds]
feats = feats[inds]
labels = labels[inds]
# Prevoxel transformations
if self.prevoxel_transform is not None:
coords, feats, labels = self.prevoxel_transform(
coords, feats, labels)
coords, feats, labels, transformation = self.voxelizer.voxelize(
coords, feats, labels, center=center)
# map labels not used for evaluation to ignore_label
if self.input_transform is not None:
coords, feats, labels = self.input_transform(coords, feats, labels)
if self.target_transform is not None:
coords, feats, labels = self.target_transform(
coords, feats, labels)
if self.IGNORE_LABELS is not None:
labels = np.array([self.label_map[x] for x in labels],
dtype=np.int)
# Use coordinate features if config is set
if self.AUGMENT_COORDS_TO_FEATS:
coords, feats, labels = self._augment_coords_to_feats(
coords, feats, labels)
return_args = [coords, feats, labels]
if self.return_transformation:
return_args.append(transformation.astype(np.float32))
return tuple(return_args)
def main():
pcd = o3d.io.read_point_cloud(INPUT_PCD)
pcd_xyz, pcd_feats = np.asarray(pcd.points), np.asarray(pcd.colors)
print(f'Finished reading {INPUT_PCD}:')
print(f'# points: {pcd_xyz.shape[0]} points')
print(f'volume: {np.prod(pcd_xyz.max(0) - pcd_xyz.min(0))} m^3')
sparse_voxelizer = Voxelizer(voxel_size=0.05)
height = pcd_xyz[:, LOCFEAT_IDX].copy()
height -= np.percentile(height, 0.99)
pcd_feats = np.hstack((pcd_feats, height[:, None]))
preprocess = []
for i in range(7):
start = time.time()
coords, feats, labels, transformation = sparse_voxelizer.voxelize(
pcd_xyz, pcd_feats, None)
preprocess.append(time.time() - start)
print('Voxelization time average: ', np.mean(preprocess[2:]))
coords = ME.utils.batched_coordinates([torch.from_numpy(coords).int()])
feats = torch.from_numpy(feats.astype(np.float32)).to('cuda')
config = get_config()
DatasetClass = load_dataset(config.dataset)
dataloader = initialize_data_loader(DatasetClass,
config,
threads=config.threads,
phase=config.test_phase,
augment_data=False,
shuffle=False,
repeat=False,
batch_size=config.test_batch_size,
limit_numpoints=False)
pipeline_model = load_pipeline(config, dataloader.dataset)
if config.weights.lower() != 'none':
state = torch.load(config.weights)
pipeline_model.load_state_dict(
state['state_dict'], strict=(not config.lenient_weight_loading))
pipeline_model.eval()
evaltime = []
for i in range(7):
start = time.time()
sinput = ME.SparseTensor(feats, coords).to('cuda')
datum = {'sinput': sinput, 'anchor_match_coords': None}
outputs = pipeline_model(datum, False)
evaltime.append(time.time() - start)
print('Network runtime average: ', np.mean(evaltime[2:]))
pred = outputs['detection'][0]
pred_mask = pred[:, -1] > MIN_CONF
pred = pred[pred_mask]
print(f'Detected {pred.shape[0]} instances')
bbox_xyz = pred[:, :6]
bbox_xyz += 0.5
bbox_xyz[:, :3] += 0.5
bbox_xyz[:, 3:] -= 0.5
bbox_xyz[:, 3:] = np.maximum(bbox_xyz[:, 3:], bbox_xyz[:, :3] + 0.1)
bbox_xyz = bbox_xyz.reshape(-1, 3)
bbox_xyz1 = np.hstack((bbox_xyz, np.ones((bbox_xyz.shape[0], 1))))
bbox_xyz = np.linalg.solve(transformation.reshape(4, 4),
bbox_xyz1.T).T[:, :3].reshape(-1, 6)
pred = np.hstack((bbox_xyz, pred[:, 6:]))
pred_pcd = pc_utils.visualize_bboxes(pred[:, :6],
pred[:, 6],
num_points=1000)
mesh = o3d.io.read_triangle_mesh(INPUT_MESH)
mesh.compute_vertex_normals()
pc_utils.visualize_pcd(mesh, pred_pcd)
class SparseVoxelizationDataset(VoxelizationDatasetBase):
"""This dataset loads RGB point clouds and their labels as a list of points
and voxelizes the pointcloud with sufficient data augmentation.
"""
IS_ROTATION_BBOX = False
HAS_GT_BBOX = False
# Voxelization arguments
VOXEL_SIZE = 0.05 # 5cm
# Augmentation arguments
SCALE_AUGMENTATION_BOUND = (0.9, 1.1)
ROTATION_AUGMENTATION_BOUND = ((-np.pi / 6, np.pi / 6), (-np.pi, np.pi),
(-np.pi / 6, np.pi / 6))
TRANSLATION_AUGMENTATION_RATIO_BOUND = ((-0.2, 0.2), (-0.05, 0.05), (-0.2,
0.2))
ELASTIC_DISTORT_PARAMS = None
def __init__(self,
data_paths,
input_transform=None,
target_transform=None,
data_root='/',
explicit_rotation=-1,
ignore_label=255,
return_transformation=False,
augment_data=False,
config=None,
**kwargs):
if explicit_rotation > 0:
raise NotImplementedError
self.augment_data = augment_data
self.config = config
super().__init__(data_paths,
input_transform=input_transform,
target_transform=target_transform,
cache=cache,
data_root=data_root,
explicit_rotation=explicit_rotation,
ignore_mask=ignore_label,
return_transformation=return_transformation)
self.sparse_voxelizer = Voxelizer(
voxel_size=self.VOXEL_SIZE,
clip_bound=self.CLIP_BOUND,
use_augmentation=augment_data,
scale_augmentation_bound=self.SCALE_AUGMENTATION_BOUND,
rotation_augmentation_bound=self.ROTATION_AUGMENTATION_BOUND,
translation_augmentation_ratio_bound=self.
TRANSLATION_AUGMENTATION_RATIO_BOUND,
ignore_label=ignore_label)
# map labels not evaluated to ignore_label
if self.IGNORE_LABELS is not None:
label_map = {}
n_used = 0
for l in range(self.NUM_LABELS):
if l in self.IGNORE_LABELS:
label_map[l] = self.ignore_mask
elif l not in self.INSTANCE_LABELS:
label_map[l] = n_used
n_used += 1
for l in self.INSTANCE_LABELS:
label_map[l] = n_used
n_used += 1
label_map[self.ignore_mask] = self.ignore_mask
self.label_map = label_map
self.NUM_LABELS -= len(self.IGNORE_LABELS)
def convert_mat2cfl(self, mat):
# Generally, xyz, rgb, label
return mat[:, :3], mat[:, 3:-2], mat[:, -2:]
def _augment_elastic_distortion(self, pointcloud):
if self.ELASTIC_DISTORT_PARAMS is not None:
if random.random() < 0.95:
for granularity, magnitude in self.ELASTIC_DISTORT_PARAMS:
pointcloud = elastic_distortion(pointcloud, granularity,
magnitude)
return pointcloud
def __getitem__(self, index):
pointcloud, bboxes, center = self.load_datafile(index)
if self.augment_data and self.config.elastic_distortion:
pointcloud = self._augment_elastic_distortion(pointcloud)
coords, feats, labels = self.convert_mat2cfl(pointcloud)
feats = self._augment_locfeat(coords, feats)
coords, feats, labels, transformation = self.sparse_voxelizer.voxelize(
coords, feats, labels, center=center)
# transform bboxes to match the voxelization transformation.
# TODO(jgwak): Support rotation augmentation.
if self.HAS_GT_BBOX:
bboxes_xyz = np.hstack(
(bboxes[:, :6].reshape(-1, 3), np.ones(
(bboxes.shape[0] * 2, 1))))
bboxes_xyz = bboxes_xyz @ transformation.reshape(4, 4).T
bboxes_xyz = (bboxes_xyz[:, :3] / bboxes_xyz[:, 3, None]).reshape(
-1, 2, 3)
bboxes_xyz = np.hstack((bboxes_xyz.min(1), bboxes_xyz.max(1)))
bboxes = np.hstack((bboxes_xyz, bboxes[:, 6:]))
if labels is None:
semantic_labels = np.ones(
coords.shape[0]) * self.config.ignore_label
instance_labels = np.ones(
coords.shape[0]) * self.config.ignore_label
else:
semantic_labels, instance_labels = labels.T
# map labels not used for evaluation to ignore_label
if self.input_transform is not None:
coords, feats, bboxes = self.input_transform(coords, feats, bboxes)
if self.target_transform is not None:
coords, feats, bboxes = self.target_transform(
coords, feats, bboxes)
if self.IGNORE_LABELS is not None:
semantic_labels = np.array(
[self.label_map[x] for x in semantic_labels], dtype=np.int)
if self.HAS_GT_BBOX:
bboxes[:, -1] = np.array(
[self.label_map[int(x)] for x in bboxes[:, -1]])
# Normalize rotation.
if self.augment_data and self.HAS_GT_BBOX and self.IS_ROTATION_BBOX:
if self.config.normalize_rotation:
rot_bin = np.floor((bboxes[:, 6] + np.pi / 4) / (np.pi / 2))
idxs = list(range(3))
idxs.remove(self.LOCFEAT_IDX)
assert len(idxs) == 2
bbox_size = (bboxes[:, 3:6] - bboxes[:, :3]) / 2
bbox_center = (bboxes[:, 3:6] + bboxes[:, :3]) / 2
bboxes[:, 6] = bboxes[:, 6] - rot_bin * np.pi / 2
is_rotated = np.where(rot_bin % 2)[0]
bbox_size[np.ix_(is_rotated, idxs)] = bbox_size[np.ix_(
is_rotated, list(reversed(idxs)))]
bboxes = np.hstack((bbox_center - bbox_size,
bbox_center + bbox_size, bboxes[:, 6:]))
elif self.config.normalize_rotation2:
bboxes[:, 6] = (bboxes[:, 6] % np.pi - np.pi / 2) * 2
if not self.HAS_GT_BBOX:
instance_mask = self.get_instance_mask(semantic_labels,
instance_labels)
bboxes, instance_labels = get_bbox(coords, semantic_labels,
instance_labels, instance_mask,
self.ignore_mask)
labels = np.hstack(
(np.expand_dims(semantic_labels,
1), np.expand_dims(instance_labels, 1)))
return_args = [coords, feats, labels, bboxes]
if self.return_transformation:
transformation = np.expand_dims(transformation, 0)
return_args.extend([
pointcloud.astype(np.float32),
transformation.astype(np.float32)
])
return tuple(return_args)