Rewrite pattern-dither code

- Artifact reduction: sort colors by luminance
- Use NumPy to speed up dithering
- Allow dither amount to be customized

dither.py

@@ -1,6 +1,6 @@
 import numpy as np
 from PIL import Image
-import tqdm
+from tqdm import tqdm
 
 img = Image.open('peacock.jpg')
 resized = img.resize((128,128))
@@ -58,13 +58,57 @@ PICO_RGB_LINEAR = [rgb_to_linear(x) for x in PICO_RGB]
 def make_bayer_matrix(size):
     assert size > 0
     assert size & (size - 1) == 0 # power of two
-    m = np.array([[0.0]])
+    m = np.array([[0]])
     while m.shape[0] < size:
         m = np.block([[4*m, 4*m+3],[4*m+2, 4*m+1]])
     return m
 
 bayer = make_bayer_matrix(4)
 
+def pattern_dither(img, palette, *, pat_size=4, amount=0.75):
+    assert img.mode == 'RGB'
+    bayer = make_bayer_matrix(pat_size)
+
+    # Reorder palette by luminance
+    old_palette = palette
+    reordered = list(zip(palette, range(len(palette))))
+    reordered.sort(key=lambda c: 3*c[0][0] + 6*c[0][1] + c[0][2])
+    palette = [c for c, _ in reordered]
+    old_palette_indexes = [i for _, i in reordered]
+
+    # Convert palette colors to linear RGB
+    palette = [rgb_to_linear(c) for c in palette]
+    palette_matrix = np.array(palette)
+
+    def nearest(lc):
+        """Takes a linear RGB color c and returns a palette index"""
+        squared_errors = (palette_matrix - lc)**2
+        return np.argmin(np.sum(squared_errors, axis=1))
+
+    # Choose palette indexes for every pixel in the image
+    output_indexes = []
+    for y in tqdm(range(img.height)):
+        for x in range(img.width):
+            color = np.array(rgb_to_linear(img.getpixel((x, y))))
+            err = np.zeros((3,)) + color
+            counts = [0] * len(palette)
+            for _ in range(bayer.size):
+                i = nearest(err)
+                counts[i] += 1
+                err = (err - palette_matrix[i]) * amount + color
+            thresh = bayer[y % pat_size][x % pat_size]
+            for i, count in enumerate(counts):
+                thresh -= count
+                if thresh < 0:
+                    output_indexes.append(i)
+                    break
+
+    # Convert palette indexes to an indexed image
+    result = 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 dither_pixel(img, xy, pal):
     global bayer
     pal2 = [np.array(rgb_to_linear(c)) for c in pal]
@@ -92,10 +136,10 @@ dithered = Image.new('P', resized.size)
 dithered.putpalette([x for rgb in PICO_RGB[:16] for x in rgb])
 def do_it():
     global dithered, resized
-    for x in tqdm.tqdm(range(128)):
+    for x in tqdm(range(128)):
         for y in range(128):
             i = dither_pixel(resized, (x,y), PICO_RGB[:16])
             dithered.putpixel((x,y), i)
 # import cProfile
 # cProfile.run('do_it()')
-do_it()
+#do_it()