Implement two-tone for 8x8 blocks of an image

two-tone.py

@@ -0,0 +1,140 @@
+import math
+import PIL
+from PIL.Image import Image
+from functools import lru_cache
+
+RGB = tuple[int, int, int]
+
+# Load image, palette
+img: Image = PIL.Image.open('peacock-bayer.png')
+
+# Get palette info
+assert img.mode == 'P'
+raw_pal = img.getpalette()
+palette: list[RGB] = []
+for i in range(0, len(raw_pal), 3):
+    palette.append(tuple(raw_pal[i:i+3]))
+palette_indexes = {c: i for i, c in enumerate(palette)}
+
+# Get darkest/lightest colors
+@lru_cache
+def brightness(c: RGB):
+    return c[0]*0.3 + c[1]*0.6 + c[2]*0.1
+
+darkest_color = min(palette, key=brightness)
+lightest_color = max(palette, key=brightness)
+
+@lru_cache
+def distance_sq(a: RGB, b: RGB) -> int:
+    result = 0
+    for i in range(3):
+        result += (a[i] - b[i])**2
+    return result
+
+def nearest_color(color: RGB, colors: list[RGB]) -> RGB:
+    return min(colors, key=lambda c: distance_sq(c, color))
+
+def choose_color_pair_kmeans(counts: dict[RGB,int]) -> tuple[RGB, RGB]:
+    colors = [k for k, v in counts.items() if v > 0]
+    total_count = sum(counts.values())
+    if len(colors) == 0:
+        raise ValueError('no colors present')
+    elif len(colors) == 1:
+        return colors[0], None
+
+    def weighted_avg(colors: list[RGB]) -> RGB:
+        result = [0, 0, 0]
+        for color in colors:
+            weight = counts[color]
+            for i in range(3):
+                result[i] += weight*color[i]
+        return tuple(round(x/total_count) for x in result)
+
+    # k-means to choose the two colors
+    old_result = (None, None)
+    result = (darkest_color, lightest_color)
+    while True:
+        lists = ([], [])
+        for color in colors:
+            d_dark = distance_sq(color, result[0])
+            d_light = distance_sq(color, result[1])
+            if d_dark < d_light:
+                lists[0].append(color)
+            else:
+                lists[1].append(color)
+        old_result = result
+        result = tuple(weighted_avg(cs) for cs in lists)
+        if old_result == result:
+            break
+    return tuple(nearest_color(x, palette) for x in result)
+
+
+def choose_color_pair_luminance(counts: dict[RGB, int]) -> tuple[RGB,RGB]:
+    # Calculate brightness threshold
+    mean_brightness = 0
+    total_count = 0
+    for rgb, count in counts.items():
+        mean_brightness += brightness(rgb)*count
+        total_count += count
+    mean_brightness = mean_brightness/total_count
+
+    # Separate colors into two lists
+    lists = ([], [])
+    for color in counts.keys():
+        if brightness(color) <= mean_brightness:
+            lists[0].append(color)
+        else:
+            lists[1].append(color)
+    print(lists)
+
+    # Find least-bad approximations for each color
+    def choose_color(cs: list[RGB]) -> RGB:
+        best = None
+        best_cost = math.inf
+        for rgb in palette:
+            cost = sum(distance_sq(rgb, c) * counts[c] for c in cs)
+            if cost < best_cost:
+                best = rgb
+                best_cost = cost
+        return best
+    dark, light = tuple(choose_color(l) for l in lists)
+    if dark == light:
+        light = None
+    print(dark, light)
+    return dark, light
+
+
+def two_tone(img: Image, xy1: tuple[int,int] = None, xy2: tuple[int,int] = None) -> None:
+    # choose the two colors
+    x1, y1 = xy1 or (0, 0)
+    x2, y2 = xy2 or (img.width, img.height)
+    x1 = max(x1, 0)
+    y1 = max(y1, 0)
+    x2 = min(x2, img.width)
+    y2 = min(y2, img.height)
+    counts = {}
+    for x in range(x1, x2):
+        for y in range(y1, y2):
+            color = palette[img.getpixel((x,y))]
+            if color not in counts:
+                counts[color] = 0
+            counts[color] += 1
+    dark, light = choose_color_pair_luminance(counts)
+    if light == None:
+        # Only one color, *ought* to be a no-op
+        light = dark
+    colors = [dark, light]
+
+    # change each pixel's color to nearest
+    for x in range(x1, x2):
+        for y in range(y1, y2):
+            color = palette[img.getpixel((x, y))]
+            replacement = nearest_color(color, colors)
+            img.putpixel((x, y), palette_indexes[replacement])
+
+
+BLOCK_SIZE = 10
+blocked = img.copy()
+for x in range(0, img.width, BLOCK_SIZE):
+    for y in range(0, img.height, BLOCK_SIZE):
+        two_tone(blocked, (x, y), (x + BLOCK_SIZE, y + BLOCK_SIZE))