Major rework: more pythonic API, code rework

pixelmatch.py

@@ -1,45 +1,105 @@
-DEFAULT_OPTIONS = {
-    "threshold": 0.1,  # matching threshold (0 to 1); smaller is more sensitive
-    "includeAA": False,  # whether to skip anti-aliasing detection
-    "alpha": 0.1,  # opacity of original image in diff ouput
-    "aa_color": [255, 255, 0],  # color of anti-aliased pixels in diff output
-    "diff_color": [255, 0, 0],  # color of different pixels in diff output
-    "diff_mask": False,  # draw the diff over a transparent background (a mask)
-}
+import functools
+from typing import Sequence, Union, List, Tuple
+
+
+ImageData = List[int]
+PILImageOrImageArray = Union["Image", Sequence[float]]
+RGBTuple = Tuple[int, int, int]
+
+
+def pixelmatch(
+    img1: PILImageOrImageArray,
+    img2: PILImageOrImageArray,
+    width: int = None,
+    height: int = None,
+    output: PILImageOrImageArray = None,
+    threshold: float = 0.1,
+    includeAA: bool = False,
+    alpha: float = 0.1,
+    aa_color: RGBTuple = (255, 255, 0),
+    diff_color: RGBTuple = (255, 0, 0),
+    diff_mask: bool = False,
+):
+    """
+    Compares two images, writes the output diff and returns the number of mismatched pixels.
+    'Raw image data' refers to a 1D, indexable collection of image data in the
+    format [R1, G1, B1, A1, R2, G2, ...].
+
+    :param img1: Image data to compare with img2. Can be PIL.Image or raw image data.
+        Must be the same size as img2
+    :param img2: Image data to compare with img2. Can be PIL.Image or raw image data.
+        Must be the same size as img1
+    :param width: Width of both images (they should be the same). If img1 or img2 is an instance of PIL.Image,
+        the width will be extracted and used.
+    :param height: Height of both images (they should be the same). If img1 or img2 is an instance of PIL.Image,
+        the height will be extracted and used.
+    :param output: Image data to write the diff to. Should be the same size as
+    :param threshold: matching threshold (0 to 1); smaller is more sensitive, defaults to 1
+    :param includeAA: whether or not to skip anti-aliasing detection, ie if includeAA is True,
+        detecting and ignoring anti-aliased pixels is disabled. Defaults to False
+    :param alpha: opacity of original image in diff output, defaults to 0.1
+    :param aa_color: tuple of RGB color of anti-aliased pixels in diff output,
+        defaults to (255, 255, 0) (yellow)
+    :param diff_color: tuple of RGB color of the color of different pixels in diff output,
+        defaults to (255, 0, 0) (red)
+    :param diff_mask: whether or not to draw the diff over a transparent background (a mask),
+        defaults to False
+    :return: number of pixels that are different
+    """
+
+    img1_size, img1 = extract_size_and_convert_to_image_array(img1)
+    img2_size, img2 = extract_size_and_convert_to_image_array(img2)
 
+    width1 = height1 = width2 = height2 = None
+    if img1_size:
+        width1, height1 = img1_size
 
-def pixelmatch(img1, img2, width: int, height: int, output=None, options=None):
+    if img2_size:
+        width2, height2 = img2_size
 
-    if len(img1) != len(img2) or (output and len(output) != len(img1)):
-        raise ValueError("Image sizes do not match.", len(img1), len(img2), len(output))
+    output_image = None
+    if is_PIL_image(output):
+        output_image = output
+        _, output = extract_size_and_convert_to_image_array(output_image)
+
+    if len(img1) != len(img2):
+        raise ValueError("Image sizes do not match.", len(img1), len(img2))
+    if output and len(output) != len(img1):
+        raise ValueError(
+            "Diff image size does not match img1 & img2.", len(img1), len(output)
+        )
 
+    width = width or width1 or width2
+    height = height or height1 or height2
+    if width is None or height is None:
+        raise ValueError(
+            "Width or height couldn't be determined from the image input. Width and Height are only "
+            "automatically calculated when img1 or img2 is a PIL.Image"
+        )
     if len(img1) != width * height * 4:
         raise ValueError(
             "Image data size does not match width/height.",
             len(img1),
             width * height * 4,
         )
 
-    if options:
-        options = {**DEFAULT_OPTIONS, **options}
-    else:
-        options = DEFAULT_OPTIONS
-
     # fast path if identical
     if img1 == img2:
-        if output and not options["diff_mask"]:
+        if output and not diff_mask:
             for i in range(width * height):
