Group points by boxes v2 (keras-team#1090)

* add v2 of group_points_by_boxes.

* fix format.

* add test.

keras_cv/layers/preprocessing3d/group_points_by_bounding_boxes.py

@@ -15,6 +15,7 @@
 import tensorflow as tf
 
 from keras_cv.layers.preprocessing3d import base_augmentation_layer_3d
+from keras_cv.ops.point_cloud import group_points_by_boxes
 from keras_cv.ops.point_cloud import is_within_box3d
 
 POINT_CLOUDS = base_augmentation_layer_3d.POINT_CLOUDS
@@ -98,6 +99,7 @@ def augment_point_clouds_bounding_boxes(
             point_clouds[:, :, :3], object_bounding_boxes[:, :, :7]
         )
         # Filter bounding boxes using the current frame.
+        # [num_boxes]
         min_points_filter = (
             tf.reduce_sum(tf.cast(points_in_bounding_boxes[0], dtype=tf.int32), axis=0)
             >= self._min_points_per_bounding_boxes
@@ -138,6 +140,49 @@ def augment_point_clouds_bounding_boxes(
             object_bounding_boxes,
         )
 
+    def augment_point_clouds_bounding_boxes_v2(
+        self, point_clouds, bounding_boxes, **kwargs
+    ):
+        if self._label_index:
+            bounding_boxes_mask = (
+                bounding_boxes[0, :, BOX_LABEL_INDEX] == self._label_index
+            )
+            object_bounding_boxes = tf.boolean_mask(
+                bounding_boxes, bounding_boxes_mask, axis=1
+            )
+        else:
+            bounding_boxes_mask = bounding_boxes[0, :, BOX_LABEL_INDEX] > 0.0
+            object_bounding_boxes = tf.boolean_mask(
+                bounding_boxes, bounding_boxes_mask, axis=1
+            )
+
+        # [frames, num_boxes, ragged_points]
+        points_in_bounding_boxes = group_points_by_boxes(
+            point_clouds[:, :, :3], object_bounding_boxes[:, :, :7]
+        )
+        # Filter bounding boxes using the current frame.
+        # [num_boxes]
+        min_points_filter = (
+            points_in_bounding_boxes.row_lengths(-1)
+            >= self._min_points_per_bounding_boxes
+        )
+
+        # [frames, num_valid_boxes, box_feature]
+        object_bounding_boxes = tf.ragged.boolean_mask(
+            object_bounding_boxes, min_points_filter
+        )
+        # [frames, num_valid_boxes, ragged_points]
+        points_in_bounding_boxes = tf.ragged.boolean_mask(
+            points_in_bounding_boxes, min_points_filter
+        )
+        # point_clouds: [frames, num_points, point_feature]
+        # object_point_clouds: [frames, num_valid_boxes, ragged_points, point_feature]
+        object_point_clouds = tf.gather(
+            point_clouds, points_in_bounding_boxes, axis=1, batch_dims=1
+        )
+
+        return (object_point_clouds, object_bounding_boxes)
+
     def _augment(self, inputs):
         result = inputs
         point_clouds = inputs[POINT_CLOUDS]

keras_cv/layers/preprocessing3d/group_points_by_bounding_boxes_test.py

@@ -152,3 +152,43 @@ def test_augment_batch_point_clouds_and_bounding_boxes(self):
         object_point_clouds = tf.sort(object_point_clouds, axis=-2)
         self.assertAllClose(outputs[OBJECT_POINT_CLOUDS], object_point_clouds)
         self.assertAllClose(outputs[OBJECT_BOUNDING_BOXES], object_bounding_boxes)
+
+    def test_augment_point_clouds_and_bounding_boxes_v2(self):
+        add_layer = GroupPointsByBoundingBoxes(
+            label_index=1,
+            min_points_per_bounding_boxes=1,
+            max_points_per_bounding_boxes=2,
+        )
+        point_clouds = np.array(
+            [[[0, 1, 2, 3, 4], [10, 1, 2, 3, 4], [0, -1, 2, 3, 4], [100, 100, 2, 3, 4]]]
+            * 2
+        ).astype("float32")
+        bounding_boxes = np.array(
+            [
+                [
+                    [0, 0, 0, 4, 4, 4, 0, 1],
+                    [10, 1, 2, 2, 2, 2, 0, 1],
+                    [20, 20, 20, 1, 1, 1, 0, 1],
+                ]
+            ]
+            * 2
+        ).astype("float32")
+        point_clouds = tf.convert_to_tensor(point_clouds)
+        bounding_boxes = tf.convert_to_tensor(bounding_boxes)
+        outputs = add_layer.augment_point_clouds_bounding_boxes_v2(
+            point_clouds=point_clouds, bounding_boxes=bounding_boxes
+        )
+        object_point_clouds, object_bounding_boxes = outputs[0], outputs[1]
+        expected_object_point_clouds = np.array(
+            [[[[0, 1, 2, 3, 4], [0, -1, 2, 3, 4]], [[10, 1, 2, 3, 4], [0, 0, 0, 0, 0]]]]
+            * 2
+        ).astype("float32")
+        expected_object_bounding_boxes = np.array(
+            [[[0, 0, 0, 4, 4, 4, 0, 1], [10, 1, 2, 2, 2, 2, 0, 1]]] * 2
+        ).astype("float32")
+        self.assertAllClose(
+            expected_object_point_clouds, object_point_clouds.to_tensor()
+        )
+        self.assertAllClose(
+            expected_object_bounding_boxes, object_bounding_boxes.to_tensor()
+        )

keras_cv/ops/point_cloud.py

@@ -71,19 +71,34 @@ def group_points_by_boxes(points, boxes):
       the point indices that belong to the box.
 
     """
-    num_boxes = boxes.get_shape().as_list()[-2]
-    if not num_boxes:
-        raise ValueError(
-            f"boxes cannot have unknown shape, got {boxes.get_shape().as_list()}"
-        )
+    num_boxes = boxes.get_shape().as_list()[-2] or tf.shape(boxes)[-2]
+    # [..., num_points]
     box_indices = within_box3d_index(points, boxes)
     num_points = points.get_shape().as_list()[-2] or tf.shape(points)[-2]
     point_indices = tf.range(num_points, dtype=tf.int32)
-    point_mask = tf.math.greater_equal(box_indices, 0)
-    point_indices = tf.boolean_mask(point_indices, point_mask)
-    box_indices = tf.boolean_mask(box_indices, point_mask)
-    res = tf.dynamic_partition(point_indices, box_indices, num_partitions=num_boxes)
-    return tf.ragged.stack(res, axis=0)
+
+    def group_per_sample(box_index):
+        point_mask = tf.math.greater_equal(box_index, 0)
+        valid_point_indices = tf.boolean_mask(point_indices, point_mask)
+        valid_box_index = tf.boolean_mask(box_index, point_mask)
+        res = tf.ragged.stack_dynamic_partitions(
+            valid_point_indices, valid_box_index, num_partitions=num_boxes
+        )
+        return res
+
+    boxes_rank = boxes.shape.rank
+    if boxes_rank == 2:
+        return group_per_sample(box_indices)
+    elif boxes_rank == 3:
+        num_samples = boxes.get_shape().as_list()[0]
+        res_list = []
+        for i in range(num_samples):
+            res_list.append(group_per_sample(box_indices[i]))
+        return tf.ragged.stack(res_list)
+    else:
+        raise ValueError(
+            f"Does not support box rank > 3, got boxes shape {boxes.shape}"
+        )
 
 
 # TODO(lengzhaoqi/tanzhenyu): compare the performance with v1