-
Notifications
You must be signed in to change notification settings - Fork 0
/
voxelizer.py
141 lines (123 loc) · 6.18 KB
/
voxelizer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
# Codes are taken from BPNet, CVPR'21
# https://github.com/wbhu/BPNet/blob/main/dataset/voxelizer.py
import collections
import numpy as np
from dataset.voxelization_utils import sparse_quantize
from scipy.linalg import expm, norm
# Rotation matrix along axis with angle theta
def M(axis, theta):
return expm(np.cross(np.eye(3), axis / norm(axis) * theta))
class Voxelizer:
def __init__(self,
voxel_size=1,
clip_bound=None,
use_augmentation=False,
scale_augmentation_bound=None,
rotation_augmentation_bound=None,
translation_augmentation_ratio_bound=None,
ignore_label=255):
'''
Args:
voxel_size: side length of a voxel
clip_bound: boundary of the voxelizer. Points outside the bound will be deleted
expects either None or an array like ((-100, 100), (-100, 100), (-100, 100)).
scale_augmentation_bound: None or (0.9, 1.1)
rotation_augmentation_bound: None or ((np.pi / 6, np.pi / 6), None, None) for 3 axis.
Use random order of x, y, z to prevent bias.
translation_augmentation_bound: ((-5, 5), (0, 0), (-10, 10))
ignore_label: label assigned for ignore (not a training label).
'''
self.voxel_size = voxel_size
self.clip_bound = clip_bound
self.ignore_label = ignore_label
# Augmentation
self.use_augmentation = use_augmentation
self.scale_augmentation_bound = scale_augmentation_bound
self.rotation_augmentation_bound = rotation_augmentation_bound
self.translation_augmentation_ratio_bound = translation_augmentation_ratio_bound
def get_transformation_matrix(self):
voxelization_matrix, rotation_matrix = np.eye(4), np.eye(4)
# Get clip boundary from config or pointcloud.
# Get inner clip bound to crop from.
# Transform pointcloud coordinate to voxel coordinate.
# 1. Random rotation
rot_mat = np.eye(3)
if self.use_augmentation and self.rotation_augmentation_bound is not None:
if isinstance(self.rotation_augmentation_bound, collections.Iterable):
rot_mats = []
for axis_ind, rot_bound in enumerate(self.rotation_augmentation_bound):
theta = 0
axis = np.zeros(3)
axis[axis_ind] = 1
if rot_bound is not None:
theta = np.random.uniform(*rot_bound)
rot_mats.append(M(axis, theta))
# Use random order
np.random.shuffle(rot_mats)
rot_mat = rot_mats[0] @ rot_mats[1] @ rot_mats[2]
else:
raise ValueError()
rotation_matrix[:3, :3] = rot_mat
# 2. Scale and translate to the voxel space.
scale = 1 / self.voxel_size
if self.use_augmentation and self.scale_augmentation_bound is not None:
scale *= np.random.uniform(*self.scale_augmentation_bound)
np.fill_diagonal(voxelization_matrix[:3, :3], scale)
# Get final transformation matrix.
return voxelization_matrix, rotation_matrix
def clip(self, coords, center=None, trans_aug_ratio=None):
bound_min = np.min(coords, 0).astype(float)
bound_max = np.max(coords, 0).astype(float)
bound_size = bound_max - bound_min
if center is None:
center = bound_min + bound_size * 0.5
lim = self.clip_bound
if trans_aug_ratio is not None:
trans = np.multiply(trans_aug_ratio, bound_size)
center += trans
# Clip points outside the limit
clip_inds = ((coords[:, 0] >= (lim[0][0] + center[0])) &
(coords[:, 0] < (lim[0][1] + center[0])) &
(coords[:, 1] >= (lim[1][0] + center[1])) &
(coords[:, 1] < (lim[1][1] + center[1])) &
(coords[:, 2] >= (lim[2][0] + center[2])) &
(coords[:, 2] < (lim[2][1] + center[2])))
return clip_inds
def voxelize(self, coords, feats, labels = None, center=None, link=None, return_ind=False):
assert coords.shape[1] == 3 and coords.shape[0] == feats.shape[0] and coords.shape[0]
if self.clip_bound is not None:
trans_aug_ratio = np.zeros(3)
if self.use_augmentation and self.translation_augmentation_ratio_bound is not None:
for axis_ind, trans_ratio_bound in enumerate(self.translation_augmentation_ratio_bound):
trans_aug_ratio[axis_ind] = np.random.uniform(*trans_ratio_bound)
clip_inds = self.clip(coords, center, trans_aug_ratio)
if clip_inds.sum():
coords, feats = coords[clip_inds], feats[clip_inds]
if labels is not None:
labels = labels[clip_inds]
# Get rotation and scale
M_v, M_r = self.get_transformation_matrix()
# Apply transformations
rigid_transformation = M_v
if self.use_augmentation:
rigid_transformation = M_r @ rigid_transformation
homo_coords = np.hstack((coords, np.ones((coords.shape[0], 1), dtype=coords.dtype)))
coords_aug = np.floor(homo_coords @ rigid_transformation.T[:, :3])
# Align all coordinates to the origin.
min_coords = coords_aug.min(0)
M_t = np.eye(4)
M_t[:3, -1] = -min_coords
rigid_transformation = M_t @ rigid_transformation
coords_aug = np.floor(coords_aug - min_coords)
inds, inds_reconstruct = sparse_quantize(coords_aug, return_index=True)
# coords_aug, feats, labels = coords_aug[inds], feats[inds], labels[inds]
coords_aug, feats = coords_aug[inds], feats[inds]
# Normal rotation
if feats.shape[1] > 6:
feats[:, 3:6] = feats[:, 3:6] @ (M_r[:3, :3].T)
if return_ind:
return coords_aug, feats, labels, np.array(inds_reconstruct), inds
if link is not None:
return coords_aug, feats, labels, np.array(inds_reconstruct), link[inds]
#return coords_aug, feats, labels, np.array(inds_reconstruct)
return coords_aug, feats, np.array(inds_reconstruct)