def test_assign_track_targets_weights(self):
    """Test the assign_track_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]),
      ]
      track_id_batch = [
          tf.constant([0, 1]),
          tf.constant([1, 0]),
          tf.constant([0, 2]),
      ]
      weights_batch = [
          tf.constant([0.0, 1.0]),
          tf.constant([1.0, 1.0]),
          tf.constant([0.0, 0.0])
      ]

      assigner = targetassigner.CenterNetTrackTargetAssigner(
          stride=4, num_track_ids=3)

      (batch_indices, batch_weights,
       track_targets) = assigner.assign_track_targets(
           height=80,
           width=80,
           gt_track_ids_list=track_id_batch,
           gt_boxes_list=box_batch,
           gt_weights_list=weights_batch)
      return batch_indices, batch_weights, track_targets

    indices, weights, track_ids = self.execute(graph_fn, [])

    self.assertEqual(indices.shape, (3, 2, 3))
    self.assertEqual(track_ids.shape, (3, 2, 3))
    self.assertEqual(weights.shape, (3, 2))

    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(track_ids,
                                  [[[1, 0, 0], [0, 1, 0]],
                                   [[0, 1, 0], [1, 0, 0]],
                                   [[1, 0, 0], [0, 0, 1]]])
    np.testing.assert_array_equal(weights, [[0, 1], [1, 1], [0, 0]])
    # TODO(xwwang): Add a test for the case when no objects are detected.
Test the assign_track_targets function with box weights.
inference_api/src/main/object_detection/core/target_assigner_test.py