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