Screenshot Differ for Visual Regression Testing

Screenshot Differ for Visual Regression Testing

Problem Statement

Front-end applications break in subtle ways that unit tests never catch: a CSS refactor shifts a button 4 pixels, a font swap changes text rendering, an image asset fails to load, or a responsive breakpoint collapses a layout. These are "visual regressions" — the code still runs, tests still pass, but the UI looks wrong. Manually eyeballing screenshots across every pull request is tedious and error-prone, and human reviewers miss small changes after looking at dozens of images.

This script automates that comparison. Given a directory of baseline screenshots and a directory of freshly captured screenshots (typically produced by a headless browser like Playwright or Puppeteer in CI), it computes a pixel-level diff for each matching pair, quantifies the percentage of changed pixels, and produces a report with per-image verdicts plus annotated diff images highlighting exactly where the UI drifted. It exits with a non-zero status when regressions exceed the allowed threshold, making it a drop-in gate for any CI pipeline.

How It Works

The script follows a directory-pairing → per-image diff → aggregate report pipeline. First, it discovers images in both directories (recursively, matching by relative path) and computes the intersection of filenames. For each matched pair, it opens both images with Pillow, converts them to RGB, and — critically — normalizes their sizes. If dimensions differ, it pads the smaller image with a configurable background color so the diff is meaningful rather than aborting. The two images are then cast to NumPy arrays, and a per-pixel absolute difference is computed across the RGB channels.

A pixel is counted as "changed" when its maximum channel delta exceeds a per-pixel sensitivity threshold (to ignore anti-aliasing and JPEG noise). The fraction of changed pixels is compared against a global --threshold. If it exceeds the threshold, the pair is flagged as a failure and a diff visualization is written: unchanged regions are dimmed to grayscale while changed pixels are painted in a bright highlight color. The script uses a small dataclass to carry per-image results and returns a proper exit code (0 = pass, 1 = regression detected, 2 = runtime error) so CI systems can react correctly.

Full Script

