File size: 4,152 Bytes
b30c1d8 |
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 |
import numpy as np
import torch
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()
|