136 lines
5.2 KiB
Python
136 lines
5.2 KiB
Python
"""Palette-constrained dithering.
|
|
|
|
Every routine takes the working image in CIELAB plus a per-pixel table of the
|
|
palette indices that pixel is *allowed* to use (because the VIC-II only lets a
|
|
given screen cell show a small set of colours), and returns an (H,W) image of
|
|
chosen palette indices (0..15). Because the allowed set is per-pixel, error that
|
|
diffuses across a cell boundary is automatically re-clamped to the neighbour
|
|
cell's own colours -- exactly the constraint real C64 hardware imposes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
DITHER_MODES = ["bayer", "floyd", "atkinson", "stucki", "jarvis", "none"]
|
|
|
|
|
|
def bayer_matrix(n: int) -> np.ndarray:
|
|
"""Normalised (0..1) Bayer threshold matrix of size n x n (n power of two)."""
|
|
if n == 1:
|
|
return np.array([[0.0]])
|
|
smaller = bayer_matrix(n // 2)
|
|
m = np.block([
|
|
[4 * smaller + 0, 4 * smaller + 2],
|
|
[4 * smaller + 3, 4 * smaller + 1],
|
|
])
|
|
return m / (n * n)
|
|
|
|
|
|
def _gather_colors(palette_lab: np.ndarray, allowed: np.ndarray) -> np.ndarray:
|
|
# allowed: (H,W,K) palette indices -> (H,W,K,3) Lab
|
|
return palette_lab[allowed]
|
|
|
|
|
|
def quantize_ordered(img_lab, allowed, palette_lab, strength=1.0, n=8):
|
|
"""Ordered (Bayer) dithering between the two best colours of each pixel's set.
|
|
|
|
For every pixel we find its nearest and second-nearest allowed colour, project
|
|
the pixel onto the segment between them, and use the Bayer threshold to decide
|
|
which of the two to emit -- giving smooth ordered blends without ever leaving
|
|
the cell's legal colour set.
|
|
"""
|
|
H, W, _ = img_lab.shape
|
|
colors = _gather_colors(palette_lab, allowed) # (H,W,K,3)
|
|
d = np.sum((img_lab[:, :, None, :] - colors) ** 2, axis=-1) # (H,W,K)
|
|
|
|
i1 = np.argmin(d, axis=-1)
|
|
d2 = np.array(d)
|
|
np.put_along_axis(d2, i1[..., None], np.inf, axis=-1)
|
|
i2 = np.argmin(d2, axis=-1)
|
|
|
|
yy, xx = np.indices((H, W))
|
|
c1 = colors[yy, xx, i1] # (H,W,3)
|
|
c2 = colors[yy, xx, i2]
|
|
seg = c2 - c1
|
|
seg_len2 = np.sum(seg * seg, axis=-1) + 1e-9
|
|
t = np.sum((img_lab - c1) * seg, axis=-1) / seg_len2 # projection 0..1
|
|
t = np.clip(t * strength, 0.0, 1.0)
|
|
|
|
thr = bayer_matrix(n)
|
|
thr_full = thr[yy % n, xx % n]
|
|
pick2 = t > thr_full
|
|
chosen = np.where(pick2, i2, i1)
|
|
return np.take_along_axis(allowed, chosen[..., None], axis=-1)[..., 0]
|
|
|
|
|
|
def _quantize_diffusion(img_lab, allowed, palette_lab, kernel, divisor):
|
|
"""Generic serpentine error-diffusion constrained to per-pixel allowed sets."""
|
|
H, W, _ = img_lab.shape
|
|
work = img_lab.astype(np.float64).copy()
|
|
out = np.zeros((H, W), dtype=np.int64)
|
|
pal = palette_lab
|
|
for y in range(H):
|
|
cols = range(W) if (y % 2 == 0) else range(W - 1, -1, -1)
|
|
flip = 1 if (y % 2 == 0) else -1
|
|
for x in cols:
|
|
allow = allowed[y, x]
|
|
cand = pal[allow]
|
|
diff = cand - work[y, x]
|
|
k = int(allow[np.argmin(np.sum(diff * diff, axis=-1))])
|
|
out[y, x] = k
|
|
err = work[y, x] - pal[k]
|
|
for dx, dy, w in kernel:
|
|
nx, ny = x + dx * flip, y + dy
|
|
if 0 <= nx < W and 0 <= ny < H:
|
|
work[ny, nx] += err * (w / divisor)
|
|
return out
|
|
|
|
|
|
# (dx, dy, weight) relative to current pixel, assuming left-to-right scan.
|
|
_FLOYD = [(1, 0, 7), (-1, 1, 3), (0, 1, 5), (1, 1, 1)]
|
|
_ATKINSON = [(1, 0, 1), (2, 0, 1), (-1, 1, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1)]
|
|
# Larger kernels spread error further -> smoother gradients (best for grayscale).
|
|
_STUCKI = [(1, 0, 8), (2, 0, 4),
|
|
(-2, 1, 2), (-1, 1, 4), (0, 1, 8), (1, 1, 4), (2, 1, 2),
|
|
(-2, 2, 1), (-1, 2, 2), (0, 2, 4), (1, 2, 2), (2, 2, 1)]
|
|
_JARVIS = [(1, 0, 7), (2, 0, 5),
|
|
(-2, 1, 3), (-1, 1, 5), (0, 1, 7), (1, 1, 5), (2, 1, 3),
|
|
(-2, 2, 1), (-1, 2, 3), (0, 2, 5), (1, 2, 3), (2, 2, 1)]
|
|
|
|
|
|
def quantize_floyd(img_lab, allowed, palette_lab):
|
|
return _quantize_diffusion(img_lab, allowed, palette_lab, _FLOYD, 16)
|
|
|
|
|
|
def quantize_atkinson(img_lab, allowed, palette_lab):
|
|
return _quantize_diffusion(img_lab, allowed, palette_lab, _ATKINSON, 8)
|
|
|
|
|
|
def quantize_stucki(img_lab, allowed, palette_lab):
|
|
return _quantize_diffusion(img_lab, allowed, palette_lab, _STUCKI, 42)
|
|
|
|
|
|
def quantize_jarvis(img_lab, allowed, palette_lab):
|
|
return _quantize_diffusion(img_lab, allowed, palette_lab, _JARVIS, 48)
|
|
|
|
|
|
def quantize_none(img_lab, allowed, palette_lab):
|
|
colors = _gather_colors(palette_lab, allowed)
|
|
d = np.sum((img_lab[:, :, None, :] - colors) ** 2, axis=-1)
|
|
i1 = np.argmin(d, axis=-1)
|
|
return np.take_along_axis(allowed, i1[..., None], axis=-1)[..., 0]
|
|
|
|
|
|
def quantize(img_lab, allowed, palette_lab, mode="bayer"):
|
|
if mode == "bayer":
|
|
return quantize_ordered(img_lab, allowed, palette_lab)
|
|
if mode == "floyd":
|
|
return quantize_floyd(img_lab, allowed, palette_lab)
|
|
if mode == "atkinson":
|
|
return quantize_atkinson(img_lab, allowed, palette_lab)
|
|
if mode == "stucki":
|
|
return quantize_stucki(img_lab, allowed, palette_lab)
|
|
if mode == "jarvis":
|
|
return quantize_jarvis(img_lab, allowed, palette_lab)
|
|
return quantize_none(img_lab, allowed, palette_lab)
|