1. Add global_z_rotation data augmentation.

2. Add wrap_angle_rad helper function.

keras_cv/layers/preprocessing3d/global_z_rotation.py

@@ -0,0 +1,76 @@
+
+# Copyright 2022 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tensorflow as tf
+
+from keras_cv.layers.preprocessing3d.base_augmentation_layer_3d import (
+    BaseAugmentationLayer3D,
+)
+from keras_cv.ops.point_cloud import coordinate_transform, wrap_angle_rad
+
+
+POINT_CLOUDS = "point_clouds"
+BOUNDING_BOXES = "bounding_boxes"
+
+class GlobalZRotation(BaseAugmentationLayer3D):
+    """A preprocessing layer which randomly rotates point clouds and bounding boxes along
+    Z axis during training.
+
+    This layer will randomly rotate the whole scene along the Z axis based on a randomly sampled 
+    rotation angle between [-max_rotation_angle, max_rotation_angle] following a uniform distribution.
+    During inference time, the output will be identical to input. Call the layer with `training=True` to rotate the input.
+
+    Input shape:
+      point_clouds: 2D (single frame) or 3D (multi frames) float32 Tensor with shape 
+        [..., num of points, num of point features]. 
+        The first 5 features are [x, y, z, class, range].
+      bounding_boxes: 2D (single frame) or 3D (multi frames) float32 Tensor with shape
+        [..., num of boxes, num of box features]. 
+        The first 7 features are [x, y, z, dx, dy, dz, phi].
+
+    Output shape:
+      A tuple of two Tensors (point_clouds, bounding_boxes) with the same shape as input Tensors.
+
+    Arguments:
+      max_rotation_angle: A float scaler or Tensor sets the maximum rotation angle. 
+    """
+    def __init__(self, max_rotation_angle, **kwargs):
+        super().__init__(**kwargs)
+        if max_rotation_angle<0:
+            raise ValueError(
+                "max_rotation_angle must be >=0."
+            )
+        self._max_rotation_angle = max_rotation_angle
+
+    def get_random_transformation(self, **kwargs):
+        random_rotation_z = self._random_generator.random_uniform(
+            (), minval=-self._max_rotation_angle, maxval=self._max_rotation_angle
+        )
+        return {"pose": tf.stack([0, 0, 0, random_rotation_z, 0, 0], axis=0)}
+
+    def augment_point_clouds_bounding_boxes(
+        self, point_clouds, bounding_boxes, transformation, **kwargs
+    ):
+        pose = transformation["pose"]
+        point_clouds_xyz = coordinate_transform(point_clouds[..., :3], pose)
+        point_clouds = tf.concat([point_clouds_xyz, point_clouds[..., 3:]], axis=-1)
+
+        bounding_boxes_xyz = coordinate_transform(bounding_boxes[..., :3], pose)
+        bounding_boxes_heading = wrap_angle_rad(bounding_boxes[..., 6:7] - pose[3])
+        bounding_boxes = tf.concat([
+            bounding_boxes_xyz, bounding_boxes[..., 3:6], 
+            bounding_boxes_heading, bounding_boxes[..., 7:]], axis=-1)
+
+        return (point_clouds, bounding_boxes)

keras_cv/layers/preprocessing3d/global_z_rotation_test.py

@@ -0,0 +1,56 @@
+# Copyright 2022 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import numpy as np
+import tensorflow as tf
+
+from keras_cv.layers.preprocessing3d.global_z_rotation import GlobalZRotation
+
+
+POINT_CLOUDS = "point_clouds"
+BOUNDING_BOXES = "bounding_boxes"
+
+
+class GlobalZRotationTest(tf.test.TestCase):
+
+    def test_augment_point_clouds_and_bounding_boxes(self):
+        add_layer = GlobalZRotation(max_rotation_angle=1.0)
+        point_clouds = np.random.random(size=(2, 50, 10)).astype("float32")
+        bounding_boxes = np.random.random(size=(2, 10, 7)).astype("float32")
+        inputs = {POINT_CLOUDS: point_clouds, BOUNDING_BOXES: bounding_boxes}
+        outputs = add_layer(inputs)
+        self.assertNotAllClose(inputs, outputs)
+
+    def test_not_augment_point_clouds_and_bounding_boxes(self):
+        add_layer = GlobalZRotation(max_rotation_angle=0.0)
+        point_clouds = np.random.random(size=(2, 50, 10)).astype("float32")
+        bounding_boxes = np.random.random(size=(2, 10, 7)).astype("float32")
+        inputs = {POINT_CLOUDS: point_clouds, BOUNDING_BOXES: bounding_boxes}
+        outputs = add_layer(inputs)
+        self.assertAllClose(inputs, outputs)
+
+    def test_augment_batch_point_clouds_and_bounding_boxes(self):
+        add_layer = GlobalZRotation(max_rotation_angle=1.0)
+        point_clouds = np.random.random(size=(3, 2, 50, 10)).astype("float32")
+        bounding_boxes = np.random.random(size=(3, 2, 10, 7)).astype("float32")
+        inputs = {POINT_CLOUDS: point_clouds, BOUNDING_BOXES: bounding_boxes}
+        outputs = add_layer(inputs)
+        self.assertNotAllClose(inputs, outputs)
+
+    def test_not_augment_batch_point_clouds_and_bounding_boxes(self):
+        add_layer = GlobalZRotation(max_rotation_angle=0.0)
+        point_clouds = np.random.random(size=(3, 2, 50, 10)).astype("float32")
+        bounding_boxes = np.random.random(size=(3, 2, 10, 7)).astype("float32")
+        inputs = {POINT_CLOUDS: point_clouds, BOUNDING_BOXES: bounding_boxes}
+        outputs = add_layer(inputs)
+        self.assertAllClose(inputs, outputs)

keras_cv/ops/point_cloud.py

@@ -13,11 +13,15 @@
 # limitations under the License.
 
 import tensorflow as tf
-
+import numpy as np
 
 def get_rank(tensor):
     return tensor.shape.ndims or tf.rank(tensor)
 
+def wrap_angle_rad(angles_rad, min_val=-np.pi, max_val=np.pi):
+  """Wrap the value of `angles_rad` to the range [min_val, max_val]."""
+  max_min_diff = max_val - min_val
+  return min_val + tf.math.floormod(angles_rad + max_val, max_min_diff)
 
 def _get_3d_rotation_matrix(yaw, roll, pitch):
     """Creates 3x3 rotation matrix from yaw, roll, pitch (angles in radians).

keras_cv/ops/point_cloud_test.py

@@ -18,6 +18,11 @@
 
 from keras_cv import ops
 
+class AngleTest(tf.test.TestCase):
+    def test_wrap_angle_rad(self):
+        self.assertAllClose(-np.pi+0.1, ops.point_cloud.wrap_angle_rad(np.pi+0.1))
+        self.assertAllClose(0.0, ops.point_cloud.wrap_angle_rad(2 * np.pi))
+
 
 class Boxes3DTestCase(tf.test.TestCase, parameterized.TestCase):
     def test_convert_center_to_corners(self):