def test_assign_corner_offset_multiple_objects(self):
    """Test corner offsets are correct with multiple objects."""
    assigner = targetassigner.CenterNetCornerOffsetTargetAssigner(stride=1)

    def graph_fn():
      boxes = [
          tf.constant([[0., 0., 1., 1.], [0., 0., 0., 0.]]),
          tf.constant([[0., 0., .25, .25], [.25, .25, 1., 1.]])
      ]
      mask1 = np.zeros((2, 4, 4), dtype=np.float32)
      mask1[0, 0, 0] = 1.0
      mask1[0, 3, 3] = 1.0

      mask2 = np.zeros((2, 4, 4), dtype=np.float32)
      mask2[0, :2, :2] = 1.0
      mask2[1, 1:, 1:] = 1.0

      masks = [tf.constant(mask1), tf.constant(mask2)]
      return assigner.assign_corner_offset_targets(boxes, masks)

    corner_offsets, foreground = self.execute(graph_fn, [])
    self.assertEqual(corner_offsets.shape, (2, 4, 4, 4))
    self.assertEqual(foreground.shape, (2, 4, 4))

    self.assertAllClose(foreground[0],
                        [[1, 0, 0, 0],
                         [0, 0, 0, 0],
                         [0, 0, 0, 0],
                         [0, 0, 0, 1]])

    self.assertAllClose(corner_offsets[0, :, :, 0],
                        [[0, 0, 0, 0],
                         [0, 0, 0, 0],
                         [0, 0, 0, 0],
                         [0, 0, 0, -3]])
    self.assertAllClose(corner_offsets[0, :, :, 1],
                        [[0, 0, 0, 0],
                         [0, 0, 0, 0],
                         [0, 0, 0, 0],
                         [0, 0, 0, -3]])
    self.assertAllClose(corner_offsets[0, :, :, 2],
                        [[4, 0, 0, 0],
                         [0, 0, 0, 0],
                         [0, 0, 0, 0],
                         [0, 0, 0, 1]])
    self.assertAllClose(corner_offsets[0, :, :, 3],
                        [[4, 0, 0, 0],
                         [0, 0, 0, 0],
                         [0, 0, 0, 0],
                         [0, 0, 0, 1]])

    self.assertAllClose(foreground[1],
                        [[1, 1, 0, 0],
                         [1, 1, 1, 1],
                         [0, 1, 1, 1],
                         [0, 1, 1, 1]])

    self.assertAllClose(corner_offsets[1, :, :, 0],
                        [[0, 0, 0, 0],
                         [-1, -1, 0, 0],
                         [0, -1, -1, -1],
                         [0, -2, -2, -2]])
    self.assertAllClose(corner_offsets[1, :, :, 1],
                        [[0, -1, 0, 0],
                         [0, -1, -1, -2],
                         [0, 0, -1, -2],
                         [0, 0, -1, -2]])
    self.assertAllClose(corner_offsets[1, :, :, 2],
                        [[1, 1, 0, 0],
                         [0, 0, 3, 3],
                         [0, 2, 2, 2],
                         [0, 1, 1, 1]])
    self.assertAllClose(corner_offsets[1, :, :, 3],
                        [[1, 0, 0, 0],
                         [1, 0, 2, 1],
                         [0, 3, 2, 1],
                         [0, 3, 2, 1]])
Test corner offsets are correct with multiple objects.
inference_api/src/main/object_detection/core/target_assigner_test.py