def augment(
coords,
num_points: int = 1024,
shuffle: bool = True,
rotate_scheme: str = rigid.RotateScheme.NONE,
angle_stddev: Optional[float] = None,
angle_clip: Optional[float] = None,
jitter_stddev: float = None,
jitter_clip: float = None,
maybe_reflect_x: bool = False,
uniform_scale_range: Optional[Tuple[float, float]] = None,
drop_prob_limits: Optional[Tuple[float, float]] = None,
up_dim: int = 2,
shuffle_first: bool = False,
):
coords.shape.assert_has_rank(2)
assert coords.shape[1] == 3
if shuffle_first and shuffle:
coords = tfrng.shuffle(coords)
if num_points is not None:
coords = coords[:num_points]
if not shuffle_first and shuffle:
coords = tfrng.shuffle(coords)
if drop_prob_limits is not None:
lower, upper = drop_prob_limits
drop_prob = tfrng.uniform((), lower, upper)
if shuffle:
size = tf.cast(
tf.cast(tf.shape(coords)[0], tf.float32) * (1 - drop_prob),
tf.int64)
coords = coords[:size]
else:
mask = tfrng.uniform((tf.shape(coords)[0], )) > drop_prob
coords = tf.boolean_mask(coords, mask)
if up_dim != 2:
roll = (2 - up_dim) % 3
coords = tf.roll(coords, roll, axis=-1)
if jitter_stddev is not None:
coords = jitter.jitter_positions(coords, jitter_stddev, jitter_clip)
if angle_stddev is not None:
coords, _ = rigid.random_rotation(coords,
angle_stddev=angle_stddev,
angle_clip=angle_clip)
if rotate_scheme is not None:
coords, _ = rigid.rotate_by_scheme(coords, None, scheme=rotate_scheme)
if uniform_scale_range is not None:
coords = rigid.random_scale(coords, uniform_range=uniform_scale_range)
if maybe_reflect_x:
prob = tfrng.uniform(()) < 0.5
coords = tf.cond(prob, lambda: coords, lambda: rigid.reflect(coords))
assert coords.shape[0] is None
return coords
def to_bool(maybe: MaybeBool) -> Union[bool, BoolTensor, FloatTensor]:
if isinstance(maybe, (tf.Tensor, np.ndarray)):
# dtype = maybe.dtype
dtype = getattr(maybe, "dtype")
if dtype.is_bool:
return maybe
elif dtype.is_floating:
return tfrng.uniform(()) < maybe
else:
raise ValueError(
f"maybe must be a bool or float if a tensor, got dtype {maybe.dtype}"
)
elif isinstance(maybe, bool):
return maybe
else:
return tfrng.uniform(()) < maybe
def _maybe_reflect(positions, axis=-1, dim=0, prob=0.5):
should_reflect = tfrng.uniform(shape=(), dtype=tf.float32) > prob
return tf.cond(
should_reflect,
lambda: tuple(reflect(p, dim=dim, axis=axis) for p in positions),
lambda: tuple(positions),
)
def _rotate(positions, normals=None, angle=None, rotation_dim=2, impl=tf):
"""
Randomly rotate the point cloud about the z-axis.
Args:
positions: (n, 3) float array
normals (optional): (n, 3) float array
angle: float scalar. If None, a uniform random angle in [0, 2pi) is
used.
rotation_dim: int denoting x (0), y (1), or z (2) axis about which to
rotate
impl: tf or np
Returns:
rotated (`positions`, `normals`). `normals` will be None if not
provided. shape and dtype is the same as provided.
"""
dtype = positions.dtype
if angle is None:
angle = tfrng.uniform((), dtype=dtype) * (2 * np.pi)
if normals is not None:
assert normals.dtype == dtype
c = impl.cos(angle)
s = impl.sin(angle)
# multiply on right, use non-standard rotation matrix (-s and s swapped)
rotation_matrix = impl.reshape(
impl.stack(_pack_rotation_matrix(c, s, rotation_dim=rotation_dim)),
(3, 3))
positions = impl.matmul(positions, rotation_matrix)
if normals is not None:
normals = impl.matmul(normals, rotation_matrix)
return positions, normals
def augment(
time: IntTensor,
coords: IntTensor,
polarity: BoolTensor,
grid_shape: GridShape,
flip_lr: MaybeBool = False,
flip_ud: MaybeBool = False,
flip_time: MaybeBool = False,
rotate_limits: Optional[Tuple[float, float]] = None,
) -> Tuple[IntTensor, IntTensor, BoolTensor, Optional[BoolTensor]]:
mask = None
if all(b is False for b in (flip_lr, flip_ud, flip_time)) and rotate_limits is None:
return time, coords, polarity, mask
if flip_lr is not False or flip_ud is not False:
# autograph complains about lambdas
def flipx():
return grid_shape[0] - 1 - x
def no_flipx():
return x
def flipy():
return grid_shape[1] - 1 - y
def no_flipy():
return y
x, y = tf.unstack(coords, axis=-1)
x = smart_cond(to_bool(flip_lr), flipx, no_flipx)
y = smart_cond(to_bool(flip_ud), flipy, no_flipy)
coords = tf.stack((x, y), axis=-1)
def flipt():
return reverse_time(time, coords, polarity)
def no_flipt():
return (time, coords, polarity)
time, coords, polarity = smart_cond(to_bool(flip_time), flipt, no_flipt)
if rotate_limits is not None:
min_rot, max_rot = rotate_limits
radians = tfrng.uniform((), minval=min_rot, maxval=max_rot)
dtype = coords.dtype
coords = tf.cast(coords, tf.float32)
coords = rotate2d(coords, radians, center=tf.cast(grid_shape, tf.float32) / 2)
coords = tf.cast(tf.round(coords), dtype)
mask = tf.reduce_all(tf.logical_and(coords >= 0, coords < grid_shape), axis=-1)
time = tf.boolean_mask(time, mask)
coords = tf.boolean_mask(coords, mask)
polarity = tf.boolean_mask(polarity, mask)
return time, coords, polarity, mask
def random_scale(positions: tf.Tensor,
stddev=None,
uniform_range=None) -> tf.Tensor:
if stddev is not None:
scale = tfrng.truncated_normal(shape=(), mean=1.0, stddev=stddev)
elif uniform_range is not None:
minval, maxval = uniform_range
scale = tfrng.uniform(shape=(), minval=minval, maxval=maxval)
else:
raise NotImplementedError(
"One of stddev or uniform_range must be defined")
return positions * scale
def rotate_by_scheme(positions,
normals=None,
scheme: str = RotateScheme.RANDOM
) -> Tuple[tf.Tensor, Optional[tf.Tensor]]:
"""scheme should be in ("random", "pca-xy", "none")."""
RotateScheme.validate(scheme)
if scheme == RotateScheme.NONE:
return positions, normals
if scheme == RotateScheme.PCA_XY:
angle = get_pca_xy_angle(positions)
elif scheme == RotateScheme.RANDOM:
angle = tfrng.uniform(shape=(), dtype=positions.dtype) * (2 * np.pi)
else:
raise ValueError('Unrecognized scheme "%s"' % scheme)
return rotate(positions, normals, angle)
def augment_segmentation(
coords,
normals,
labels,
num_points=None,
shuffle=True,
rotate_scheme="none",
angle_stddev: Optional[float] = None,
angle_clip: Optional[float] = None,
jitter_stddev: float = None,
jitter_clip: float = None,
maybe_reflect_x=False,
uniform_scale_range: Optional[Tuple[float, float]] = None,
drop_prob_limits: Optional[Tuple[float, float]] = None,
up_dim: int = 2,
):
coords.shape.assert_has_rank(2)
assert coords.shape[1] == 3
normals.shape.assert_has_rank(2)
assert normals.shape[1] == 3
if num_points is not None:
coords = coords[:num_points]
normals = normals[:num_points]
labels = labels[:num_points]
size = tf.size(labels)
if shuffle:
order = tfrng.shuffle(tf.range(size))
coords, normals, labels = (tf.gather(x, order)
for x in (coords, normals, labels))
if drop_prob_limits is not None:
lower, upper = drop_prob_limits
drop_prob = tfrng.uniform((), lower, upper)
if shuffle:
size = tf.cast(
tf.cast(tf.shape(coords)[0], tf.float32) * (1 - drop_prob),
tf.int64)
def f(x):
return x[:size]
else:
mask = tfrng.uniform((tf.shape(coords)[0], )) > drop_prob
def f(x):
return tf.boolean_mask(x, mask)
coords, normals, labels = (f(x) for x in (coords, normals, labels))
if up_dim != 2:
roll = (2 - up_dim) % 3
coords, normals = (tf.roll(x, roll, axis=-1)
for x in (coords, normals))
if jitter_stddev is not None:
coords = jitter.jitter_positions(coords, jitter_stddev, jitter_clip)
if angle_stddev is not None:
coords, normals = rigid.random_rotation(coords,
normals,
angle_stddev=angle_stddev,
angle_clip=angle_clip)
if rotate_scheme is not None:
coords, normals = rigid.rotate_by_scheme(coords,
normals,
scheme=rotate_scheme)
if uniform_scale_range is not None:
coords = rigid.random_scale(coords, uniform_range=uniform_scale_range)
if maybe_reflect_x:
prob = tfrng.uniform(()) < 0.5
coords, normals = tf.cond(
prob,
lambda: (coords, normals),
lambda: (rigid.reflect(coords), rigid.reflect(normals)),
)
return coords, normals, labels
def map_func(x):
return x + tfrng.uniform(())
def tfrng_map_func(x):
scale = tfrng.normal((), stddev=0.1, mean=1.0)
shift = tfrng.uniform(())
return transform(x, scale, shift)
