def test_diagonal_gradient_image(self):
    """Tests if a good pyramid image is created."""
    pyramid_image = test_utils.create_diagonal_gradient_image(3, 4, 2)

    # Test which is easy to understand.
    expected_first_channel = np.array([[3, 2, 1, 0],
                                       [4, 3, 2, 1],
                                       [5, 4, 3, 2]], dtype=np.float32)
    self.assertAllEqual(np.squeeze(pyramid_image[:, :, 0]),
                        expected_first_channel)

    # Actual test.
    expected_image = np.array([[[3, 30],
                                [2, 20],
                                [1, 10],
                                [0, 0]],
                               [[4, 40],
                                [3, 30],
                                [2, 20],
                                [1, 10]],
                               [[5, 50],
                                [4, 40],
                                [3, 30],
                                [2, 20]]], dtype=np.float32)

    self.assertAllEqual(pyramid_image, expected_image)
Tests if a good pyramid image is created.
inference_api/src/main/object_detection/utils/test_utils_test.py