def test_weights(self):
    """Test that the weights correctly ignore ground truth."""
    def graph1_fn():
      box_batch = [
          tf.constant([self._box_center, self._box_lower_left]),
          tf.constant([self._box_center]),
          tf.constant([self._box_center_small]),
      ]
      classes = [
          tf.one_hot([0, 1], depth=4),
          tf.one_hot([2], depth=4),
          tf.one_hot([3], depth=4),
      ]
      assigner = targetassigner.CenterNetCenterHeatmapTargetAssigner(4)
      targets = assigner.assign_center_targets_from_boxes(80, 80, box_batch,
                                                          classes)
      return targets

    targets = self.execute(graph1_fn, [])
    self.assertAlmostEqual(1.0, targets[0, :, :, 0].max())
    self.assertAlmostEqual(1.0, targets[0, :, :, 1].max())
    self.assertAlmostEqual(1.0, targets[1, :, :, 2].max())
    self.assertAlmostEqual(1.0, targets[2, :, :, 3].max())
    self.assertAlmostEqual(0.0, targets[0, :, :, [2, 3]].max())
    self.assertAlmostEqual(0.0, targets[1, :, :, [0, 1, 3]].max())
    self.assertAlmostEqual(0.0, targets[2, :, :, :3].max())

    def graph2_fn():
      weights = [
          tf.constant([0., 1.]),
          tf.constant([1.]),
          tf.constant([1.]),
      ]
      box_batch = [
          tf.constant([self._box_center, self._box_lower_left]),
          tf.constant([self._box_center]),
          tf.constant([self._box_center_small]),
      ]
      classes = [
          tf.one_hot([0, 1], depth=4),
          tf.one_hot([2], depth=4),
          tf.one_hot([3], depth=4),
      ]
      assigner = targetassigner.CenterNetCenterHeatmapTargetAssigner(4)
      targets = assigner.assign_center_targets_from_boxes(80, 80, box_batch,
                                                          classes,
                                                          weights)
      return targets
    targets = self.execute(graph2_fn, [])
    self.assertAlmostEqual(1.0, targets[0, :, :, 1].max())
    self.assertAlmostEqual(1.0, targets[1, :, :, 2].max())
    self.assertAlmostEqual(1.0, targets[2, :, :, 3].max())
    self.assertAlmostEqual(0.0, targets[0, :, :, [0, 2, 3]].max())
    self.assertAlmostEqual(0.0, targets[1, :, :, [0, 1, 3]].max())
    self.assertAlmostEqual(0.0, targets[2, :, :, :3].max())
Test that the weights correctly ignore ground truth.
inference_api/src/main/object_detection/core/target_assigner_test.py