Account for harmless pixel shifts in image checks

conda/dev.yml

@@ -18,6 +18,7 @@ dependencies:
   - jinja2>=3.1.2     # Changed from =3.1.2
   - mache>=1.5.0
   - mpas_tools>=0.15.0
+  - numpy
   - pillow>=9.2.0     # Changed from =9.2.0
   # Testing
   # =======================

tests/images/test_image_checker.py

@@ -1,15 +1,24 @@
 import os
+from pathlib import Path
 from typing import List, Optional
 
 from mache import MachineInfo
+from PIL import Image, ImageDraw
 
 from tests.integration.image_checker import _compare_actual_and_expected
 
 
+def _write_image(path: str) -> None:
+    image = Image.new("RGB", (100, 100), "white")
+    draw = ImageDraw.Draw(image)
+    draw.rectangle((20, 20, 80, 80), fill="black")
+    image.save(path)
+
+
 # Run this test with:
 # cd zppy
 # pytest tests/images/test_image_checker.py
-def test_compare():
+def test_compare() -> None:
     missing_images: List[str] = []
     mismatched_images: List[str] = []
 
@@ -42,3 +51,29 @@ def test_compare():
     )
     assert missing_images == []
     assert mismatched_images == ["CRU-TREFHT-ANN-land_60S90N"]
+
+
+def test_compare_ignores_small_pixel_shift(tmp_path: Path) -> None:
+    expected_path = tmp_path / "expected.png"
+    actual_path = tmp_path / "actual.png"
+    _write_image(str(expected_path))
+
+    with Image.open(expected_path) as expected_image:
+        expected_image = expected_image.convert("RGB")
+        actual_image = Image.new("RGB", expected_image.size, "white")
+        actual_image.paste(expected_image, (1, -2))
+        actual_image.save(actual_path)
+
+    missing_images: List[str] = []
+    mismatched_images: List[str] = []
+    _compare_actual_and_expected(
+        missing_images,
+        mismatched_images,
+        "shifted.png",
+        str(actual_path),
+        str(expected_path),
+        str(tmp_path / "diffs"),
+    )
+
+    assert missing_images == []
+    assert mismatched_images == []

tests/integration/image_checker.py

@@ -1,14 +1,21 @@
 import os
 import shutil
 from math import ceil
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import matplotlib.backends.backend_pdf
 import matplotlib.image as mpimg
+import numpy as np
 from mache import MachineInfo
 from matplotlib import pyplot as plt
 from PIL import Image, ImageChops, ImageDraw
 
+# The FFT-based shift estimate in `_estimate_shift` can in principle return
+# any shift up to half the image's width/height, but shifts larger than this
+# are treated as genuine mismatches rather than harmless translations.
+MAXIMUM_PIXEL_SHIFT = 10
+MAXIMUM_MISMATCH_FRACTION = 0.0002
+
 
 # Classes #####################################################################
 class Parameters(object):