-                draw_gray_pixel(img1, 4 * i, options["alpha"], output)
+                draw_gray_pixel(img1, 4 * i, alpha, output)
+            if output_image is not None:
+                output_image.putdata(to_PIL_image_data(output))
 
         return 0
 
     # maximum acceptable square distance between two colors;
     # 35215 is the maximum possible value for the YIQ difference metric
-    maxDelta = 35215 * options["threshold"] * options["threshold"]
+    maxDelta = 35215 * threshold * threshold
 
     diff = 0
-    [aaR, aaG, aaB] = options["aa_color"]
-    [diffR, diffG, diffB] = options["diff_color"]
+    aaR, aaG, aaB = aa_color
+    diffR, diffG, diffB = diff_color
 
     # compare each pixel of one image against the other one
     for y in range(height):
@@ -52,13 +112,13 @@ def pixelmatch(img1, img2, width: int, height: int, output=None, options=None):
             # the color difference is above the threshold
             if delta > maxDelta:
                 # check it's a real rendering difference or just anti-aliasing
-                if not options["includeAA"] and (
+                if not includeAA and (
                     antialiased(img1, x, y, width, height, img2)
                     or antialiased(img2, x, y, width, height, img1)
                 ):
                     # one of the pixels is anti-aliasing; draw as yellow and do not count as difference
                     # note that we do not include such pixels in a mask
-                    if output and not options["diff_mask"]:
+                    if output and not diff_mask:
                         draw_pixel(output, pos, aaR, aaG, aaB)
                 else:
                     # found substantial difference not caused by anti-aliasing; draw it as red
@@ -68,8 +128,11 @@ def pixelmatch(img1, img2, width: int, height: int, output=None, options=None):
 
             elif output:
                 # pixels are similar; draw background as grayscale image blended with white
-                if not options["diff_mask"]:
-                    draw_gray_pixel(img1, pos, options["alpha"], output)
+                if not diff_mask:
+                    draw_gray_pixel(img1, pos, alpha, output)
+
+    if output_image is not None:
+        output_image.putdata(to_PIL_image_data(output))
 
     # return the number of different pixels
     return diff
@@ -86,12 +149,7 @@ def antialiased(img, x1, y1, width, height, img2):
     y2 = min(y1 + 1, height - 1)
     pos = (y1 * width + x1) * 4
     zeroes = (x1 == x0 or x1 == x2 or y1 == y0 or y1 == y2) and 1 or 0
-    min_delta = 0
-    max_delta = 0
-    min_x = 0
-    min_y = 0
-    max_x = 0
-    max_y = 0
+    min_delta = max_delta = min_x = min_y = max_x = max_y = 0
 
     # go through 8 adjacent pixels
     for x in range(x0, x2 + 1):
@@ -154,12 +212,7 @@ def has_many_siblings(img, x1, y1, width, height):
                 continue
 
             pos2 = (y * width + x) * 4
-            if (
-                img[pos] == img[pos2]
-                and img[pos + 1] == img[pos2 + 1]
-                and img[pos + 2] == img[pos2 + 2]
-                and img[pos + 3] == img[pos2 + 3]
-            ):
+            if all(img[pos + offset] == img[pos2 + offset] for offset in range(4)):
                 zeroes += 1
 
             if zeroes > 2:
@@ -168,36 +221,24 @@ def has_many_siblings(img, x1, y1, width, height):
     return False
 
 
-def color_delta(img1, img2, k, m, y_only=False):
+def color_delta(img1: ImageData, img2: ImageData, k: int, m: int, y_only: bool = False):
     """
     calculate color difference according to the paper "Measuring perceived color difference
     using YIQ NTSC transmission color space in mobile applications" by Y. Kotsarenko and F. Ramos
     """
-
-    r1 = img1[k + 0]
-    g1 = img1[k + 1]
-    b1 = img1[k + 2]
-    a1 = img1[k + 3]
-
-    r2 = img2[m + 0]
-    g2 = img2[m + 1]
-    b2 = img2[m + 2]
-    a2 = img2[m + 3]
+    r1, g1, b1, a1 = [img1[k + offset] for offset in range(4)]
+    r2, g2, b2, a2 = [img2[m + offset] for offset in range(4)]
 
     if a1 == a2 and r1 == r2 and g1 == g2 and b1 == b2:
         return 0
 
     if a1 < 255:
         a1 /= 255
-        r1 = blend(r1, a1)
-        g1 = blend(g1, a1)
-        b1 = blend(b1, a1)
+        r1, b1, g1 = blendRGB(r1, b1, g1, a1)
 
     if a2 < 255:
         a2 /= 255
-        r2 = blend(r2, a2)
-        g2 = blend(g2, a2)
-        b2 = blend(b2, a2)
+        r2, b2, g2 = blendRGB(r2, b2, g2, a2)
 
     y = rgb2y(r1, g1, b1) - rgb2y(r2, g2, b2)
 
@@ -223,7 +264,11 @@ def rgb2q(r: int, g: int, b: int):
     return r * 0.21147017 - g * 0.52261711 + b * 0.31114694
 
 
-def blend(c, a):
+def blendRGB(r: int, g: int, b: int, a):
+    return blend(r, a), blend(g, a), blend(b, a)
+
+
+def blend(c: float, a):
     """blend semi-transparent color with white"""
     return 255 + (c - 255) * a
 
@@ -241,3 +286,25 @@ def draw_gray_pixel(img, i: int, alpha, output):
     b = img[i + 2]
     val = blend(rgb2y(r, g, b), alpha * img[i + 3] / 255)
     draw_pixel(output, i, val, val, val)
+
+
+def to_PIL_image_data(raw: ImageData):
+    return [*zip(raw[::4], raw[1::4], raw[2::4], raw[3::4])]
+
+
+def extract_size_and_convert_to_image_array(img: PILImageOrImageArray):
+    """
+    Takes a img of type PILImageOrImageArray and extracts the size information from it if possible.
+    :param img: PIL.Image or sequence in the format [R1, G1, B1, A1, R2, ...]
+    :return: tuple of size, list of image data in the format [R1, G1, B1, A1, R2, ...]
+    """
+    if is_PIL_image(img):
+        return img.size, [x for p in img.convert("RGBA").getdata() for x in p]
+    return None, img
+
+
+def is_PIL_image(img: PILImageOrImageArray):
+    return (
+        all(hasattr(img, attr) for attr in ["convert", "getdata", "putdata"])
+        and "Image" in type(img).__name__
+    )

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pixelmatch"
-version = "0.1.1"
+version = "1.0.0"
 description = "A pixel-level image comparison library. Python port of https://github.com/mapbox/pixelmatch"
 authors = ["Wu Haotian <whtsky@gmail.com>"]
 license = "ISC"

test_pixelmatch.py

@@ -1,3 +1,4 @@
+import itertools
 from pathlib import Path
 
 import pytest
@@ -65,10 +66,46 @@ def test_pixelmatch(
     img2_data = pil_to_flatten_data(img2)
     diff_data = [0] * len(img1_data)
 
-    mismatch = pixelmatch(img1_data, img2_data, width, height, diff_data, options)
-    mismatch2 = pixelmatch(img1_data, img2_data, width, height, None, options)
+    mismatch = pixelmatch(img1_data, img2_data, width, height, diff_data, **options)
+    mismatch2 = pixelmatch(img1_data, img2_data, width, height, None, **options)
 
     expected_diff = read_img(diff_path)
     assert diff_data == pil_to_flatten_data(expected_diff), "diff image"
     assert mismatch == expected_mismatch, "number of mismatched pixels"
     assert mismatch == mismatch2, "number of mismatched pixels without diff"
+
+
+@pytest.mark.parametrize(
+    "img_a_is_PIL, img_b_is_PIL, output_is_PIL, specify_size",
+    itertools.product([True, False], repeat=4),
+)
+def test_works_with_PIL_Image(img_a_is_PIL, img_b_is_PIL, output_is_PIL, specify_size):
+    img_a_path, img_b_path, diff1_path, options, expected_mismatch = testdata[0]
+
+    img_a_data = read_img(img_a_path)
+    img_sizes = img_a_data.size
+    if not img_a_is_PIL:
+        img_a_data = pil_to_flatten_data(img_a_data)
+
+    img_b_data = read_img(img_b_path)
+    if not img_b_is_PIL:
+        img_b_data = pil_to_flatten_data(img_b_data)
+
+    diff_data = Image.new("RGBA", img_sizes)
+    if not output_is_PIL:
+        diff_data = pil_to_flatten_data(diff_data)
+
+    if specify_size or (not img_a_is_PIL and not img_b_is_PIL):
+        options["width"], options["height"] = img_sizes
+
+    mismatch = pixelmatch(img_a_data, img_b_data, output=diff_data, **options)
+    mismatch2 = pixelmatch(img_a_data, img_b_data, **options)
+
+    expected_diff = pil_to_flatten_data(read_img(diff1_path))
+
+    if output_is_PIL:
+        diff_data = pil_to_flatten_data(diff_data)
+    assert diff_data == expected_diff, "diff image"
+
+    assert mismatch == expected_mismatch, "number of mismatched pixels"
+    assert mismatch == mismatch2, "number of mismatched pixels without diff"