A pytest fixture for image similarity

When testing codepaths that generate images, one might want to ensure that the generated image is what is expected. Matplotlib has a nice decorator @image_comparison that can be applied for this purpose, but looking at the implementation, it’s pretty tied to the matplotlib Figure object. I wanted something generic to use with PNGs.

I ended up writing a pytest fixture that would compare the image generated during the test with a baseline image (in tests/baseline_images as in the matplotlib implementatino). Here are the contents of conftest.py, which contain the fixture and its related image similarity assert:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33


import os
import pytest
import numpy as np
from PIL import Image


def assert_images_equal(image_1: str, image_2: str):
    img1 = Image.open(image_1)
    img2 = Image.open(image_2)

    # Convert to same mode and size for comparison
    img2 = img2.convert(img1.mode)
    img2 = img2.resize(img1.size)

    sum_sq_diff = np.sum((np.asarray(img1).astype('float') - np.asarray(img2).astype('float'))**2)

    if sum_sq_diff == 0:
        # Images are exactly the same
        pass
    else:
        normalized_sum_sq_diff = sum_sq_diff / np.sqrt(sum_sq_diff)
        assert normalized_sum_sq_diff < 0.001


@pytest.fixture
def image_similarity(request, tmpdir):
    testname = request.node.name
    filename = "{}.png".format(testname)
    generated_file = os.path.join(str(tmpdir), "{}.png".format(testname))

    yield {'filename': generated_file}

    assert_images_equal("tests/baseline_images/{}.png".format(testname), generated_file)

The assert rescales the images to be the same size, as well as the same mode, and then computes the sum of the squared differences as an image similarity metric.

You can use the above fixture in a test via:

1
2
3
4


def test_example(image_similarity):
    # test logic goes here and should generate an image in the
    # path given by image_similarity['filename']
    pass

When you add a new test, you need to add the expected image to tests/baseline_images/<testname>.png.

A pytest fixture for image similarity

See Also: