-
Notifications
You must be signed in to change notification settings - Fork 0
/
data_utils.py
148 lines (113 loc) · 4.18 KB
/
data_utils.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
142
143
144
145
146
147
148
from __future__ import (
division,
absolute_import,
with_statement,
print_function,
unicode_literals,
)
import torch
import numpy as np
def angle_axis(angle, axis):
# type: (float, np.ndarray) -> float
r"""Returns a 4x4 rotation matrix that performs a rotation around axis by angle
Parameters
----------
angle : float
Angle to rotate by
axis: np.ndarray
Axis to rotate about
Returns
-------
torch.Tensor
3x3 rotation matrix
"""
u = axis / np.linalg.norm(axis)
cosval, sinval = np.cos(angle), np.sin(angle)
# yapf: disable
cross_prod_mat = np.array([[0.0, -u[2], u[1]],
[u[2], 0.0, -u[0]],
[-u[1], u[0], 0.0]])
R = torch.from_numpy(
cosval * np.eye(3)
+ sinval * cross_prod_mat
+ (1.0 - cosval) * np.outer(u, u)
)
# yapf: enable
return R.float()
class PointcloudScale(object):
def __init__(self, lo=0.8, hi=1.25):
self.lo, self.hi = lo, hi
def __call__(self, points):
scaler = np.random.uniform(self.lo, self.hi)
points[:, 0:3] *= scaler
return points
class PointcloudRotate(object):
def __init__(self, axis=np.array([0.0, 1.0, 0.0])):
self.axis = axis
def __call__(self, points):
rotation_angle = np.random.uniform() * 2 * np.pi
rotation_matrix = angle_axis(rotation_angle, self.axis)
normals = points.size(1) > 3
if not normals:
return torch.matmul(points, rotation_matrix.t())
else:
pc_xyz = points[:, 0:3]
pc_normals = points[:, 3:]
points[:, 0:3] = torch.matmul(pc_xyz, rotation_matrix.t())
points[:, 3:] = torch.matmul(pc_normals, rotation_matrix.t())
return points
class PointcloudRotatePerturbation(object):
def __init__(self, angle_sigma=0.06, angle_clip=0.18):
self.angle_sigma, self.angle_clip = angle_sigma, angle_clip
def _get_angles(self):
angles = np.clip(
self.angle_sigma * np.random.randn(3), -self.angle_clip, self.angle_clip
)
return angles
def __call__(self, points):
angles = self._get_angles()
Rx = angle_axis(angles[0], np.array([1.0, 0.0, 0.0]))
Ry = angle_axis(angles[1], np.array([0.0, 1.0, 0.0]))
Rz = angle_axis(angles[2], np.array([0.0, 0.0, 1.0]))
rotation_matrix = torch.matmul(torch.matmul(Rz, Ry), Rx)
normals = points.size(1) > 3
if not normals:
return torch.matmul(points, rotation_matrix.t())
else:
pc_xyz = points[:, 0:3]
pc_normals = points[:, 3:]
points[:, 0:3] = torch.matmul(pc_xyz, rotation_matrix.t())
points[:, 3:] = torch.matmul(pc_normals, rotation_matrix.t())
return points
class PointcloudJitter(object):
def __init__(self, std=0.01, clip=0.05):
self.std, self.clip = std, clip
def __call__(self, points):
jittered_data = (
points.new(points.size(0), 3)
.normal_(mean=0.0, std=self.std)
.clamp_(-self.clip, self.clip)
)
points[:, 0:3] += jittered_data
return points
class PointcloudTranslate(object):
def __init__(self, translate_range=0.1):
self.translate_range = translate_range
def __call__(self, points):
translation = np.random.uniform(-self.translate_range, self.translate_range)
points[:, 0:3] += translation
return points
class PointcloudToTensor(object):
def __call__(self, points):
return torch.from_numpy(points).float()
class PointcloudRandomInputDropout(object):
def __init__(self, max_dropout_ratio=0.875):
assert max_dropout_ratio >= 0 and max_dropout_ratio < 1
self.max_dropout_ratio = max_dropout_ratio
def __call__(self, points):
pc = points.numpy()
dropout_ratio = np.random.random() * self.max_dropout_ratio # 0~0.875
drop_idx = np.where(np.random.random((pc.shape[0])) <= dropout_ratio)[0]
if len(drop_idx) > 0:
pc[drop_idx] = pc[0] # set to the first point
return torch.from_numpy(pc).float()