Remove gt_ prefix from OD models compute_loss (keras-team#1304)

Author: ianstenbit

keras_cv/models/object_detection/faster_rcnn.py

@@ -437,7 +437,7 @@ def compile(
         }
         super().compile(loss=losses, **kwargs)
 
-    def compute_loss(self, images, gt_boxes, gt_classes, training):
+    def compute_loss(self, images, boxes, classes, training):
         image_shape = tf.shape(images[0])
         local_batch = images.get_shape().as_list()[0]
         if tf.distribute.has_strategy():
@@ -452,7 +452,7 @@ def compute_loss(self, images, gt_boxes, gt_classes, training):
             rpn_cls_targets,
             rpn_cls_weights,
         ) = self.rpn_labeler(
-            tf.concat(tf.nest.flatten(anchors), axis=0), gt_boxes, gt_classes
+            tf.concat(tf.nest.flatten(anchors), axis=0), boxes, classes
         )
         rpn_box_weights /= self.rpn_labeler.samples_per_image * global_batch * 0.25
         rpn_cls_weights /= self.rpn_labeler.samples_per_image * global_batch
@@ -461,7 +461,7 @@ def compute_loss(self, images, gt_boxes, gt_classes, training):
         )
         rois = tf.stop_gradient(rois)
         rois, box_targets, box_weights, cls_targets, cls_weights = self.roi_sampler(
-            rois, gt_boxes, gt_classes
+            rois, boxes, classes
         )
         box_weights /= self.roi_sampler.num_sampled_rois * global_batch * 0.25
         cls_weights /= self.roi_sampler.num_sampled_rois * global_batch
@@ -492,16 +492,16 @@ def train_step(self, data):
         images, y, sample_weight = tf.keras.utils.unpack_x_y_sample_weight(data)
         if sample_weight is not None:
             raise ValueError("`sample_weight` is currently not supported.")
-        gt_boxes = y["boxes"]
+        boxes = y["boxes"]
         if len(y["classes"].shape) != 2:
             raise ValueError(
                 "Expected 'classes' to be a tf.Tensor of rank 2. "
                 f"Got y['classes'].shape={y['classes'].shape}."
             )
         # TODO(tanzhenyu): remove this hack and perform broadcasting elsewhere
-        gt_classes = tf.expand_dims(y["classes"], axis=-1)
+        classes = tf.expand_dims(y["classes"], axis=-1)
         with tf.GradientTape() as tape:
-            total_loss = self.compute_loss(images, gt_boxes, gt_classes, training=True)
+            total_loss = self.compute_loss(images, boxes, classes, training=True)
             reg_losses = []
             if self.weight_decay:
                 for var in self.trainable_variables:
@@ -516,14 +516,14 @@ def test_step(self, data):
         images, y, sample_weight = tf.keras.utils.unpack_x_y_sample_weight(data)
         if sample_weight is not None:
             raise ValueError("`sample_weight` is currently not supported.")
-        gt_boxes = y["boxes"]
+        boxes = y["boxes"]
         if len(y["classes"].shape) != 2:
             raise ValueError(
                 "Expected 'classes' to be a tf.Tensor of rank 2. "
                 f"Got y['classes'].shape={y['classes'].shape}."
             )
-        gt_classes = tf.expand_dims(y["classes"], axis=-1)
-        self.compute_loss(images, gt_boxes, gt_classes, training=False)
+        classes = tf.expand_dims(y["classes"], axis=-1)
+        self.compute_loss(images, boxes, classes, training=False)
         return self.compute_metrics(images, {}, {}, sample_weight={})
 
     def make_predict_function(self, force=False):

keras_cv/models/object_detection/retina_net/retina_net.py

@@ -316,12 +316,12 @@ def compile(
         }
         super().compile(loss=losses, **kwargs)
 
-    def compute_loss(self, images, gt_boxes, gt_classes, training):
+    def compute_loss(self, images, boxes, classes, training):
         box_pred, cls_pred = self._forward(images, training=training)
-        if gt_boxes.shape[-1] != 4:
+        if boxes.shape[-1] != 4:
             raise ValueError(
-                "gt_boxes should have shape (None, None, 4).  Got "
-                f"gt_boxes.shape={tuple(gt_boxes.shape)}"
+                "boxes should have shape (None, None, 4).  Got "
+                f"boxes.shape={tuple(boxes.shape)}"
             )
 
         if box_pred.shape[-1] != 4:
@@ -338,18 +338,18 @@ def compute_loss(self, images, gt_boxes, gt_classes, training):
             )
 
         cls_labels = tf.one_hot(
-            tf.cast(gt_classes, dtype=tf.int32),
+            tf.cast(classes, dtype=tf.int32),
             depth=self.classes,
             dtype=tf.float32,
         )
 
-        positive_mask = tf.cast(tf.greater(gt_classes, -1.0), dtype=tf.float32)
+        positive_mask = tf.cast(tf.greater(classes, -1.0), dtype=tf.float32)
         normalizer = tf.reduce_sum(positive_mask)
-        cls_weights = tf.cast(tf.math.not_equal(gt_classes, -2.0), dtype=tf.float32)
+        cls_weights = tf.cast(tf.math.not_equal(classes, -2.0), dtype=tf.float32)
         cls_weights /= normalizer
         box_weights = positive_mask / normalizer
         y_true = {
-            "box": gt_boxes,
+            "box": boxes,
             "cls": cls_labels,
         }
         y_pred = {
@@ -372,16 +372,16 @@ def train_step(self, data):
             target=self.label_encoder.bounding_box_format,
             images=x,
         )
-        gt_boxes, gt_classes = self.label_encoder(x, y)
-        gt_boxes = bounding_box.convert_format(
-            gt_boxes,
+        boxes, classes = self.label_encoder(x, y)
+        boxes = bounding_box.convert_format(
+            boxes,
             source=self.label_encoder.bounding_box_format,
             target=self.bounding_box_format,
             images=x,
         )
 
         with tf.GradientTape() as tape:
-            total_loss = self.compute_loss(x, gt_boxes, gt_classes, training=True)
+            total_loss = self.compute_loss(x, boxes, classes, training=True)
 
             reg_losses = []
             if self.weight_decay:
@@ -405,14 +405,14 @@ def test_step(self, data):
             target=self.label_encoder.bounding_box_format,
             images=x,
         )
-        gt_boxes, gt_classes = self.label_encoder(x, y)
-        gt_boxes = bounding_box.convert_format(
-            gt_boxes,
+        boxes, classes = self.label_encoder(x, y)
+        boxes = bounding_box.convert_format(
+            boxes,
             source=self.label_encoder.bounding_box_format,
             target=self.bounding_box_format,
             images=x,
         )
-        _ = self.compute_loss(x, gt_boxes, gt_classes, training=False)
+        _ = self.compute_loss(x, boxes, classes, training=False)
 
         return self.compute_metrics(x, {}, {}, sample_weight={})