@@ -45,6 +52,22 @@ def __init__(
 # Specialized setup ###########################################################
 
 
+def _get_unused_diff_dir(base_diff_dir: str) -> str:
+    """Return a diff_dir path that does not already exist.
+
+    If base_diff_dir does not exist, it is returned unchanged. Otherwise,
+    a "_tryN" suffix (starting at N=2) is appended until an unused
+    directory name is found. This prevents reruns of the image checker
+    from overwriting the results of previous runs.
+    """
+    if not os.path.exists(base_diff_dir):
+        return base_diff_dir
+    try_number = 2
+    while os.path.exists(f"{base_diff_dir}_try{try_number}"):
+        try_number += 1
+    return f"{base_diff_dir}_try{try_number}"
+
+
 def set_up_and_run_image_checker(
     cfg_specifier: str,
     case_name: str,
@@ -55,22 +78,19 @@ def set_up_and_run_image_checker(
 ):
     print(f"Image checking {cfg_specifier}")
     actual_images_dir = f"{expansions['user_www']}zppy_weekly_{cfg_specifier}_www/{expansions['unique_id']}/{case_name}/"
+    base_diff_dir = (
+        f"{actual_images_dir}image_check_failures_{cfg_specifier}{diff_dir_suffix}"
+    )
     d: Dict[str, str] = {
         "actual_images_dir": actual_images_dir,
         "expected_images_dir": f"{expansions['expected_dir']}expected_{cfg_specifier}",
-        "diff_dir": f"{actual_images_dir}image_check_failures_{cfg_specifier}{diff_dir_suffix}",
+        "diff_dir": _get_unused_diff_dir(base_diff_dir),
         "expected_images_list": f"{expansions['expected_dir']}image_list_expected_{cfg_specifier}.txt",
     }
-    print(f"Removing diff_dir={d['diff_dir']} to produce new results")
-    if os.path.exists(d["diff_dir"]):
-        try:
-            shutil.rmtree(d["diff_dir"])
-        except PermissionError:
-            print(
-                f"{d['diff_dir']} cannot be removed. Execute permissions are needed to remove files. Adding execute permission and trying again."
-            )
-            _chmod_recursive(d["diff_dir"], 0o744)
-            shutil.rmtree(d["diff_dir"])
+    if d["diff_dir"] != base_diff_dir:
+        print(
+            f"diff_dir={base_diff_dir} already exists; using {d['diff_dir']} instead to avoid overwriting previous results"
+        )
     print("Image checking dict:")
     for key in d:
         print(f"{key}: {d[key]}")
@@ -271,27 +291,15 @@ def _compare_actual_and_expected(
         # If `diff.getbbox()` is None, then the images are in theory equal
         assert diff.getbbox() is None
     else:
-        # Sometimes, a few pixels will differ, but the two images appear identical.
-        # https://codereview.stackexchange.com/questions/55902/fastest-way-to-count-non-zero-pixels-using-python-and-pillow
-        nonzero_pixels = (
-            diff.crop(bbox)
-            .point(lambda x: 255 if x else 0)
-            .convert("L")
-            .point(bool)
-            .getdata()
-        )
-        num_nonzero_pixels = sum(nonzero_pixels)
-        width, height = expected_png.size
-        num_pixels = width * height
-        fraction = num_nonzero_pixels / num_pixels
+        fraction = _get_mismatched_fraction(diff, expected_png.size)
         # Fraction of mismatched pixels should be less than 0.02%
-        if fraction >= 0.0002:
+        if fraction >= MAXIMUM_MISMATCH_FRACTION and not _images_match_after_shift(
+            actual_png, expected_png
+        ):
             verbose = False
             if verbose:
                 print("\npath_to_actual_png={}".format(path_to_actual_png))
                 print("path_to_expected_png={}".format(path_to_expected_png))
-                print("diff has {} nonzero pixels.".format(num_nonzero_pixels))
-                print("total number of pixels={}".format(num_pixels))
                 print("num_nonzero_pixels/num_pixels fraction={}".format(fraction))
 
             mismatched_images.append(image_name)
@@ -324,6 +332,120 @@ def _compare_actual_and_expected(
             )
 
 
+def _get_mismatched_fraction(diff: Image.Image, size: Tuple[int, int]) -> float:
+    # Sometimes, a few pixels will differ, but the two images appear identical.
+    # https://codereview.stackexchange.com/questions/55902/fastest-way-to-count-non-zero-pixels-using-python-and-pillow
+    bbox = diff.getbbox()
+    if bbox is None:
+        return 0.0
+    nonzero_pixels = (
+        diff.crop(bbox)
+        .point(lambda x: 255 if x else 0)
+        .convert("L")
+        .point(bool)
+        .getdata()
+    )
+    return sum(nonzero_pixels) / (size[0] * size[1])
+
+
+def _images_match_after_shift(
+    actual_png: Image.Image, expected_png: Image.Image
+) -> bool:
+    if actual_png.size != expected_png.size:
+        return False
+
+    # Rather than exhaustively trying every candidate shift (which costs up
+    # to (2 * MAXIMUM_PIXEL_SHIFT + 1)^2 image diffs), use phase correlation
+    # (an FFT-based technique) to find the single best-aligning shift in one
+    # pass. See: https://en.wikipedia.org/wiki/Phase_correlation
+    horizontal_shift, vertical_shift = _estimate_shift(actual_png, expected_png)
+    if horizontal_shift == 0 and vertical_shift == 0:
+        return False
+    if (
+        abs(horizontal_shift) > MAXIMUM_PIXEL_SHIFT
+        or abs(vertical_shift) > MAXIMUM_PIXEL_SHIFT
+    ):
+        return False
+
+    actual_overlap, expected_overlap = _get_overlapping_images(
+        actual_png, expected_png, horizontal_shift, vertical_shift
+    )
+    shifted_diff = ImageChops.difference(actual_overlap, expected_overlap)
+    return (
+        _get_mismatched_fraction(shifted_diff, shifted_diff.size)
+        < MAXIMUM_MISMATCH_FRACTION
+    )
+
+
+def _estimate_shift(
+    actual_png: Image.Image, expected_png: Image.Image
+) -> Tuple[int, int]:
+    # Use phase correlation to estimate the (horizontal, vertical) pixel
+    # shift that best aligns `actual_png` with `expected_png`. This finds the
+    # shift in a single FFT-based operation, regardless of how large the
+    # shift is, rather than searching over every candidate shift.
+    actual_array = np.asarray(actual_png.convert("L"), dtype=np.float64)
+    expected_array = np.asarray(expected_png.convert("L"), dtype=np.float64)
+
+    actual_fft = np.fft.fft2(actual_array)
+    expected_fft = np.fft.fft2(expected_array)
+    # The order here matters: this yields a peak at (horizontal_shift,
+    # vertical_shift) such that actual(x, y) == expected(x + horizontal_shift,
+    # y + vertical_shift), matching the convention used by
+    # `_get_overlapping_images`.
+    cross_power = expected_fft * np.conj(actual_fft)
+    magnitude = np.abs(cross_power)
+    # Avoid division by zero where the cross power is (near) zero.
+    magnitude[magnitude < 1e-10] = 1e-10
+    correlation = np.abs(np.fft.ifft2(cross_power / magnitude))
+
+    height, width = correlation.shape
+    peak_row, peak_col = np.unravel_index(np.argmax(correlation), correlation.shape)
+
+    # The correlation surface wraps around at the image boundaries, so a peak
+    # in the second half of the axis corresponds to a negative shift.
+    horizontal_shift = int(peak_col)
+    if horizontal_shift > width // 2:
+        horizontal_shift -= width
+    vertical_shift = int(peak_row)
+    if vertical_shift > height // 2:
+        vertical_shift -= height
+
+    return horizontal_shift, vertical_shift
+
+
+def _get_overlapping_images(
+    actual_png: Image.Image,
+    expected_png: Image.Image,
+    horizontal_shift: int,
+    vertical_shift: int,
+) -> Tuple[Image.Image, Image.Image]:
+    width, height = actual_png.size
+    actual_left = max(0, -horizontal_shift)
+    actual_upper = max(0, -vertical_shift)
+    expected_left = max(0, horizontal_shift)
+    expected_upper = max(0, vertical_shift)
+    overlap_width = width - abs(horizontal_shift)
+    overlap_height = height - abs(vertical_shift)
+    actual_overlap = actual_png.crop(
+        (
+            actual_left,
+            actual_upper,
+            actual_left + overlap_width,
+            actual_upper + overlap_height,
+        )
+    )
+    expected_overlap = expected_png.crop(
+        (
+            expected_left,
+            expected_upper,
+            expected_left + overlap_width,
+            expected_upper + overlap_height,
+        )
+    )
+    return actual_overlap, expected_overlap
+
+
 def _draw_box(image, diff, output_path: str):
     # https://stackoverflow.com/questions/41405632/draw-a-rectangle-and-a-text-in-it-using-pil
     draw = ImageDraw.Draw(image)