Pre-quantize the corpus

python/util.py

@@ -0,0 +1,76 @@
+from PIL import Image
+from pathlib import Path
+import numpy as np
+import math
+
+def load_palette(path: Path) -> np.ndarray:
+    palette = []
+    with open(path, 'r') as f:
+        for line in f:
+            hex_color = line.strip()
+            # Convert hex to RGB tuple
+            rgb_color = [int(hex_color[i:i+2], 16) for i in (0, 2, 4)]
+            palette.append(rgb_color)
+    return np.array(palette)
+
+pico8_palette = load_palette('../pico-8-secret-palette.hex')
+
+def load_image_128(path: Path) -> Image.Image:
+    img = Image.open(path).convert('RGB')
+    w, h = img.size
+    if w > 128 or h > 128:
+        if w % 128 == 0 and h % 128 == 0:
+            resample = Image.Resampling.NEAREST
+        else:
+            resample = Image.Resampling.BICUBIC
+        if w > h:
+            w, h = 128, int(round(h * (128 / w)))
+        else:
+            w, h = int(round(w * (128 / h))), 128
+        img = img.resize((w, h), resample=resample)
+    return img
+
+
+def quantize_image(img: Image.Image) -> Image.Image:
+    pal_idxs = choose_subpalette_indexes(img, 16, pico8_palette)
+    pal = np.array([pico8_palette[i] for i in pal_idxs]).astype(np.float32)
+    img_np = np.array(img)
+    pixels = img_np.reshape(-1, 3).astype(np.float32)
+    distances = np.sum((pixels[:, None, :] - pal[None, :, :]) ** 2, axis=2)
+    closest_idxs = np.argmin(distances, axis=1).astype(np.uint8).reshape(img_np.shape[1], -1)
+    result = Image.fromarray(closest_idxs, mode="P")
+    result.putpalette(pal.astype(np.uint8).reshape(-1).tolist())
+    return result
+
+
+def choose_subpalette_indexes(img: Image.Image, n: int, palette: np.ndarray):
+    assert len(palette.shape) == 2;
+    assert palette.shape[1] == 3
+    img_np = np.array(img)
+    result = list(range(n))
+    num_steps = 0
+    pal_idx = 0
+    while num_steps < n:
+        num_steps += 1
+        old_val = result[pal_idx]
+        result[pal_idx] = None
+        best_err = math.inf
+        best_val = None
+        for val in [i for i in range(len(palette)) if i not in result]:
+            result[pal_idx] = val
+            err = _squared_quantization_error(img_np, [palette[i] for i in result])
+            if err < best_err:
+                best_err = err
+                best_val = val
+        result[pal_idx] = best_val
+        if best_val != old_val:
+            num_steps = 0
+        pal_idx = (pal_idx + 1) % len(result)
+    return result
+
+def _squared_quantization_error(img_np: np.ndarray, palette: np.ndarray):
+    pixels = img_np.reshape(-1, 3).astype(np.float32)
+    pal = np.array(palette, dtype=np.float32)
+    sq_distances = np.sum((pixels[:, None, :] - pal[None, :, :]) ** 2, axis=2)
+    min_errs = np.min(sq_distances, axis=1)
+    return np.sum(min_errs)