[Port] OD layers to Keras 3 (#2295)

* chore: porting roi aling to keras 3

* chore: fixing the scope, using ones in place of constant

* chore: porting roi generation to keras 3 with test

note: the nms bit reproduces -1 instead of 0

* chore: port roi pooling

* chore: fix pool and port sampler

* chore: port label encoder

* chore: swap get_shape with ops.shape

* lint error

* chore: porting sampling to keras 3

* lint fix

* chore: using random from backend

* chore: disabling flaky test

* chore: disable roi sampler test

* chore: ignore lint

* chore: skipping test the right way

* chore: using ops shape

* chore: tests pass for all backends

removed vectorized map as it was not working for jax and torch
used ops convert_to_numpy in tests to make np operations work on torch tensor

* chore: explicit type cast to int32

Author: ariG23498

keras_cv/layers/object_detection/roi_align.py

@@ -12,17 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Mapping
 from typing import Optional
-from typing import Tuple
-from typing import Union
 
-import tensorflow as tf
-from tensorflow import keras
-
-from keras_cv import bounding_box
 from keras_cv.api_export import keras_cv_export
-from keras_cv.backend import assert_tf_keras
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+from keras_cv.layers import NonMaxSuppression
 
 
 @keras_cv_export("keras_cv.layers.ROIGenerator")
@@ -97,7 +92,6 @@ def __init__(
         post_nms_topk_test: int = 1000,
         **kwargs,
     ):
-        assert_tf_keras("keras_cv.layers.ROIGenerator")
         super().__init__(**kwargs)
         self.bounding_box_format = bounding_box_format
         self.pre_nms_topk_train = pre_nms_topk_train
@@ -112,10 +106,10 @@ def __init__(
 
     def call(
         self,
-        multi_level_boxes: Union[tf.Tensor, Mapping[int, tf.Tensor]],
-        multi_level_scores: Union[tf.Tensor, Mapping[int, tf.Tensor]],
+        multi_level_boxes,
+        multi_level_scores,
         training: Optional[bool] = None,
-    ) -> Tuple[tf.Tensor, tf.Tensor]:
+    ):
         """
         Args:
           multi_level_boxes: float Tensor. A dictionary or single Tensor of
@@ -131,7 +125,6 @@ def call(
           rois: float Tensor of [batch_size, post_nms_topk, 4]
           roi_scores: float Tensor of [batch_size, post_nms_topk]
         """
-
         if training:
             pre_nms_topk = self.pre_nms_topk_train
             post_nms_topk = self.post_nms_topk_train
@@ -144,53 +137,35 @@ def call(
             nms_iou_threshold = self.nms_iou_threshold_test
 
         def per_level_gen(boxes, scores):
-            scores_shape = scores.get_shape().as_list()
-            # scores can also be [batch_size, num_boxes, 1]
+            boxes = ops.convert_to_tensor(boxes, dtype="float32")
+            scores = ops.convert_to_tensor(scores, dtype="float32")
+            scores_shape = ops.shape(scores)
+            # Check if scores is a 3-dimensional tensor
+            # ([batch_size, num_boxes, 1])
+            # If so, remove the last dimension to make it 2D
             if len(scores_shape) == 3:
-                scores = tf.squeeze(scores, axis=-1)
-            _, num_boxes = scores.get_shape().as_list()
+                scores = ops.squeeze(scores, axis=-1)
+            _, num_boxes = scores_shape
             level_pre_nms_topk = min(num_boxes, pre_nms_topk)
             level_post_nms_topk = min(num_boxes, post_nms_topk)
-            scores, sorted_indices = tf.nn.top_k(
+            scores, sorted_indices = ops.top_k(
                 scores, k=level_pre_nms_topk, sorted=True
             )
-            boxes = tf.gather(boxes, sorted_indices, batch_dims=1)
-            # convert from input format to yxyx for the TF NMS operation
-            boxes = bounding_box.convert_format(
-                boxes,
-                source=self.bounding_box_format,
-                target="yxyx",
+            boxes = ops.take_along_axis(
+                boxes, sorted_indices[..., None], axis=1
             )
             # TODO(tanzhenyu): consider supporting soft / batched nms for accl
-            selected_indices, num_valid = tf.image.non_max_suppression_padded(
-                boxes,
-                scores,
-                max_output_size=level_post_nms_topk,
+            boxes = NonMaxSuppression(
+                bounding_box_format=self.bounding_box_format,
+                from_logits=False,
                 iou_threshold=nms_iou_threshold,
-                score_threshold=nms_score_threshold,
-                pad_to_max_output_size=True,
-                sorted_input=True,
-                canonicalized_coordinates=True,
-            )
-            # convert back to input format
-            boxes = bounding_box.convert_format(
-                boxes,
-                source="yxyx",
-                target=self.bounding_box_format,
-            )
-            level_rois = tf.gather(boxes, selected_indices, batch_dims=1)
-            level_roi_scores = tf.gather(scores, selected_indices, batch_dims=1)
-            level_rois = level_rois * tf.cast(
-                tf.reshape(tf.range(level_post_nms_topk), [1, -1, 1])
-                < tf.reshape(num_valid, [-1, 1, 1]),
-                level_rois.dtype,
-            )
-            level_roi_scores = level_roi_scores * tf.cast(
-                tf.reshape(tf.range(level_post_nms_topk), [1, -1])
-                < tf.reshape(num_valid, [-1, 1]),
-                level_roi_scores.dtype,
+                confidence_threshold=nms_score_threshold,
+                max_detections=level_post_nms_topk,
+            )(
+                box_prediction=boxes,
+                class_prediction=scores[..., None],
             )
-            return level_rois, level_roi_scores
+            return boxes["boxes"], boxes["confidence"]
 
         if not isinstance(multi_level_boxes, dict):
             return per_level_gen(multi_level_boxes, multi_level_scores)
@@ -204,14 +179,14 @@ def per_level_gen(boxes, scores):
             rois.append(level_rois)
             roi_scores.append(level_roi_scores)
 
-        rois = tf.concat(rois, axis=1)
-        roi_scores = tf.concat(roi_scores, axis=1)
-        _, num_valid_rois = roi_scores.get_shape().as_list()
+        rois = ops.concatenate(rois, axis=1)
+        roi_scores = ops.concatenate(roi_scores, axis=1)
+        _, num_valid_rois = ops.shape(roi_scores)
         overall_top_k = min(num_valid_rois, post_nms_topk)
-        roi_scores, sorted_indices = tf.nn.top_k(
+        roi_scores, sorted_indices = ops.top_k(
             roi_scores, k=overall_top_k, sorted=True
         )
-        rois = tf.gather(rois, sorted_indices, batch_dims=1)
+        rois = ops.take_along_axis(rois, sorted_indices[..., None], axis=1)
 
         return rois, roi_scores
 

keras_cv/layers/object_detection/roi_generator.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import pytest
-import tensorflow as tf
 
 from keras_cv.layers.object_detection.roi_generator import ROIGenerator
 from keras_cv.tests.test_case import TestCase
@@ -23,7 +23,7 @@
 class ROIGeneratorTest(TestCase):
     def test_single_tensor(self):
         roi_generator = ROIGenerator("xyxy", nms_iou_threshold_train=0.96)
-        rpn_boxes = tf.constant(
+        rpn_boxes = np.array(
             [
                 [
                     [0, 0, 10, 10],
@@ -33,26 +33,33 @@ def test_single_tensor(self):
                 ],
             ]
         )
-        expected_rois = tf.gather(rpn_boxes, [[1, 3, 2]], batch_dims=1)
-        expected_rois = tf.concat([expected_rois, tf.zeros([1, 1, 4])], axis=1)
-        rpn_scores = tf.constant(
+        indices = [1, 3, 2]
+        expected_rois = np.take(rpn_boxes, indices, axis=1)
+        expected_rois = np.concatenate(
+            [expected_rois, -np.ones([1, 1, 4])], axis=1
+        )
+        rpn_scores = np.array(
             [
                 [0.6, 0.9, 0.2, 0.3],
             ]
         )
         # selecting the 1st, then 3rd, then 2nd as they don't overlap
         # 0th box overlaps with 1st box
-        expected_roi_scores = tf.gather(rpn_scores, [[1, 3, 2]], batch_dims=1)
-        expected_roi_scores = tf.concat(
-            [expected_roi_scores, tf.zeros([1, 1])], axis=1
+        expected_roi_scores = np.take(rpn_scores, indices, axis=1)
+        expected_roi_scores = np.concatenate(
+            [expected_roi_scores, -np.ones([1, 1])], axis=1
+        )
+        rois, roi_scores = roi_generator(
+            multi_level_boxes=rpn_boxes,
+            multi_level_scores=rpn_scores,
+            training=True,
         )
-        rois, roi_scores = roi_generator(rpn_boxes, rpn_scores, training=True)
         self.assertAllClose(expected_rois, rois)
         self.assertAllClose(expected_roi_scores, roi_scores)
 
     def test_single_level_single_batch_roi_ignore_box(self):
         roi_generator = ROIGenerator("xyxy", nms_iou_threshold_train=0.96)
-        rpn_boxes = tf.constant(
+        rpn_boxes = np.array(
             [
                 [
                     [0, 0, 10, 10],
@@ -62,19 +69,22 @@ def test_single_level_single_batch_roi_ignore_box(self):
                 ],
             ]
         )
-        expected_rois = tf.gather(rpn_boxes, [[1, 3, 2]], batch_dims=1)
-        expected_rois = tf.concat([expected_rois, tf.zeros([1, 1, 4])], axis=1)
+        indices = [1, 3, 2]
+        expected_rois = np.take(rpn_boxes, indices, axis=1)
+        expected_rois = np.concatenate(
+            [expected_rois, -np.ones([1, 1, 4])], axis=1
+        )
         rpn_boxes = {2: rpn_boxes}
-        rpn_scores = tf.constant(
+        rpn_scores = np.array(
             [
                 [0.6, 0.9, 0.2, 0.3],
             ]
         )
         # selecting the 1st, then 3rd, then 2nd as they don't overlap
         # 0th box overlaps with 1st box
-        expected_roi_scores = tf.gather(rpn_scores, [[1, 3, 2]], batch_dims=1)
-        expected_roi_scores = tf.concat(
-            [expected_roi_scores, tf.zeros([1, 1])], axis=1
+        expected_roi_scores = np.take(rpn_scores, indices, axis=1)
+        expected_roi_scores = np.concatenate(
+            [expected_roi_scores, -np.ones([1, 1])], axis=1
         )
         rpn_scores = {2: rpn_scores}
         rois, roi_scores = roi_generator(rpn_boxes, rpn_scores, training=True)
@@ -85,7 +95,7 @@ def test_single_level_single_batch_roi_all_box(self):
         # for iou between 1st and 2nd box is 0.9604, so setting to 0.97 to
         # such that NMS would treat them as different ROIs
         roi_generator = ROIGenerator("xyxy", nms_iou_threshold_train=0.97)
-        rpn_boxes = tf.constant(
+        rpn_boxes = np.array(
             [
                 [
                     [0, 0, 10, 10],
@@ -95,25 +105,24 @@ def test_single_level_single_batch_roi_all_box(self):
                 ],
             ]
         )
-        expected_rois = tf.gather(rpn_boxes, [[1, 0, 3, 2]], batch_dims=1)
+        indices = [1, 0, 3, 2]
+        expected_rois = np.take(rpn_boxes, indices, axis=1)
         rpn_boxes = {2: rpn_boxes}
-        rpn_scores = tf.constant(
+        rpn_scores = np.array(
             [
                 [0.6, 0.9, 0.2, 0.3],
             ]
         )
         # selecting the 1st, then 0th, then 3rd, then 2nd as they don't overlap
-        expected_roi_scores = tf.gather(
-            rpn_scores, [[1, 0, 3, 2]], batch_dims=1
-        )
+        expected_roi_scores = np.take(rpn_scores, indices, axis=1)
         rpn_scores = {2: rpn_scores}
         rois, roi_scores = roi_generator(rpn_boxes, rpn_scores, training=True)
         self.assertAllClose(expected_rois, rois)
         self.assertAllClose(expected_roi_scores, roi_scores)
 
     def test_single_level_propose_rois(self):
         roi_generator = ROIGenerator("xyxy")
-        rpn_boxes = tf.constant(
+        rpn_boxes = np.array(
             [
                 [
                     [0, 0, 10, 10],
@@ -129,21 +138,22 @@ def test_single_level_propose_rois(self):
                 ],
             ]
         )
-        expected_rois = tf.gather(
-            rpn_boxes, [[1, 3, 2], [1, 3, 0]], batch_dims=1
+        indices = np.array([[1, 3, 2], [1, 3, 0]])
+        expected_rois = np.take_along_axis(
+            rpn_boxes, indices[:, :, None], axis=1
+        )
+        expected_rois = np.concatenate(
+            [expected_rois, -np.ones([2, 1, 4])], axis=1
         )
-        expected_rois = tf.concat([expected_rois, tf.zeros([2, 1, 4])], axis=1)
         rpn_boxes = {2: rpn_boxes}
-        rpn_scores = tf.constant([[0.6, 0.9, 0.2, 0.3], [0.1, 0.8, 0.3, 0.5]])
+        rpn_scores = np.array([[0.6, 0.9, 0.2, 0.3], [0.1, 0.8, 0.3, 0.5]])
         # 1st batch -- selecting the 1st, then 3rd, then 2nd as they don't
         #   overlap
         # 2nd batch -- selecting the 1st, then 3rd, then 0th as they don't
         #   overlap
-        expected_roi_scores = tf.gather(
-            rpn_scores, [[1, 3, 2], [1, 3, 0]], batch_dims=1
-        )
-        expected_roi_scores = tf.concat(
-            [expected_roi_scores, tf.zeros([2, 1])], axis=1
+        expected_roi_scores = np.take_along_axis(rpn_scores, indices, axis=1)
+        expected_roi_scores = np.concatenate(
+            [expected_roi_scores, -np.ones([2, 1])], axis=1
         )
         rpn_scores = {2: rpn_scores}
         rois, roi_scores = roi_generator(rpn_boxes, rpn_scores, training=True)
@@ -152,7 +162,7 @@ def test_single_level_propose_rois(self):
 
     def test_two_level_single_batch_propose_rois_ignore_box(self):
         roi_generator = ROIGenerator("xyxy")
-        rpn_boxes = tf.constant(
+        rpn_boxes = np.array(
             [
                 [
                     [0, 0, 10, 10],
@@ -168,7 +178,7 @@ def test_two_level_single_batch_propose_rois_ignore_box(self):
                 ],
             ]
         )
-        expected_rois = tf.constant(
+        expected_rois = np.array(
             [
                 [
                     [0.1, 0.1, 9.9, 9.9],
@@ -177,13 +187,13 @@ def test_two_level_single_batch_propose_rois_ignore_box(self):
                     [2, 2, 8, 8],
                     [5, 5, 10, 10],
                     [2, 2, 4, 4],
-                    [0, 0, 0, 0],
-                    [0, 0, 0, 0],
+                    [-1, -1, -1, -1],
+                    [-1, -1, -1, -1],
                 ]
             ]
         )
         rpn_boxes = {2: rpn_boxes[0:1], 3: rpn_boxes[1:2]}
-        rpn_scores = tf.constant([[0.6, 0.9, 0.2, 0.3], [0.1, 0.8, 0.3, 0.5]])
+        rpn_scores = np.array([[0.6, 0.9, 0.2, 0.3], [0.1, 0.8, 0.3, 0.5]])
         # 1st batch -- selecting the 1st, then 3rd, then 2nd as they don't
         #   overlap
         # 2nd batch -- selecting the 1st, then 3rd, then 0th as they don't
@@ -196,8 +206,8 @@ def test_two_level_single_batch_propose_rois_ignore_box(self):
                 0.3,
                 0.2,
                 0.1,
-                0.0,
-                0.0,
+                -1.0,
+                -1.0,
             ]
         ]
         rpn_scores = {2: rpn_scores[0:1], 3: rpn_scores[1:2]}
@@ -207,7 +217,7 @@ def test_two_level_single_batch_propose_rois_ignore_box(self):
 
     def test_two_level_single_batch_propose_rois_all_box(self):
         roi_generator = ROIGenerator("xyxy", nms_iou_threshold_train=0.99)
-        rpn_boxes = tf.constant(
+        rpn_boxes = np.array(
             [
                 [
                     [0, 0, 10, 10],
@@ -223,7 +233,7 @@ def test_two_level_single_batch_propose_rois_all_box(self):
                 ],
             ]
         )
-        expected_rois = tf.constant(
+        expected_rois = np.array(
             [
                 [
                     [0.1, 0.1, 9.9, 9.9],
@@ -238,7 +248,7 @@ def test_two_level_single_batch_propose_rois_all_box(self):
             ]
         )
         rpn_boxes = {2: rpn_boxes[0:1], 3: rpn_boxes[1:2]}
-        rpn_scores = tf.constant([[0.6, 0.9, 0.2, 0.3], [0.1, 0.8, 0.3, 0.5]])
+        rpn_scores = np.array([[0.6, 0.9, 0.2, 0.3], [0.1, 0.8, 0.3, 0.5]])
         # 1st batch -- selecting the 1st, then 0th, then 3rd, then 2nd as they
         #   don't overlap
         # 2nd batch -- selecting the 1st, then 3rd, then 2nd, then 0th as they