Pre-quantize the corpus

python/quantize.ipynb

@@ -0,0 +1,61 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "from util import load_image_128, quantize_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Processing ../corpus/iso_ray-09_000.png...\n"
+     ]
+    }
+   ],
+   "source": [
+    "corpus_dir = Path('../corpus/')\n",
+    "quantized_dir = corpus_dir / 'quantized'\n",
+    "\n",
+    "for src_path in Path('../corpus/').glob('*.*'):\n",
+    "    dest_path = quantized_dir / src_path.name\n",
+    "    if not dest_path.exists():\n",
+    "        print(f\"Processing {src_path}...\")\n",
+    "        img = load_image_128(src_path)\n",
+    "        img = quantize_image(img)\n",
+    "        img.save(dest_path)\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

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)