Bug Fix for CenterNetBoxLoss (#2432)

* fix CenterNetBoxLoss, add Testcase

* cleanup testcase

* test case code clean up

* fix Pytorch pipeline

* make test_heading_regression_loss framework agnostic

---------

Co-authored-by: Till Beemelmanns <[email protected]>

Author: TillBeemelmanns

keras_cv/src/losses/centernet_box_loss.py

@@ -65,14 +65,14 @@ def __init__(self, num_heading_bins, anchor_size, **kwargs):
         self.anchor_size = anchor_size
 
     def heading_regression_loss(self, heading_true, heading_pred):
+        heading_pred = ops.convert_to_tensor(heading_pred)
+
         # Set the heading to within 0 -> 2pi
-        heading_true = ops.floor(ops.mod(heading_true, 2 * math.pi))
+        heading_true = ops.mod(heading_true, 2 * math.pi)
 
         # Divide 2pi into bins. shifted by 0.5 * angle_per_class.
         angle_per_class = (2 * math.pi) / self.num_heading_bins
-        shift_angle = ops.floor(
-            ops.mod(heading_true + angle_per_class / 2, 2 * math.pi)
-        )
+        shift_angle = ops.mod(heading_true + angle_per_class / 2, 2 * math.pi)
 
         heading_bin_label_float = ops.floor(
             ops.divide(shift_angle, angle_per_class)

keras_cv/src/losses/centernet_box_loss_test.py

@@ -16,6 +16,7 @@
 from absl.testing import parameterized
 
 import keras_cv
+from keras_cv.src.backend import ops
 from keras_cv.src.tests.test_case import TestCase
 
 
@@ -42,3 +43,28 @@ def test_proper_output_shapes(self, reduction, target_size):
             y_pred=np.random.uniform(size=(2, 10, 6 + 2 * 4)),
         )
         self.assertEqual(result.shape, target_size)
+
+    def test_heading_regression_loss(self):
+        num_heading_bins = 4
+        loss = keras_cv.losses.CenterNetBoxLoss(
+            num_heading_bins=num_heading_bins, anchor_size=[1.0, 1.0, 1.0]
+        )
+        heading_true = np.array(
+            [[0, (1 / 2.0) * np.pi, np.pi, (3.0 / 2.0) * np.pi]]
+        )
+        heading_pred = np.array(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
+                ]
+            ]
+        )
+        heading_loss = loss.heading_regression_loss(
+            heading_true=heading_true, heading_pred=heading_pred
+        )
+        ce_loss = -np.log(np.exp(1) / np.exp([1, 0, 0, 0]).sum())
+        expected_loss = ce_loss * num_heading_bins
+        self.assertAllClose(ops.sum(heading_loss), expected_loss)