Updating FasterRCNN to use Task API

keras_cv/layers/object_detection/roi_generator.py

@@ -23,6 +23,7 @@
 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 ops
 
 
 @keras_cv_export("keras_cv.layers.ROIGenerator")
@@ -148,7 +149,7 @@ def per_level_gen(boxes, scores):
             # scores can also be [batch_size, num_boxes, 1]
             if len(scores_shape) == 3:
                 scores = tf.squeeze(scores, axis=-1)
-            _, num_boxes = scores.get_shape().as_list()
+            num_boxes = ops.shape(boxes)[1]
             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(

keras_cv/layers/object_detection/roi_sampler.py

@@ -12,11 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import tensorflow as tf
-from tensorflow import keras
+import numpy as np  # used for newaxis
 
 from keras_cv import bounding_box
 from keras_cv.backend import assert_tf_keras
+from keras_cv.backend import keras
+from keras_cv.backend import ops
 from keras_cv.bounding_box import iou
 from keras_cv.layers.object_detection import box_matcher
 from keras_cv.layers.object_detection import sampling
@@ -84,9 +85,9 @@ def __init__(
 
     def call(
         self,
-        rois: tf.Tensor,
-        gt_boxes: tf.Tensor,
-        gt_classes: tf.Tensor,
+        rois,
+        gt_boxes,
+        gt_classes,
     ):
         """
         Args:
@@ -102,8 +103,8 @@ def call(
         """
         if self.append_gt_boxes:
             # num_rois += num_gt
-            rois = tf.concat([rois, gt_boxes], axis=1)
-        num_rois = rois.get_shape().as_list()[1]
+            rois = ops.concatenate([rois, gt_boxes], axis=1)
+        num_rois = ops.shape(rois)[1]
         if num_rois is None:
             raise ValueError(
                 f"`rois` must have static shape, got {rois.get_shape()}"
@@ -126,27 +127,27 @@ def call(
         # [batch_size, num_rois] | [batch_size, num_rois]
         matched_gt_cols, matched_vals = self.roi_matcher(similarity_mat)
         # [batch_size, num_rois]
-        positive_matches = tf.math.equal(matched_vals, 1)
-        negative_matches = tf.math.equal(matched_vals, -1)
+        positive_matches = ops.equal(matched_vals, 1)
+        negative_matches = ops.equal(matched_vals, -1)
         self._positives.update_state(
-            tf.reduce_sum(tf.cast(positive_matches, tf.float32), axis=-1)
+            ops.sum(ops.cast(positive_matches, "float32"), axis=-1)
         )
         self._negatives.update_state(
-            tf.reduce_sum(tf.cast(negative_matches, tf.float32), axis=-1)
+            ops.sum(ops.cast(negative_matches, "float32"), axis=-1)
         )
         # [batch_size, num_rois, 1]
-        background_mask = tf.expand_dims(
-            tf.logical_not(positive_matches), axis=-1
+        background_mask = ops.expand_dims(
+            ops.logical_not(positive_matches), axis=-1
         )
         # [batch_size, num_rois, 1]
         matched_gt_classes = target_gather._target_gather(
             gt_classes, matched_gt_cols
         )
         # also set all background matches to `background_class`
-        matched_gt_classes = tf.where(
+        matched_gt_classes = ops.where(
             background_mask,
-            tf.cast(
-                self.background_class * tf.ones_like(matched_gt_classes),
+            ops.cast(
+                self.background_class * ops.ones_like(matched_gt_classes),
                 gt_classes.dtype,
             ),
             matched_gt_classes,
@@ -163,9 +164,9 @@ def call(
             variance=[0.1, 0.1, 0.2, 0.2],
         )
         # also set all background matches to 0 coordinates
-        encoded_matched_gt_boxes = tf.where(
+        encoded_matched_gt_boxes = ops.where(
             background_mask,
-            tf.zeros_like(matched_gt_boxes),
+            ops.zeros_like(matched_gt_boxes),
             encoded_matched_gt_boxes,
         )
         # [batch_size, num_rois]
@@ -176,7 +177,7 @@ def call(
             self.positive_fraction,
         )
         # [batch_size, num_sampled_rois] in the range of [0, num_rois)
-        sampled_indicators, sampled_indices = tf.math.top_k(
+        sampled_indicators, sampled_indices = ops.math.top_k(
             sampled_indicators, k=self.num_sampled_rois, sorted=True
         )
         # [batch_size, num_sampled_rois, 4]
@@ -192,12 +193,12 @@ def call(
         # [batch_size, num_sampled_rois, 1]
         # all negative samples will be ignored in regression
         sampled_box_weights = target_gather._target_gather(
-            tf.cast(positive_matches[..., tf.newaxis], gt_boxes.dtype),
+            ops.cast(positive_matches[..., np.newaxis], gt_boxes.dtype),
             sampled_indices,
         )
         # [batch_size, num_sampled_rois, 1]
-        sampled_indicators = sampled_indicators[..., tf.newaxis]
-        sampled_class_weights = tf.cast(sampled_indicators, gt_classes.dtype)
+        sampled_indicators = sampled_indicators[..., np.newaxis]
+        sampled_class_weights = ops.cast(sampled_indicators, gt_classes.dtype)
         return (
             sampled_rois,
             sampled_gt_boxes,

keras_cv/layers/object_detection/rpn_label_encoder.py

@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Mapping
-
-import tensorflow as tf
-from tensorflow import keras
+import numpy as np  # Used for newaxis
+import tree
 
 from keras_cv import bounding_box
 from keras_cv.backend import assert_tf_keras
+from keras_cv.backend import keras
+from keras_cv.backend import ops
 from keras_cv.bounding_box import iou
 from keras_cv.layers.object_detection import box_matcher
 from keras_cv.layers.object_detection import sampling
@@ -92,9 +92,9 @@ def __init__(
 
     def call(
         self,
-        anchors_dict: Mapping[str, tf.Tensor],
-        gt_boxes: tf.Tensor,
-        gt_classes: tf.Tensor,
+        anchors_dict,
+        gt_boxes,
+        gt_classes,
     ):
         """
         Args:
@@ -112,7 +112,7 @@ def call(
         anchors = anchors_dict
         if isinstance(anchors, dict):
             pack = True
-            anchors = tf.concat(tf.nest.flatten(anchors), axis=0)
+            anchors = ops.concatenate(tree.flatten(anchors), axis=0)
         anchors = bounding_box.convert_format(
             anchors, source=self.anchor_format, target="yxyx"
         )
@@ -126,14 +126,14 @@ def call(
         # [num_anchors] or [batch_size, num_anchors]
         matched_gt_indices, matched_vals = self.box_matcher(similarity_mat)
         # [num_anchors] or [batch_size, num_anchors]
-        positive_matches = tf.math.equal(matched_vals, 1)
+        positive_matches = ops.equal(matched_vals, 1)
         # currently SyncOnReadVariable does not support `assign_add` in
         # cross-replica.
         #    self._positives.update_state(
         #        tf.reduce_sum(tf.cast(positive_matches, tf.float32), axis=-1)
         #    )
 
-        negative_matches = tf.math.equal(matched_vals, -1)
+        negative_matches = ops.equal(matched_vals, -1)
         # [num_anchors, 4] or [batch_size, num_anchors, 4]
         matched_gt_boxes = target_gather._target_gather(
             gt_boxes, matched_gt_indices
@@ -148,18 +148,18 @@ def call(
             variance=self.box_variance,
         )
         # [num_anchors, 1] or [batch_size, num_anchors, 1]
-        box_sample_weights = tf.cast(
-            positive_matches[..., tf.newaxis], gt_boxes.dtype
+        box_sample_weights = ops.cast(
+            positive_matches[..., np.newaxis], gt_boxes.dtype
         )
 
         # [num_anchors, 1] or [batch_size, num_anchors, 1]
-        positive_mask = tf.expand_dims(positive_matches, axis=-1)
+        positive_mask = ops.expand_dims(positive_matches, axis=-1)
         # set all negative and ignored matches to 0, and all positive matches to
         # 1 [num_anchors, 1] or [batch_size, num_anchors, 1]
-        positive_classes = tf.ones_like(positive_mask, dtype=gt_classes.dtype)
-        negative_classes = tf.zeros_like(positive_mask, dtype=gt_classes.dtype)
+        positive_classes = ops.ones_like(positive_mask, dtype=gt_classes.dtype)
+        negative_classes = ops.zeros_like(positive_mask, dtype=gt_classes.dtype)
         # [num_anchors, 1] or [batch_size, num_anchors, 1]
-        class_targets = tf.where(
+        class_targets = ops.where(
             positive_mask, positive_classes, negative_classes
         )
         # [num_anchors] or [batch_size, num_anchors]
@@ -170,8 +170,8 @@ def call(
             self.positive_fraction,
         )
         # [num_anchors, 1] or [batch_size, num_anchors, 1]
-        class_sample_weights = tf.cast(
-            sampled_indicators[..., tf.newaxis], gt_classes.dtype
+        class_sample_weights = ops.cast(
+            sampled_indicators[..., np.newaxis], gt_classes.dtype
         )
         if pack:
             encoded_box_targets = self.unpack_targets(

keras_cv/models/__init__.py

@@ -183,6 +183,7 @@
 from keras_cv.models.backbones.vit_det.vit_det_aliases import ViTDetLBackbone
 from keras_cv.models.backbones.vit_det.vit_det_backbone import ViTDetBackbone
 from keras_cv.models.classification.image_classifier import ImageClassifier
+from keras_cv.models.object_detection.faster_rcnn.faster_rcnn import FasterRCNN
 from keras_cv.models.object_detection.retinanet.retinanet import RetinaNet
 from keras_cv.models.object_detection.yolo_v8.yolo_v8_backbone import (
     YOLOV8Backbone,

keras_cv/models/legacy/__init__.py

@@ -27,9 +27,6 @@
 from keras_cv.models.legacy.mlp_mixer import MLPMixerB16
 from keras_cv.models.legacy.mlp_mixer import MLPMixerB32
 from keras_cv.models.legacy.mlp_mixer import MLPMixerL16
-from keras_cv.models.legacy.object_detection.faster_rcnn.faster_rcnn import (
-    FasterRCNN,
-)
 from keras_cv.models.legacy.regnet import RegNetX002
 from keras_cv.models.legacy.regnet import RegNetX004
 from keras_cv.models.legacy.regnet import RegNetX006

keras_cv/models/legacy/object_detection/__init__.py → keras_cv/models/object_detection/faster_rcnn/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# Copyright 2023 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.
@@ -11,3 +11,8 @@
 # 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.
+from keras_cv.models.object_detection.faster_rcnn.feature_pyramid import (
+    FeaturePyramid,
+)
+from keras_cv.models.object_detection.faster_rcnn.rcnn_head import RCNNHead
+from keras_cv.models.object_detection.faster_rcnn.rpn_head import RPNHead