def test_assign_size_and_offset_targets_weights(self):
    """Test the assign_size_and_offset_targets function with box weights."""
    def graph_fn():
      box_batch = [
          tf.constant([self._box_center, self._box_lower_left]),
          tf.constant([self._box_lower_left, self._box_center_small]),
          tf.constant([self._box_center_small, self._box_odd_coordinates]),
      ]

      cn_assigner = targetassigner.CenterNetBoxTargetAssigner(4)
      weights_batch = [
          tf.constant([0.0, 1.0]),
          tf.constant([1.0, 1.0]),
          tf.constant([0.0, 0.0])
      ]
      indices, hw, yx_offset, weights = cn_assigner.assign_size_and_offset_targets(
          80, 80, box_batch, weights_batch)
      return indices, hw, yx_offset, weights
    indices, hw, yx_offset, weights = self.execute(graph_fn, [])
    self.assertEqual(indices.shape, (6, 3))
    self.assertEqual(hw.shape, (6, 2))
    self.assertEqual(yx_offset.shape, (6, 2))
    self.assertEqual(weights.shape, (6,))
    np.testing.assert_array_equal(indices,
                                  [[0, 10, 10], [0, 15, 5], [1, 15, 5],
                                   [1, 10, 10], [2, 10, 10], [2, 7, 11]])
    np.testing.assert_array_equal(
        hw, [[20, 20], [10, 10], [10, 10], [10, 10], [10, 10], [8, 15]])
    np.testing.assert_array_equal(
        yx_offset, [[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0.25, 0.75]])
    np.testing.assert_array_equal(weights, [0, 1, 1, 1, 0, 0])
Test the assign_size_and_offset_targets function with box weights.
inference_api/src/main/object_detection/core/target_assigner_test.py