Applying Atkinson dithering with Red, Yellow, Black and White, mostly for EZSign.

atkinson_dither.py

#!/usr/bin/env python3
"""
Apply Atkinson dithering to an input PNG using a fixed 4-color palette:
red, yellow, black, and white.
"""

from __future__ import annotations

import argparse
from pathlib import Path

from PIL import Image


PALETTE = [
    (255, 0, 0),      # red
    (255, 255, 0),    # yellow
    (0, 0, 0),        # black
    (255, 255, 255),  # white
]


def closest_palette_color(r: float, g: float, b: float) -> tuple[int, int, int]:
    best = PALETTE[0]
    best_dist = float("inf")
    for pr, pg, pb in PALETTE:
        dr = r - pr
        dg = g - pg
        db = b - pb
        dist = dr * dr + dg * dg + db * db
        if dist < best_dist:
            best_dist = dist
            best = (pr, pg, pb)
    return best


def atkinson_dither(image: Image.Image) -> Image.Image:
    rgb = image.convert("RGB")
    w, h = rgb.size
    src = [[[float(c) for c in rgb.getpixel((x, y))] for x in range(w)] for y in range(h)]
    out = Image.new("RGB", (w, h))

    neighbors = [(1, 0), (2, 0), (-1, 1), (0, 1), (1, 1), (0, 2)]

    for y in range(h):
        for x in range(w):
            old_r, old_g, old_b = src[y][x]
            new_r, new_g, new_b = closest_palette_color(old_r, old_g, old_b)
            out.putpixel((x, y), (new_r, new_g, new_b))

            err_r = old_r - new_r
            err_g = old_g - new_g
            err_b = old_b - new_b

            for dx, dy in neighbors:
                nx = x + dx
                ny = y + dy
                if 0 <= nx < w and 0 <= ny < h:
                    src[ny][nx][0] += err_r / 8.0
                    src[ny][nx][1] += err_g / 8.0
                    src[ny][nx][2] += err_b / 8.0

    return out


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Reduce a PNG to red/yellow/black/white using Atkinson dithering."
    )
    parser.add_argument("input_png", type=Path, help="Path to input PNG file.")
    parser.add_argument("output_png", type=Path, help="Path to output PNG file.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    with Image.open(args.input_png) as img:
        result = atkinson_dither(img)

    args.output_png.parent.mkdir(parents=True, exist_ok=True)
    result.save(args.output_png, format="PNG")


if __name__ == "__main__":
    main()