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)