# /// script # requires-python = ">=3.10" # dependencies = [ # "numpy>=2.3.0,<3", # "pillow>=11.2.1,<12", # "matplotlib>=3.10.3,<4", # "scipy>=1.15.3,<2", # "numba==0.63.1", # ] # /// """ Ordered Dither Color Selection via simplified Yliluoma's and Knoll's algorithms. Pekka Väänänen, 2026 See https://30fps.net/pages/revisiting-yliluoma-2/ Usage example: uv run color_selection.py originals/bigbuckbunny_bird.png 16 32 --algo=ema-exact License: CC0 https://creativecommons.org/public-domain/cc0/ """ import sys import numpy as np from dataclasses import dataclass import numba import argparse from PIL import Image bayer_4x4 = np.array([ [ 0, 8, 2, 10], [12, 4, 14, 6], [ 3, 11, 1, 9], [15, 7, 13, 5]]) def srgb2linear_(arr): """In-place sRGB-to-linear conversion.""" assert isinstance(arr, np.ndarray) assert arr.dtype == np.float32 or arr.dtype == np.float64 assert arr.max() <= 1.0 # from scikit-image mask = arr > 0.04045 arr[mask] = np.power((arr[mask] + 0.055) / 1.055, 2.4) arr[~mask] /= 12.92 def srgb2linear(arr): y = arr.copy() srgb2linear_(y) return y # libimagequant's luminance-weighted sRGB in 1.75 gamma space liq_weights = np.array([0.5, 1, 0.45, 0.625]) liq_exponent = 0.57 / 0.45455 liq_inv_exponent = 1.0 / liq_exponent def srgb_to_liq(srgb): weighted = srgb.copy() weighted[..., :3] **= liq_exponent return weighted * liq_weights[:3].reshape(1,-1) return weighted def dither_none(img, palette, linear=False): """ Pixel mapping via Euclidean sRGB distance without dithering. """ img_float = img / 255 palette_float = palette / 255 if linear: img_float = srgb2linear(img_float) palette_float = srgb2linear(palette_float) H, W, _ = img.shape dist = np.zeros(palette.shape[0]) def find_index_of_closest(color): c = color.reshape(1, -1) dist[:] = np.sum((palette_float - c) ** 2, axis=1) return np.argmin(dist) inds = np.zeros((H, W), dtype=np.int32) for y in range(H): for x in range(W): color = img_float[y, x] idx = find_index_of_closest(color) inds[y, x] = idx return inds def dither_offset(img, palette, linear=False, dither_strength=0.5): """ A naive ordered dither that adds a constant greyscale noise term to the input colors before a closest color lookup. Estimates the offset via an ad-hoc formula that works the best for gamma space inputs. """ img_float = img / 255 palette_float = palette / 255 if linear: img_float = srgb2linear(img_float) palette_float = srgb2linear(palette_float) H, W, _ = img.shape K, _ = palette.shape # Create a list of pairwise palette color distances dists = [] for k1 in range(K): for k2 in range(k1+1, K): assert k1 != k2 d = np.sqrt(np.sum((palette_float[k2] - palette_float[k1])**2)) dists.append(d) # Calculate a decent offset size for dithering that # works for the most common colors. offset = np.median(np.array(dists)) offset /= 2.0 offset *= dither_strength print(f"Dither offset magnitude: {float(offset):.3f}") dist = np.zeros(palette.shape[0]) def find_index_of_closest(color): c = color.reshape(1, -1) dist[:] = np.sum((palette_float - c) ** 2, axis=1) return np.argmin(dist) inds = np.zeros((H, W), dtype=np.int32) for y in range(H): for x in range(W): color = img_float[y, x] # (0, 1) range threshold threshold = (bayer_4x4[y%4,x%4] + 0.5) / 16.0 # bias threshold to be (-0.5, 0.5) and add moved = color + offset * (threshold - 0.5) idx = find_index_of_closest(moved) inds[y, x] = idx return inds def dither_knoll_slow(img, palette, N, dither_strength, linear=False): img_float = img / 255 palette_float = palette / 255 H, W, _ = img.shape if linear: img_float = srgb2linear(img_float) palette_float = srgb2linear(palette_float) luma = palette.dot([0.309, 0.609, 0.082]) dist = np.zeros(palette.shape[0]) def find_closest(color): c = color.reshape(1, -1) dist[:] = np.sum((palette_float - c) ** 2, axis=1) return np.argmin(dist) # Output array inds = np.zeros((H, W), dtype=np.int32) error = np.zeros((3,),float) for y in range(H): for x in range(W): color = img_float[y, x] candidates = [] error[:] = 0.0 for iter in range(N): idx = find_closest(color + error * dither_strength) error += color - palette_float[idx] candidates.append(idx) threshold = (bayer_4x4[y%4,x%4] + 0.5) / 16.0 luma_order = np.argsort(np.take(luma, candidates)) cand_idx = int(threshold * len(luma_order)) # Threshold is never exactly 1.0 so this works idx = candidates[luma_order[cand_idx]] inds[y, x] = idx return inds def dither_knoll(img, palette, N, dither_strength, linear=False): img_float = img / 255 palette_float = palette / 255 H, W, _ = img.shape if linear: img_float = srgb2linear(img_float) palette_float = srgb2linear(palette_float) # sort palette by approximate luminance luma = palette.dot([0.309, 0.609, 0.082]) luma_order = np.argsort(luma) dist = np.zeros(palette.shape[0]) def find_closest(color): c = color.reshape(1, -1) dist[:] = np.sum((palette_float - c) ** 2, axis=1) return np.argmin(dist) # Output array inds = np.zeros((H, W), dtype=np.int32) # Work arrays allocated outside the main loop weights = np.zeros(K, dtype=float) error = np.zeros(3, dtype=float) for y in range(H): for x in range(W): color = img_float[y, x] # Clear work arrays weights[:] = 0.0 # candidate weights weight_sum = 0.0 error[:] = 0.0 # accumulated error compensation # Find candidates, can be less than N unique colors for _ in range(N): idx = find_closest(color + error * dither_strength) weights[idx] += 1 weight_sum += 1 error += color - palette_float[idx] weights /= weight_sum threshold = (bayer_4x4[y%4,x%4] + 0.5) / 16.0 # Output the candidate index on which cumulative weight sum # first crosses the threshold read from the 4x4 matrix # If a paltte index wasn't used, its weight will be zero # and it won't be selected by this loop. cumulative_sum = 0.0 for idx in luma_order: w_ko = weights[idx] cumulative_sum += w_ko if cumulative_sum > threshold: break inds[y, x] = idx return inds def dither_ema_sweep(img, palette, N, linear=False, sweep_range=(0.2, 1.0), sweep_tests=8, use_luma=False): """ A generic version of Yliluoma's method 2. This one allows changing the search factor range from the original's [0.5, 1.0) that was implicitly encoded in its running average implementation. Uses Numba for acceleration since otherwise it's very slow. """ img_float = img / 255 palette_float = palette / 255 H, W, _ = img.shape if linear: img_float = srgb2linear(img_float) palette_float = srgb2linear(palette_float) # sort palette by approximate luminance # Yliluoma's code uses Rec.602 weights but we use a more balanced weights. # The textbook by W. Burger and M. J. Burge attributes these to Charles Poynton. luma_weights = np.array([0.309, 0.609, 0.082]) luma = palette.dot(luma_weights) order = np.argsort(luma) # Consider every mixing fraction in a quantized (0, 1) range between # the mean color and the palette color. Pick as the candidate the # color with the shortest distance to the input color. # The C++ code also accelerates the fraction search not checking every value # linearly but adding more tests as candidates are added. # Here we do a slower linear search. min_t, max_t = sweep_range if min_t == 0.0: min_t += 0.05 num_steps = sweep_tests # round((max_t - min_t) / step_size) fractions = np.linspace(min_t, max_t, num_steps, endpoint=False) print("fractions", len(fractions)) print(fractions) a_weights = 1 - fractions[:, np.newaxis] b_weights = fractions[:, np.newaxis] dist = np.zeros(palette.shape[0]) @numba.jit(cache=True) def find_closest(color, dist): c = color.reshape(1, -1) dist[:] = np.sum((palette_float - c) ** 2, axis=1) return np.argmin(dist) @numba.jit(cache=True) def compute_candidate_weights(p, dist, out_weights): closest_idx = find_closest(p, dist) mean = palette_float[closest_idx] out_weights[:] = 0.0 out_weights[closest_idx] = 1.0 K, _ = palette_float.shape # Find a candidate point 'N' times for it in range(N): best_k = -1 best_t = 0.0 best_dist = 1e9 # Consider every palette_float color for k in range(K): c_k = palette_float[k] a = a_weights * mean.reshape(1,3) b = b_weights * c_k.reshape(1,3) mix = a + b diffs = (p - mix) if use_luma: lumadiff = p.dot(luma_weights) - mix.dot(luma_weights) color_dists = np.sum((diffs ** 2) * luma_weights, axis=-1) color_dists = color_dists * 0.75 + lumadiff**2 else: color_dists = np.sum((diffs ** 2), axis=-1) fraction_idx = np.argmin(color_dists) cand_dist = color_dists[fraction_idx] if cand_dist < best_dist: best_k = k best_t = fractions[fraction_idx] best_dist = cand_dist assert best_k != -1 out_weights[best_k] += best_t mean = (1-best_t)*mean + best_t*palette_float[best_k] out_weights /= np.maximum(1e-6, np.sum(out_weights)) # Function doesn't return anything since weights were written out already inds = np.zeros((H, W), dtype=np.int32) weights = np.zeros(palette_float.shape[0]) for y in range(H): print(f"[{y+1}/{H}] {((y*W+W)/(H*W))*100:.1f} %") for x in range(W): compute_candidate_weights(img_float[y, x], dist, weights) threshold = (bayer_4x4[y%4,x%4] + 0.5) / 16.0 cumulative_sum = 0.0 idx = None for idx in order: # if a point wasn't used, its weight is zero and it won't ever be selected cumulative_sum += weights[idx] if cumulative_sum > threshold: break assert idx is not None inds[y, x] = idx return inds def dither_ema_exact(img, palette, N, linear=False, search_range=(0.2, 1.0)): """ A generic version of Yliluoma's method 2. This one allows changing the search range from the original's [0.5, 1.0) that was implicitly encoded in its running average implementation. Because luma weighting isn't used the mixing fraction 't' can be solved analytically instead of via an exhaustive search. Uses Numba for acceleration of the inner loops. """ img_float = img / 255 palette_float = palette / 255 H, W, _ = img.shape if linear: img_float = srgb2linear(img_float) palette_float = srgb2linear(palette_float) # Sort palette by approximate luminance # Yliluoma's code uses Rec.602 weights but we use a more balanced weights. # The textbook by W. Burger and M. J. Burge attributes these to Charles Poynton. luma_weights = np.array([0.309, 0.609, 0.082]) luma = palette.dot(luma_weights) luma_order = np.argsort(luma) # The original C++ code seems to begin its fraction search from a 50/50 # division towards 1.0. This encourages "more movement" in the mean point # that seems to help a bit in keeping the point set smaller. min_t, max_t = search_range # A mix fraction 't' of 0.0 would mean not moving 'mean' at all. # But here the point is to always pick _some_ palette color to move towards. # Conversely, having t==1.0 would mean completely discarding the mean accumulated # so far. It's not something we want either. if min_t == 0.0: min_t += 0.025 if max_t == 1.0: max_t -= 0.025 dist = np.zeros(palette.shape[0]) @numba.jit(cache=True) def find_closest(color, dist): c = color.reshape(1, -1) dist[:] = np.sum((palette_float - c) ** 2, axis=1) return np.argmin(dist) @numba.jit(cache=True) def compute_candidate_weights(p, dist, out_weights): closest_idx = find_closest(p, dist) mean = palette_float[closest_idx] out_weights[:] = 0.0 out_weights[closest_idx] = 1.0 K, _ = palette_float.shape # Find a candidate point 'N' times for it in range(N): best_k = -1 best_t = 0.0 best_dist = 1e9 # Consider every palette color for k in range(K): c_k = palette_float[k] delta = c_k - mean # Solve for the mixing factor 't' that minimizes the distance to the input color 'p' t = 0.0 delta2 = delta.dot(delta) if delta2 >= 1e-9: t = (p - mean).dot(delta) / delta2 # Limit the choice to the given range t = max(min_t, min(max_t, float(t))) mix = (1-t) * mean + t * c_k diff = p - mix cand_dist = diff.dot(diff) if cand_dist < best_dist: best_k = k best_t = t best_dist = cand_dist assert best_k != -1 out_weights[best_k] += best_t mean = (1-best_t)*mean + best_t*palette_float[best_k] out_weights /= np.maximum(1e-6, np.sum(out_weights)) # Doesn't return anything. Weights were written out already inds = np.zeros((H, W), dtype=np.int32) weights = np.zeros(palette_float.shape[0]) for y in range(H): print(f"[{y+1}/{H}] {((y*W+W)/(H*W))*100:.1f} %") for x in range(W): compute_candidate_weights(img_float[y, x], dist, weights) # Choose the palette entry at whose index the cumulative # sum of probabilities crosses 'threshold' read from the # Bayer matrix. threshold = (bayer_4x4[y%4,x%4] + 0.5) / 16.0 cumulative_sum = 0.0 idx = None for idx in luma_order: # if a point wasn't used, its weight is zero and it won't ever be selected cumulative_sum += weights[idx] if cumulative_sum > threshold: break assert idx is not None inds[y, x] = idx return inds def dither_ema_constant(img, palette, N, linear=False, t=0.3): """ An Yliluoma-2 variant that doesn't even try to find an optimal location on the segment from the current goal to a candidate but uses a constant value instead. Since the closest possible value of 't' is chosen 99% cases in the other variants, this version works almost the same. The method simpler than 'dither_ema_exact()' but not really faster. Included here for completeness. """ img_float = img / 255 palette_float = palette / 255 H, W, _ = img.shape if linear: img_float = srgb2linear(img_float) palette_float = srgb2linear(palette_float) # Sort palette by approximate luminance # Yliluoma's code uses Rec.602 weights but we use a more balanced weights. # The textbook by W. Burger and M. J. Burge attributes these to Charles Poynton. luma_weights = np.array([0.309, 0.609, 0.082]) luma = palette.dot(luma_weights) luma_order = np.argsort(luma) dist = np.zeros(palette.shape[0]) @numba.jit(cache=True) def find_closest(color, dist): c = color.reshape(1, -1) dist[:] = np.sum((palette_float - c) ** 2, axis=1) return np.argmin(dist) @numba.jit(cache=True) def compute_candidate_weights(p, dist, out_weights): closest_idx = find_closest(p, dist) mean = palette_float[closest_idx] out_weights[:] = 0.0 out_weights[closest_idx] = 1.0 K, _ = palette_float.shape # Find a candidate point 'N' times for iter in range(N): best_k = -1 best_t = 0.0 best_dist = 1e9 # Consider every palette color for k in range(K): c_k = palette_float[k] mix = (1-t) * mean + t * c_k diff = p - mix cand_dist = diff.dot(diff) if cand_dist < best_dist: best_k = k best_t = t best_dist = cand_dist assert best_k != -1 out_weights[best_k] += best_t mean = (1-best_t)*mean + best_t*palette_float[best_k] out_weights /= np.maximum(1e-6, np.sum(out_weights)) # Doesn't return anything. Weights were written out already inds = np.zeros((H, W), dtype=np.int32) weights = np.zeros(palette_float.shape[0]) for y in range(H): print(f"[{y}] {y*W+1}/{H*W}") for x in range(W): compute_candidate_weights(img_float[y, x], dist, weights) # Choose the palette entry at whose index the cumulative # sum of probabilities crosses 'threshold' read from the # Bayer matrix. threshold = (bayer_4x4[y%4,x%4] + 0.5) / 16.0 cumulative_sum = 0.0 idx = None for idx in luma_order: # if a point wasn't used, its weight is zero and it won't ever be selected cumulative_sum += weights[idx] if cumulative_sum > threshold: break assert idx is not None inds[y, x] = idx return inds def maximin_init_vectorized(X, K, distance=None): centers = [] centers.append(np.mean(X, axis=0)) dist = np.zeros((X.shape[0],K)) for iter in range(K-1): if distance is None: dist[:, iter] = np.sum((X - centers[iter])**2, axis=1) else: dist[:, iter] = distance(X, centers[iter]) # 'dist' has now distance to every center, for every point. # We want to first take the minimum over columns, # and then the index of the point with the largest value aka # longest distance to any center. k = len(centers) closest = dist[:, :k].min(axis=1) max_dist_i = np.argmax(closest) new_center = X[max_dist_i] centers.append(new_center) return np.array(centers) def kmeans_loop(X, init_centers, num_iter): N, _ = X.shape K = init_centers.shape[0] C = init_centers.shape[1] assert C == X.shape[1], "init and data dims must match" centers = init_centers.copy().astype(float) dist = np.zeros((N, K)) for _ in range(num_iter): for k in range(K): center_k = centers[k, np.newaxis] dist[:, k] = np.sum((X - center_k) ** 2, axis=1) indices = np.argmin(dist, axis=1) for k in range(K): k_inds = indices == k if not k_inds.any(): continue centers[k] = np.mean(X[k_inds], axis=0) return centers, indices parser = argparse.ArgumentParser(description='OKM image quantization') parser.add_argument('input', help='Input image path') parser.add_argument('num_colors', type=int, help='Number of colors in the palette (K)') parser.add_argument('N', type=int, help='Number of candidate search iterations') parser.add_argument('--algo', choices=['none', 'offset', 'ema-sweep', 'ema-exact', 'ema-constant', 'knoll', 'knoll-slow'], default='ema-exact') parser.add_argument('--kmeans', type=int, default=5, help='Number of kmeans rounds after maximin init (default: 5)') parser.add_argument('--linear', action='store_true', default=False, help='Use linear RGB instead of sRGB') parser.add_argument('--use-luma', action='store_true', default=False, help='Use luma-weighting in ema-sweep') parser.add_argument('--strength', type=float, default=0.8, help="Knoll dither strength (default: 0.8)") parser.add_argument('--sweep-tests', type=int, default=8, help="How many mixing fractions to test") parser.add_argument('--palette', choices=['maximin', 'jy16'], default='maximin') parser.add_argument('--dump-palette', action='store_true', help="Writes the computed palette as a C array and exits.") parser.add_argument('--colorspace', choices=['srgb', 'liq-rgb'], default='srgb') args = parser.parse_args() K = args.num_colors print(f"Converting {args.input} to {K=} colors") img = Image.open(args.input).convert('RGB') original = np.array(img) H, W, C = original.shape X = original.reshape(-1, 3)/255.0 if args.palette == "maximin": # Compute some sort of a palette with k-means and maximin init centers = maximin_init_vectorized(X, K) if args.kmeans: centers, _ = kmeans_loop(X, centers, args.kmeans) pal_8bit = (centers*255).round().clip(0, 255).astype(np.uint8) elif args.palette == "jy16": # A 16-color test palette from https://bisqwit.iki.fi/story/howto/dither/jy/ if args.num_colors != 16: raise ValueError(f"palette={args.palette} supports only 16 colors") pal_8bit = np.zeros((16,3), dtype=np.uint8) for i, c in enumerate([ 0x080000,0x201A0B,0x432817,0x492910, 0x234309,0x5D4F1E,0x9C6B20,0xA9220F, 0x2B347C,0x2B7409,0xD0CA40,0xE8A077, 0x6A94AB,0xD5C4B3,0xFCE76E,0xFCFAE2]): r,g,b = (c >> 16) & 255, (c >> 8) & 255, c & 255 pal_8bit[i] = (r,g,b) else: raise ValueError(f"Invalid palette {args.palette}") if args.dump_palette: print("unsigned int colors[] = {") for i in range(pal_8bit.shape[0]): c = pal_8bit[i].astype(np.uint32) print(" " + hex((c[0] << 16) | (c[1] << 8) | c[2]), end='') if i < pal_8bit.shape[0]-1: print(",") else: print() print("};") sys.exit(0) input_pal = pal_8bit input_img = original if args.colorspace == "liq-rgb": def to_8bit_liq_rgb(srgb): liq = srgb_to_liq(srgb / 255) return (liq * 255).round().clip(0,255).astype(np.uint8) input_pal = to_8bit_liq_rgb(pal_8bit) input_img = to_8bit_liq_rgb(original) if args.use_luma and args.algo != "ema-sweep": raise RuntimeError("use_luma only works with ema-sweep") if args.sweep_tests != 8 and args.algo != "ema-sweep": raise RuntimeError("Number of sweep tests can be set only for ema-sweep") if args.algo == "none": inds = dither_none(input_img, input_pal, linear=args.linear) elif args.algo == "offset": inds = dither_offset(input_img, input_pal, linear=args.linear, dither_strength=args.strength) elif args.algo == "ema-sweep": inds = dither_ema_sweep(input_img, input_pal, args.N, linear=args.linear, use_luma=args.use_luma, sweep_tests=args.sweep_tests) elif args.algo == "ema-exact": inds = dither_ema_exact(input_img, input_pal, args.N, linear=args.linear) elif args.algo == "ema-constant": inds = dither_ema_constant(input_img, input_pal, args.N, linear=args.linear) elif args.algo == "knoll": inds = dither_knoll(input_img, input_pal, args.N, linear=args.linear, dither_strength=args.strength) elif args.algo == "knoll-slow": inds = dither_knoll_slow(input_img, input_pal, args.N, linear=args.linear, dither_strength=args.strength) else: raise ValueError("Unknown color selection algorithm") y = pal_8bit[inds.reshape(-1)].reshape(H, W, C) from pathlib import Path img_name = Path(args.input).stem tail = "" if args.sweep_tests != 8: tail = f"{tail}_sweep{args.sweep_tests}" if args.linear: tail = f"{tail}_linear" if args.use_luma: tail = f"{tail}_luma" if args.colorspace != "srgb": tail = f"{tail}_{args.colorspace}" if args.algo == "knoll" or args.algo == "knoll-slow": tail = f"{tail}_str{int(round(args.strength*100))}" if args.palette == "maximin": tail = f"{tail}_km{args.kmeans}" else: tail = f"{tail}_{args.palette}" outname = f"{img_name}_K{args.num_colors}_N{args.N}_{args.algo}{tail}.png" print(f"Saving {outname}") Image.fromarray(y).save(outname)