""" Dithering algorithms for 6-color e-ink display testing. Includes multiple error-diffusion and ordered dithering algorithms for comparison testing. """ import numpy as np from PIL import Image from typing import Tuple, List from numba import jit # 6-color ACeP palette (RGB format) PALETTE_RGB = [ (0, 0, 0), # BLACK (255, 255, 255), # WHITE (255, 255, 0), # YELLOW (255, 0, 0), # RED (0, 0, 255), # BLUE (0, 255, 0), # GREEN ] PALETTE_NAMES = ['Black', 'White', 'Yellow', 'Red', 'Blue', 'Green'] def create_pil_palette_image() -> Image.Image: """Create a PIL palette image for quantization.""" pal_image = Image.new("P", (1, 1)) flat_palette = [] for color in PALETTE_RGB: flat_palette.extend(color) # Pad to 256 colors (768 values) flat_palette.extend([0] * (768 - len(flat_palette))) pal_image.putpalette(flat_palette) return pal_image def find_nearest_color(pixel: np.ndarray, palette: np.ndarray) -> Tuple[int, np.ndarray]: """Find the nearest palette color using Euclidean distance.""" distances = np.sqrt(np.sum((palette - pixel) ** 2, axis=1)) idx = np.argmin(distances) return idx, palette[idx] def find_nearest_color_weighted(pixel: np.ndarray, palette: np.ndarray) -> Tuple[int, np.ndarray]: """Find nearest color using perceptually-weighted distance (human eye sensitivity).""" # Weights based on human perception: Green > Red > Blue weights = np.array([0.299, 0.587, 0.114]) diff = palette - pixel weighted_diff = diff * weights distances = np.sqrt(np.sum(weighted_diff ** 2, axis=1)) idx = np.argmin(distances) return idx, palette[idx] # ============================================================================= # Error Diffusion Dithering Algorithms # ============================================================================= def dither_floyd_steinberg(image: Image.Image, weighted: bool = False) -> Image.Image: """ Floyd-Steinberg dithering (1976). Classic error diffusion algorithm with distribution: * 7/16 3/16 5/16 1/16 """ img = np.array(image.convert('RGB'), dtype=np.float64) height, width = img.shape[:2] palette = np.array(PALETTE_RGB, dtype=np.float64) find_fn = find_nearest_color_weighted if weighted else find_nearest_color for y in range(height): for x in range(width): old_pixel = img[y, x].copy() _, new_pixel = find_fn(old_pixel, palette) img[y, x] = new_pixel error = old_pixel - new_pixel if x + 1 < width: img[y, x + 1] += error * 7 / 16 if y + 1 < height: if x > 0: img[y + 1, x - 1] += error * 3 / 16 img[y + 1, x] += error * 5 / 16 if x + 1 < width: img[y + 1, x + 1] += error * 1 / 16 img = np.clip(img, 0, 255).astype(np.uint8) return Image.fromarray(img, 'RGB') def dither_jarvis_judice_ninke(image: Image.Image, weighted: bool = False) -> Image.Image: """ Jarvis, Judice, and Ninke dithering (1976). Spreads error over a larger area (12 pixels): * 7 5 3 5 7 5 3 1 3 5 3 1 All divided by 48. """ img = np.array(image.convert('RGB'), dtype=np.float64) height, width = img.shape[:2] palette = np.array(PALETTE_RGB, dtype=np.float64) find_fn = find_nearest_color_weighted if weighted else find_nearest_color for y in range(height): for x in range(width): old_pixel = img[y, x].copy() _, new_pixel = find_fn(old_pixel, palette) img[y, x] = new_pixel error = old_pixel - new_pixel # Row 0 (current row) if x + 1 < width: img[y, x + 1] += error * 7 / 48 if x + 2 < width: img[y, x + 2] += error * 5 / 48 # Row 1 if y + 1 < height: for dx, w in [(-2, 3), (-1, 5), (0, 7), (1, 5), (2, 3)]: if 0 <= x + dx < width: img[y + 1, x + dx] += error * w / 48 # Row 2 if y + 2 < height: for dx, w in [(-2, 1), (-1, 3), (0, 5), (1, 3), (2, 1)]: if 0 <= x + dx < width: img[y + 2, x + dx] += error * w / 48 img = np.clip(img, 0, 255).astype(np.uint8) return Image.fromarray(img, 'RGB') def dither_stucki(image: Image.Image, weighted: bool = False) -> Image.Image: """ Stucki dithering (1981). Similar to JJN but with different weights: * 8 4 2 4 8 4 2 1 2 4 2 1 All divided by 42. """ img = np.array(image.convert('RGB'), dtype=np.float64) height, width = img.shape[:2] palette = np.array(PALETTE_RGB, dtype=np.float64) find_fn = find_nearest_color_weighted if weighted else find_nearest_color for y in range(height): for x in range(width): old_pixel = img[y, x].copy() _, new_pixel = find_fn(old_pixel, palette) img[y, x] = new_pixel error = old_pixel - new_pixel if x + 1 < width: img[y, x + 1] += error * 8 / 42 if x + 2 < width: img[y, x + 2] += error * 4 / 42 if y + 1 < height: for dx, w in [(-2, 2), (-1, 4), (0, 8), (1, 4), (2, 2)]: if 0 <= x + dx < width: img[y + 1, x + dx] += error * w / 42 if y + 2 < height: for dx, w in [(-2, 1), (-1, 2), (0, 4), (1, 2), (2, 1)]: if 0 <= x + dx < width: img[y + 2, x + dx] += error * w / 42 img = np.clip(img, 0, 255).astype(np.uint8) return Image.fromarray(img, 'RGB') def dither_atkinson(image: Image.Image, weighted: bool = False) -> Image.Image: """ Atkinson dithering (Bill Atkinson, Apple). Only diffuses 6/8 of the error (loses some detail but reduces noise): * 1 1 1 1 1 1 All divided by 8 (but only 6/8 total error diffused). """ img = np.array(image.convert('RGB'), dtype=np.float64) height, width = img.shape[:2] palette = np.array(PALETTE_RGB, dtype=np.float64) find_fn = find_nearest_color_weighted if weighted else find_nearest_color for y in range(height): for x in range(width): old_pixel = img[y, x].copy() _, new_pixel = find_fn(old_pixel, palette) img[y, x] = new_pixel error = old_pixel - new_pixel # Distribute 1/8 to each of 6 neighbors if x + 1 < width: img[y, x + 1] += error / 8 if x + 2 < width: img[y, x + 2] += error / 8 if y + 1 < height: if x > 0: img[y + 1, x - 1] += error / 8 img[y + 1, x] += error / 8 if x + 1 < width: img[y + 1, x + 1] += error / 8 if y + 2 < height: img[y + 2, x] += error / 8 img = np.clip(img, 0, 255).astype(np.uint8) return Image.fromarray(img, 'RGB') def dither_sierra(image: Image.Image, weighted: bool = False) -> Image.Image: """ Sierra dithering (Frankie Sierra). Full Sierra (Sierra-3): * 5 3 2 4 5 4 2 2 3 2 All divided by 32. """ img = np.array(image.convert('RGB'), dtype=np.float64) height, width = img.shape[:2] palette = np.array(PALETTE_RGB, dtype=np.float64) find_fn = find_nearest_color_weighted if weighted else find_nearest_color for y in range(height): for x in range(width): old_pixel = img[y, x].copy() _, new_pixel = find_fn(old_pixel, palette) img[y, x] = new_pixel error = old_pixel - new_pixel if x + 1 < width: img[y, x + 1] += error * 5 / 32 if x + 2 < width: img[y, x + 2] += error * 3 / 32 if y + 1 < height: for dx, w in [(-2, 2), (-1, 4), (0, 5), (1, 4), (2, 2)]: if 0 <= x + dx < width: img[y + 1, x + dx] += error * w / 32 if y + 2 < height: for dx, w in [(-1, 2), (0, 3), (1, 2)]: if 0 <= x + dx < width: img[y + 2, x + dx] += error * w / 32 img = np.clip(img, 0, 255).astype(np.uint8) return Image.fromarray(img, 'RGB') def dither_sierra_lite(image: Image.Image, weighted: bool = False) -> Image.Image: """ Sierra Lite (Sierra-2-4A) - faster variant. * 2 1 1 All divided by 4. """ img = np.array(image.convert('RGB'), dtype=np.float64) height, width = img.shape[:2] palette = np.array(PALETTE_RGB, dtype=np.float64) find_fn = find_nearest_color_weighted if weighted else find_nearest_color for y in range(height): for x in range(width): old_pixel = img[y, x].copy() _, new_pixel = find_fn(old_pixel, palette) img[y, x] = new_pixel error = old_pixel - new_pixel if x + 1 < width: img[y, x + 1] += error * 2 / 4 if y + 1 < height: if x > 0: img[y + 1, x - 1] += error * 1 / 4 img[y + 1, x] += error * 1 / 4 img = np.clip(img, 0, 255).astype(np.uint8) return Image.fromarray(img, 'RGB') def dither_burkes(image: Image.Image, weighted: bool = False) -> Image.Image: """ Burkes dithering. * 8 4 2 4 8 4 2 All divided by 32. """ img = np.array(image.convert('RGB'), dtype=np.float64) height, width = img.shape[:2] palette = np.array(PALETTE_RGB, dtype=np.float64) find_fn = find_nearest_color_weighted if weighted else find_nearest_color for y in range(height): for x in range(width): old_pixel = img[y, x].copy() _, new_pixel = find_fn(old_pixel, palette) img[y, x] = new_pixel error = old_pixel - new_pixel if x + 1 < width: img[y, x + 1] += error * 8 / 32 if x + 2 < width: img[y, x + 2] += error * 4 / 32 if y + 1 < height: for dx, w in [(-2, 2), (-1, 4), (0, 8), (1, 4), (2, 2)]: if 0 <= x + dx < width: img[y + 1, x + dx] += error * w / 32 img = np.clip(img, 0, 255).astype(np.uint8) return Image.fromarray(img, 'RGB') # ============================================================================= # Ordered Dithering Algorithms # ============================================================================= def get_bayer_matrix(n: int) -> np.ndarray: """Generate a Bayer matrix of size 2^n x 2^n.""" if n == 0: return np.array([[0]]) smaller = get_bayer_matrix(n - 1) size = 2 ** (n - 1) result = np.zeros((2 ** n, 2 ** n)) result[:size, :size] = 4 * smaller result[:size, size:] = 4 * smaller + 2 result[size:, :size] = 4 * smaller + 3 result[size:, size:] = 4 * smaller + 1 return result def dither_ordered_bayer(image: Image.Image, matrix_size: int = 4, strength: float = 1.0) -> Image.Image: """ Ordered dithering using Bayer matrix. Args: image: Input image matrix_size: Size of Bayer matrix (2, 4, 8, or 16) strength: Dithering strength multiplier (0.0-2.0) """ img = np.array(image.convert('RGB'), dtype=np.float64) height, width = img.shape[:2] palette = np.array(PALETTE_RGB, dtype=np.float64) # Get appropriate Bayer matrix n = {2: 1, 4: 2, 8: 3, 16: 4}.get(matrix_size, 2) bayer = get_bayer_matrix(n) bayer_size = bayer.shape[0] # Normalize Bayer matrix to -0.5 to 0.5 range, then scale bayer_normalized = (bayer / (bayer_size ** 2) - 0.5) * strength * 128 result = np.zeros_like(img) for y in range(height): for x in range(width): threshold = bayer_normalized[y % bayer_size, x % bayer_size] adjusted_pixel = img[y, x] + threshold adjusted_pixel = np.clip(adjusted_pixel, 0, 255) _, new_pixel = find_nearest_color(adjusted_pixel, palette) result[y, x] = new_pixel return Image.fromarray(result.astype(np.uint8), 'RGB') _NUMBA_PALETTE = np.array([ [0, 0, 0], [255, 255, 255], [255, 255, 0], [255, 0, 0], [0, 0, 255], [0, 255, 0], ], dtype=np.float64) _NUMBA_WEIGHTS = np.array([0.299, 0.587, 0.114], dtype=np.float64) @jit(nopython=True) def _numba_find_nearest(r, g, b, palette, weights): best_idx = 0 best_dist = 1e10 for i in range(palette.shape[0]): dr = (palette[i, 0] - r) * weights[0] dg = (palette[i, 1] - g) * weights[1] db = (palette[i, 2] - b) * weights[2] dist = dr * dr + dg * dg + db * db if dist < best_dist: best_dist = dist best_idx = i return best_idx @jit(nopython=True) def _numba_atkinson(img, palette, weights): height, width = img.shape[0], img.shape[1] for y in range(height): for x in range(width): old_r, old_g, old_b = img[y, x, 0], img[y, x, 1], img[y, x, 2] idx = _numba_find_nearest(old_r, old_g, old_b, palette, weights) new_r, new_g, new_b = palette[idx, 0], palette[idx, 1], palette[idx, 2] img[y, x, 0], img[y, x, 1], img[y, x, 2] = new_r, new_g, new_b err_r = (old_r - new_r) / 8.0 err_g = (old_g - new_g) / 8.0 err_b = (old_b - new_b) / 8.0 if x + 1 < width: img[y, x + 1, 0] += err_r img[y, x + 1, 1] += err_g img[y, x + 1, 2] += err_b if x + 2 < width: img[y, x + 2, 0] += err_r img[y, x + 2, 1] += err_g img[y, x + 2, 2] += err_b if y + 1 < height: if x > 0: img[y + 1, x - 1, 0] += err_r img[y + 1, x - 1, 1] += err_g img[y + 1, x - 1, 2] += err_b img[y + 1, x, 0] += err_r img[y + 1, x, 1] += err_g img[y + 1, x, 2] += err_b if x + 1 < width: img[y + 1, x + 1, 0] += err_r img[y + 1, x + 1, 1] += err_g img[y + 1, x + 1, 2] += err_b if y + 2 < height: img[y + 2, x, 0] += err_r img[y + 2, x, 1] += err_g img[y + 2, x, 2] += err_b return img def dither_atkinson_numba(image: Image.Image) -> Image.Image: """Numba-accelerated Atkinson dithering with perceptual weighting (~150x faster).""" img = np.array(image.convert('RGB'), dtype=np.float64) img = _numba_atkinson(img, _NUMBA_PALETTE, _NUMBA_WEIGHTS) img = np.clip(img, 0, 255).astype(np.uint8) return Image.fromarray(img, 'RGB') # ============================================================================= # PIL Built-in (for comparison) # ============================================================================= def dither_pil_floyd_steinberg(image: Image.Image) -> Image.Image: """PIL's built-in Floyd-Steinberg dithering for comparison.""" pal_image = create_pil_palette_image() img = image.convert('RGB') quantized = img.quantize(dither=Image.Dither.FLOYDSTEINBERG, palette=pal_image) return quantized.convert('RGB') def dither_pil_none(image: Image.Image) -> Image.Image: """PIL quantization with no dithering (nearest color only).""" pal_image = create_pil_palette_image() img = image.convert('RGB') quantized = img.quantize(dither=Image.Dither.NONE, palette=pal_image) return quantized.convert('RGB') # ============================================================================= # Algorithm Registry # ============================================================================= DITHER_ALGORITHMS = { 'none': { 'name': 'No Dithering (PIL)', 'func': dither_pil_none, 'description': 'Simple nearest-color quantization without error diffusion', }, 'pil_fs': { 'name': 'Floyd-Steinberg (PIL)', 'func': dither_pil_floyd_steinberg, 'description': 'PIL built-in Floyd-Steinberg implementation', }, 'floyd_steinberg': { 'name': 'Floyd-Steinberg', 'func': dither_floyd_steinberg, 'description': 'Classic error diffusion (1976), good balance of speed and quality', }, 'floyd_steinberg_weighted': { 'name': 'Floyd-Steinberg (Weighted)', 'func': lambda img: dither_floyd_steinberg(img, weighted=True), 'description': 'Floyd-Steinberg with perceptual color weighting', }, 'atkinson': { 'name': 'Atkinson', 'func': dither_atkinson, 'description': 'Bill Atkinson (Apple), diffuses only 75% of error for cleaner results', }, 'atkinson_weighted': { 'name': 'Atkinson (Weighted)', 'func': lambda img: dither_atkinson(img, weighted=True), 'description': 'Atkinson with perceptual color weighting', }, 'atkinson_fast': { 'name': 'Atkinson (Numba Fast)', 'func': dither_atkinson_numba, 'description': 'Numba-accelerated Atkinson (~150x faster, requires numba)', }, 'jarvis': { 'name': 'Jarvis-Judice-Ninke', 'func': dither_jarvis_judice_ninke, 'description': 'Larger diffusion kernel (1976), smoother gradients but slower', }, 'stucki': { 'name': 'Stucki', 'func': dither_stucki, 'description': 'Similar to JJN with modified weights (1981)', }, 'sierra': { 'name': 'Sierra', 'func': dither_sierra, 'description': 'Full Sierra dithering, balanced results', }, 'sierra_lite': { 'name': 'Sierra Lite', 'func': dither_sierra_lite, 'description': 'Faster Sierra variant with smaller kernel', }, 'burkes': { 'name': 'Burkes', 'func': dither_burkes, 'description': 'Simplified two-row error diffusion', }, 'bayer2': { 'name': 'Ordered (Bayer 2x2)', 'func': lambda img: dither_ordered_bayer(img, matrix_size=2), 'description': 'Ordered dithering with 2x2 Bayer matrix', }, 'bayer4': { 'name': 'Ordered (Bayer 4x4)', 'func': lambda img: dither_ordered_bayer(img, matrix_size=4), 'description': 'Ordered dithering with 4x4 Bayer matrix', }, 'bayer8': { 'name': 'Ordered (Bayer 8x8)', 'func': lambda img: dither_ordered_bayer(img, matrix_size=8), 'description': 'Ordered dithering with 8x8 Bayer matrix', }, 'bayer4_strong': { 'name': 'Ordered (Bayer 4x4 Strong)', 'func': lambda img: dither_ordered_bayer(img, matrix_size=4, strength=1.5), 'description': 'Bayer 4x4 with increased dithering strength', }, } def get_algorithm_names() -> List[str]: """Return list of available algorithm names.""" return list(DITHER_ALGORITHMS.keys()) def apply_dithering(image: Image.Image, algorithm: str) -> Image.Image: """Apply the specified dithering algorithm to an image.""" if algorithm not in DITHER_ALGORITHMS: raise ValueError(f"Unknown algorithm: {algorithm}. Available: {get_algorithm_names()}") return DITHER_ALGORITHMS[algorithm]['func'](image)