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)
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)