#!/usr/bin/env python3
"""
Screenshot Differ for Visual Regression Testing.

Compares two directories of screenshots pixel-by-pixel, reports the fraction
of changed pixels per image, and writes annotated diff images. Exits non-zero
when any image exceeds the allowed change threshold.

Usage:
    python3 screenshot_diff.py --base baseline/ --current current/ \
        --threshold 0.02 --out report/
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable

import numpy as np
from PIL import Image, UnidentifiedImageError

# ---------------------------------------------------------------------------
# Configuration (overridable via environment variables)
# ---------------------------------------------------------------------------
DEFAULT_THRESHOLD = float(os.environ.get("SDIFF_THRESHOLD", "0.02"))
# Per-channel delta (0-255) below which a pixel is considered "unchanged".
# Filters out anti-aliasing and lossy-compression noise.
PIXEL_SENSITIVITY = int(os.environ.get("SDIFF_PIXEL_SENSITIVITY", "10"))
# Max width/height difference in pixels tolerated before we treat a size
# mismatch as a hard failure instead of padding and diffing.
MAX_SIZE_DRIFT = int(os.environ.get("SDIFF_MAX_SIZE_DRIFT", "200"))
# Background color used when padding a smaller image to match the larger one.
PAD_COLOR = tuple(
    int(c) for c in os.environ.get("SDIFF_PAD_COLOR", "255,255,255").split(",")
)
SUPPORTED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tiff"}

logging.basicConfig(
    level=os.environ.get("SDIFF_LOG_LEVEL", "INFO").upper(),
    format="%(levelname)s: %(message)s",
)
log = logging.getLogger("screenshot_diff")


# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------
@dataclass
class DiffResult:
    """Outcome of comparing one baseline/current image pair."""

    name: str
    changed_fraction: float
    width: int
    height: int
    passed: bool
    diff_path: Path | None = None
    error: str | None = None
    size_mismatch: bool = False

    def as_row(self) -> str:
        status = "PASS" if self.passed else "FAIL"
        if self.error:
            status = "ERR "
        pct = f"{self.changed_fraction * 100:6.2f}%"
        return f"  [{status}] {pct}  {self.name}"


@dataclass
class Report:
    """Aggregate results for the whole run."""

    results: list[DiffResult] = field(default_factory=list)

    @property
    def failures(self) -> list[DiffResult]:
        return [r for r in self.results if not r.passed]

    @property
    def errors(self) -> list[DiffResult]:
        return [r for r in self.results if r.error]

    def summary(self) -> str:
        total = len(self.results)
        failed = len(self.failures)
        errored = len(self.errors)
        return (
            f"\n{'-' * 60}\n"
            f"Total: {total}  Passed: {total - failed - errored}  "
            f"Failed: {failed}  Errors: {errored}\n"
            f"{'-' * 60}"
        )


# ---------------------------------------------------------------------------
# Image discovery & loading
# ---------------------------------------------------------------------------
def discover_images(root: Path) -> dict[str, Path]:
    """Map relative POSIX paths -> absolute Path for all supported images."""
    if not root.is_dir():
        raise FileNotFoundError(f"Directory not found: {root}")
    found: dict[str, Path] = {}
    for path in sorted(root.rglob("*")):
        if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS:
            rel = path.relative_to(root).as_posix()
            found[rel] = path
    return found


def load_rgb(path: Path) -> np.ndarray:
    """Load an image as a contiguous RGB uint8 NumPy array."""
    try:
        with Image.open(path) as img:
            return np.asarray(img.convert("RGB"), dtype=np.uint8)
    except UnidentifiedImageError as exc:
        raise ValueError(f"Not a valid image: {path}") from exc


def pad_to_match(
    a: np.ndarray, b: np.ndarray, color: tuple[int, int, int]
) -> tuple[np.ndarray, np.ndarray]:
    """Pad the smaller array with `color` so both share the same shape."""
    h = max(a.shape[0], b.shape[0])
    w = max(a.shape[1], b.shape[1])

    def _pad(arr: np.ndarray) -> np.ndarray:
        canvas = np.full((h, w, 3), color, dtype=np.uint8)
        canvas[: arr.shape[0], : arr.shape[1]] = arr
        return canvas

    return _pad(a), _pad(b)


# ---------------------------------------------------------------------------
# Core diff logic
# ---------------------------------------------------------------------------
def compute_diff(
    base_arr: np.ndarray,
    curr_arr: np.ndarray,
    sensitivity: int,
) -> tuple[float, np.ndarray]:
    """
    Return (changed_fraction, changed_mask).

    changed_mask is a 2D boolean array, True where the pixel differs by more
    than `sensitivity` on any RGB channel.
    """
    if base_arr.shape != curr_arr.shape:
        raise ValueError("Arrays must be padded to identical shapes first.")
    # Absolute per-channel difference; collapse channels via max.
    delta = np.abs(
        base_arr.astype(np.int16) - curr_arr.astype(np.int16)
    )
    changed_mask = delta.max(axis=2) > sensitivity
    changed_fraction = float(changed_mask.mean())
    return changed_fraction, changed_mask


def render_diff_image(
    base_arr: np.ndarray,
    curr_arr: np.ndarray,
    mask: np.ndarray,
    highlight: tuple[int, int, int] = (255, 0, 0),
) -> Image.Image:
    """Grayscale the current image and paint changed pixels in `highlight`."""
    gray = (curr_arr @ np.array([0.299, 0.587, 0.114])).astype(np.uint8)
    out = np.stack([gray, gray, gray], axis=2)
    out[mask] = np.array(highlight, dtype=np.uint8)
    return Image.fromarray(out, mode="RGB")


def compare_pair(
    name: str,
    base_path: Path,
    curr_path: Path,
    threshold: float,
    out_dir: Path,
) -> DiffResult:
    """Compare a single image pair and persist a diff image if it changed."""
    try:
        base_arr = load_rgb(base_path)
        curr_arr = load_rgb(curr_path)
    except (ValueError, OSError) as exc:
        log.error("Failed to load %s: %s", name, exc)
        return DiffResult(name, 0.0, 0, 0, passed=False, error=str(exc))

    size_mismatch = base_arr.shape != curr_arr.shape
    if size_mismatch:
        drift = max(
            abs(base_arr.shape[0] - curr_arr.shape[0]),
            abs(base_arr.shape[1] - curr_arr.shape[1]),
        )
        if drift > MAX_SIZE_DRIFT:
            msg = f"size drift {drift}px exceeds limit {MAX_SIZE_DRIFT}px"
            log.error("%s: %s", name, msg)
            return DiffResult(
                name, 1.0, 0, 0, passed=False, error=msg, size_mismatch=True
            )
        log.warning("%s: size mismatch, padding to compare (drift=%dpx)", name, drift)
        base_arr, curr_arr = pad_to_match(base_arr, curr_arr, PAD_COLOR)

    changed_fraction, mask = compute_diff(base_arr, curr_arr, PIXEL_SENSITIVITY)
    passed = changed_fraction <= threshold

    diff_path: Path | None = None
    if not passed or changed_fraction > 0:
        diff_path = out_dir / f"{Path(name).stem}.diff.png"
        diff_path.parent.mkdir(parents=True, exist_ok=True)
        render_diff_image(base_arr, curr_arr, mask).save(diff_path)

    h, w = base_arr.shape[:2]
    return DiffResult(
        name=name,
        changed_fraction=changed_fraction,
        width=w,
        height=h,
        passed=passed,
        diff_path=diff_path,
        size_mismatch=size_mismatch,
    )


# ---------------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------------
def run(
    base_dir: Path,
    curr_dir: Path,
    out_dir: Path,
    threshold: float,
) -> Report:
    """Compare every matching image pair across the two directories."""
    base_images = discover_images(base_dir)
    curr_images = discover_images(curr_dir)

    if not base_images:
        raise FileNotFoundError(f"No supported images found in {base_dir}")

    out_dir.mkdir(parents=True, exist_ok=True)
    report = Report()

    only_base = set(base_images) - set(curr_images)
    only_curr = set(curr_images) - set(base_images)
    for missing in sorted(only_base):
        log.warning("Missing in current: %s", missing)
    for extra in sorted(only_curr):
        log.warning("Extra in current (no baseline): %s", extra)

    shared = sorted(set(base_images) & set(curr_images))
    if not shared:
        raise RuntimeError("No matching filenames between base and current dirs.")

    log.info("Comparing %d image pair(s)...", len(shared))
    for name in shared:
        result = compare_pair(
            name, base_images[name], curr_images[name], threshold, out_dir
        )
        report.results.append(result)
        log.info(result.as_row())

    return report


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def parse_args(argv: Iterable[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Pixel-level screenshot diff for visual regression testing."
    )
    parser.add_argument("--base", required=True, type=Path, help="Baseline dir")
    parser.add_argument("--current", required=True, type=Path, help="New screenshots dir")
    parser.add_argument(
        "--threshold",
        type=float,
        default=DEFAULT_THRESHOLD,
        help=f"Max changed fraction before failing (default {DEFAULT_THRESHOLD})",
    )
    parser.add_argument("--out", required=True, type=Path, help="Report output dir")
    parser.add_argument(
        "--fail-on-missing",
        action="store_true",
        help="Fail if an image exists in only one directory.",
    )
    return parser.parse_args(argv)


def main(argv: Iterable[str] | None = None) -> int:
    args = parse_args(argv)
    if not 0.0 <= args.threshold <= 1.0:
        log.error("--threshold must be between 0.0 and 1.0")
        return 2

    try:
        report = run(args.base, args.current, args.out, args.threshold)
    except (FileNotFoundError, RuntimeError) as exc:
        log.error("%s", exc)
        return 2

    print(report.summary())
    if report.errors:
        log.error("Encountered %d image error(s).", len(report.errors))
        return 2
    if report.failures:
        log.error(
            "%d image(s) exceeded threshold %.4f. See %s",
            len(report.failures),
            args.threshold,
            args.out,
        )
        return 1
    log.info("All screenshots within threshold. Visual regression check passed.")
    return 0


if __name__ == "__main__":
    sys.exit(main())

Configuration

CLI flags: --base, --current, --out (all required), --threshold, and --fail-on-missing (treat files present in only one directory as a failure).

Usage Examples

1. Standard CI run — fail if more than 2% of pixels changed:

$ python3 screenshot_diff.py --base baseline/ --current current/ --threshold 0.02 --out report/
INFO: Comparing 3 image pair(s)...
INFO:   [PASS]   0.00%  home.png
INFO:   [FAIL]   4.31%  login.png
INFO:   [PASS]   0.12%  dashboard.png
ERROR: 1 image(s) exceeded threshold 0.0200. See report/
$ echo $?
1

2. Strict comparison with a tighter sensitivity for pixel-perfect matching:

$ SDIFF_PIXEL_SENSITIVITY=2 python3 screenshot_diff.py \
    --base baseline/ --current current/ --threshold 0.001 --out report/
INFO:   [FAIL]   0.05%  checkout.png
ERROR: 1 image(s) exceeded threshold 0.0010. See report/

3. Loose smoke test that tolerates layout jitter and missing assets:

$ SDIFF_THRESHOLD=0.15 SDIFF_MAX_SIZE_DRIFT=500 python3 screenshot_diff.py \
    --base baseline/ --current current/ --threshold 0.15 --out report/
WARNING: home.png: size mismatch, padding to compare (drift=12px)
INFO:   [PASS]   0.00%  home.png
INFO: All screenshots within threshold. Visual regression check passed.
$ echo $?
0

Error Handling

• Missing directories → discover_images raises FileNotFoundError, caught in main, logged, and returns exit code 2.
• No overlapping filenames → RuntimeError is raised (nothing to compare), exit code 2.
• Corrupt or non-image files → load_rgb raises ValueError; the pair is recorded with an error field, logged, and the run returns 2 so CI never silently passes.
• Dimension mismatches → padded up to SDIFF_MAX_SIZE_DRIFT; beyond that, a hard error is recorded rather than producing a meaningless diff.
• Files present in only one directory → logged as warnings; use --fail-on-missing to escalate them to failures.
• Invalid --threshold → validated against [0.0, 1.0], returns 2 before doing any work.

Exit codes are distinct on purpose: 0 = pass, 1 = visual regression, 2 = infrastructure/configuration error.

Extending Ideas

• Per-image thresholds: add a thresholds.json mapping filenames to custom tolerances (e.g., allow 5% on animated charts, 0.1% on static pages).
• Region masking: accept a --ignore-regions file of bounding boxes (timestamps, ads, live data) that are excluded from the diff mask before computing the changed fraction.
• HTML report: emit a self-contained report/index.html with side-by-side base/current/diff thumbnails for fast human review in CI artifacts.
• Baseline promotion: add a --update-baseline mode that copies passing current/ images into baseline/ so approved changes become the new reference automatically.
• SSIM instead of raw pixels: swap compute_diff for a structural-similarity metric (via scikit-image) to better match human perception of "looks the same."