1import io
2from functools import partial
3
4from pytest_regressions.common import perform_regression_check, import_error_message
5
6
7class ImageRegressionFixture:
8    """
9    Regression test for image objects, accounting for small differences.
10    """
11
12    def __init__(self, datadir, original_datadir, request):
13        """
14        :type datadir: Path
15        :type original_datadir: Path
16        :type request: FixtureRequest
17        """
18        self.request = request
19        self.datadir = datadir
20        self.original_datadir = original_datadir
21        self.force_regen = False
22
23    def _load_image(self, filename):
24        """
25        Reads the image from the given file and convert it to RGB if necessary.
26        This is necessary to be used with the ImageChops module operations.
27        At this time, in this module, channel operations are only implemented
28        for 8-bit images (e.g. "L" and "RGB").
29
30        :param Path filename:
31            The name of the file
32        """
33        try:
34            from PIL import Image
35        except ModuleNotFoundError:
36            raise ModuleNotFoundError(import_error_message("Pillow"))
37
38        img = Image.open(str(filename), "r")
39        if img.mode not in ("L" or "RGB"):
40            return img.convert("RGB")
41        else:
42            return img
43
44    def _compute_manhattan_distance(self, diff_image):
45        """
46        Computes a percentage of similarity of the difference image given.
47
48        :param PIL.Image diff_image:
49            An image in RGB mode computed from ImageChops.difference
50        """
51        try:
52            import numpy
53        except ModuleNotFoundError:
54            raise ModuleNotFoundError(import_error_message("Numpy"))
55
56        number_of_pixels = diff_image.size[0] * diff_image.size[1]
57        return (
58            # To obtain a number in 0.0 -> 100.0
59            100.0
60            * (
61                # Compute the sum of differences
62                numpy.sum(diff_image)
63                /
64                # Divide by the number of channel differences RGB * Pixels
65                float(3 * number_of_pixels)
66            )
67            # Normalize between 0.0 -> 1.0
68            / 255.0
69        )
70
71    def _check_images_manhattan_distance(
72        self, obtained_file, expected_file, expect_equal, diff_threshold
73    ):
74        """
75        Compare two image by computing the differences spatially, pixel by pixel.
76
77        The Manhattan Distance is used to compute how much two images differ.
78
79        :param str obtained_file:
80            The image with the obtained image
81
82        :param str expected_files:
83            The image with the expected image
84
85        :param bool expected_equal:
86            If True, the images are expected to be equal, otherwise, they're expected to be
87            different.
88
89        :param float diff_threshold:
90            The maximum percentage of difference accepted.
91            A value between 0.0 and 100.0
92
93        :raises AssertionError:
94            raised if they are actually different and expect_equal is False or
95            if they are equal and expect_equal is True.
96        """
97        try:
98            from PIL import ImageChops
99        except ModuleNotFoundError:
100            raise ModuleNotFoundError(import_error_message("Pillow"))
101
102        __tracebackhide__ = True
103
104        obtained_img = self._load_image(obtained_file)
105        expected_img = self._load_image(expected_file)
106
107        def check_result(equal, manhattan_distance):
108            if equal != expect_equal:
109                params = manhattan_distance, expected_file, obtained_file
110                if expect_equal:
111                    assert 0, (
112                        "Difference between images too high: %.2f %%\n%s\n%s" % params
113                    )
114                else:
115                    assert 0, (
116                        "Difference between images too small: %.2f %%\n%s\n%s" % params
117                    )
118
119        # 1st check: identical
120        diff_img = ImageChops.difference(obtained_img, expected_img)
121
122        if diff_img.getbbox() is None:  # Equal
123            check_result(True, None)
124
125        manhattan_distance = self._compute_manhattan_distance(diff_img)
126        equal = manhattan_distance <= diff_threshold
127        check_result(equal, manhattan_distance)
128
129    def check(self, image_data, diff_threshold=0.1, expect_equal=True, basename=None):
130        """
131        Checks that the given image contents are comparable with the ones stored in the data directory.
132
133        :param bytes image_data: image data
134        :param str|None basename: basename to store the information in the data directory. If none, use the name
135            of the test function.
136        :param bool expect_equal: if the image should considered equal below of the given threshold. If False, the
137            image should be considered different at least above the threshold.
138        :param float diff_threshold:
139            Tolerage as a percentage (1 to 100) on how the images are allowed to differ.
140
141        """
142        __tracebackhide__ = True
143
144        try:
145            from PIL import Image
146        except ModuleNotFoundError:
147            raise ModuleNotFoundError(import_error_message("Pillow"))
148
149        def dump_fn(target):
150            image = Image.open(io.BytesIO(image_data))
151            image.save(str(target), "PNG")
152
153        perform_regression_check(
154            datadir=self.datadir,
155            original_datadir=self.original_datadir,
156            request=self.request,
157            check_fn=partial(
158                self._check_images_manhattan_distance,
159                diff_threshold=diff_threshold,
160                expect_equal=expect_equal,
161            ),
162            dump_fn=dump_fn,
163            extension=".png",
164            basename=basename,
165            force_regen=self.force_regen,
166        )
167