Experiment: dither with only two colors per blend

dither.py

@@ -1,14 +1,19 @@
 import numpy as np
-from PIL import Image
+import scipy.spatial
+import PIL
+from PIL.Image import Image
 from tqdm import tqdm
 
-img = Image.open('peacock.jpg')
+img = PIL.Image.open('peacock.jpg')
 resized = img.resize((128,128))
 
+RGB = tuple[int, int, int]
+RGBLinear = tuple[float, float, float]
+
 # sRGB values for the PICO-8 palette.
 # The first 16 entries correspond with the default palette (0-15).
 # The last 16 entries correspond with the "secret" palette (128-143).
-PICO_RGB = [
+PICO_RGB: list[RGB] = [
     (0, 0, 0),
     (29, 43, 83),
     (126, 37, 83),
@@ -43,19 +48,36 @@ PICO_RGB = [
     (255, 157, 129)
 ]
 
-def component_to_linear(x):
+def component_to_linear(x: int) -> float:
     x /= 255
     if x <= 0.04045:
         return x/12.92
     else:
         return ((x + 0.055)/1.055)**2.4
 
-def rgb_to_linear(rgb):
+def component_to_gamma(x: float) -> int:
+    if x <= 0.04045/12.92:
+        result = x * 12.92
+    else:
+        result = (x**(1/2.4) * 1.055) - 0.055
+    result *= 255
+    return round(result)
+
+def rgb_to_linear(rgb: RGB) -> RGBLinear:
     return tuple(component_to_linear(c) for c in rgb)
 
+def linear_to_rgb(linear: RGBLinear) -> RGB:
+    return tuple(component_to_gamma(c) for c in linear)
+
 PICO_RGB_LINEAR = [rgb_to_linear(x) for x in PICO_RGB]
 
-def make_bayer_matrix(size):
+def distance_sq(a: RGB, b: RGB) -> int:
+    result = 0
+    for i in range(3):
+        result += (a[i] - b[i])**2
+    return result
+
+def make_bayer_matrix(size: int) -> np.ndarray:
     assert size > 0
     assert size & (size - 1) == 0 # power of two
     m = np.array([[0]])
@@ -104,7 +126,68 @@ def pattern_dither(img, palette, *, pat_size=4, amount=0.75):
                     break
 
     # Convert palette indexes to an indexed image
-    result = Image.new(mode='P', size=img.size)
+    result = PIL.Image.new(mode='P', size=img.size)
     result.putpalette([c for rgb in old_palette for c in rgb])
     result.putdata([old_palette_indexes[i] for i in output_indexes])
     return result
+
+
+def luminance(rgb: RGB) -> int:
+    return 3*rgb[0] + 6*rgb[1] + rgb[2]
+
+def pattern_dither_twotone(img: Image, palette: list[RGB]) -> Image:
+    assert img.mode == 'RGB'
+    bayer = make_bayer_matrix(4)
+
+    # Sort palette indexes by luminance
+    # Used to make sure blends consistently order their colors
+    palette_indexes = list(range(len(palette)))
+    palette_indexes.sort(key=lambda i: luminance(palette[i]))
+
+    # Build lookup for all two-color blends
+    linear_palette = [np.array(rgb_to_linear(c)) for c in palette]
+    blends: list[RGB] = []
+    recipes: list[tuple[int,int,int]] = [] # index 1, index 2, ratio/16
+    def make_blend(a_idx: int, b_idx: int, ratio: int):
+        if distance_sq(palette[a_idx], palette[b_idx]) > 3*128**2:
+            return
+        a = linear_palette[a_idx]
+        b = linear_palette[b_idx]
+        blend = (a*(16-ratio) + b*ratio)/16
+        blend = linear_to_rgb(blend)
+        blends.append(blend)
+        recipes.append((a_idx, b_idx, ratio))
+
+    for i in range(len(linear_palette)):
+        make_blend(i, i, 0)
+    for a_idx in range(len(linear_palette)):
+        for b_idx in range(a_idx + 1, len(linear_palette)):
+            for ratio in range(1, 16):
+                make_blend(a_idx, b_idx, ratio)
+    kd = scipy.spatial.KDTree(blends)
+    print(len(blends), len(set(blends)))
+    print(max(blends))
+    print(min(blends))
+
+
+    # Convert image pixels
+    result_indexes = []
+    for y in tqdm(range(img.height)):
+        # row = [rgb_to_linear(img.getpixel((x,y))) for x in range(img.width)]
+        # _, blends_indexes = kd.query(row)
+        for x in range(img.width):
+            color = img.getpixel((x, y))
+            _dist, blends_idx = kd.query(color)
+            recipe = recipes[blends_idx]
+            thresh = bayer[y % 4][x % 4]
+            if recipe[2] <= thresh:
+                idx = recipe[0]
+            else:
+                idx = recipe[1]
+            result_indexes.append(idx)
+
+    # Output image
+    result = PIL.Image.new(mode='P', size=img.size)
+    result.putpalette([c for rgb in palette for c in rgb])
+    result.putdata(result_indexes)
+    return result