Add decay to ANS predictor and further optimize it

The overall result: some compression, but less than PX9. It's still
possible this could beat PX9 with some tweaks.

pico-8-secret-palette.hex

@@ -0,0 +1,32 @@
+000000
+1d2b53
+7e2553
+008751
+ab5236
+5f574f
+c2c3c7
+fff1e8
+ff004d
+ffa300
+ffec27
+00e436
+29adff
+83769c
+ff77a8
+ffccaa
+291814
+111d35
+422136
+125359
+742f29
+49333b
+a28879
+f3ef7d
+be1250
+ff6c24
+a8e72e
+00b543
+065ab5
+754665
+ff6e59
+ff9d81

python/ans-coding.ipynb

@@ -0,0 +1,405 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from bitarray import bitarray\n",
+    "from itertools import cycle\n",
+    "from util import walk_corpus\n",
+    "from PIL import Image\n",
+    "import numpy as np\n",
+    "from collections import Counter"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "123\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "bitarray('1000')"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "BitSequence = list[int] | bitarray\n",
+    "\n",
+    "def bans_encode_bignum(bits: BitSequence, zero_probabilities: int | list[int], n=8):\n",
+    "    if type(zero_probabilities) is int:\n",
+    "        zero_probabilities = [zero_probabilities] * len(bits)\n",
+    "    else:\n",
+    "        assert len(bits) == len(zero_probabilities)\n",
+    "    probability_total = 1 << n\n",
+    "    assert all(0 < p < probability_total for p in zero_probabilities)\n",
+    "\n",
+    "    result = probability_total - 1\n",
+    "    for bit, p0 in zip(reversed(bits), reversed(zero_probabilities)):\n",
+    "        p = p0 if bit == 0 else probability_total - p0\n",
+    "        quotient, remainder = divmod(result, p)\n",
+    "        symbol_idx = remainder + (0 if bit == 0 else p0)\n",
+    "        result = (quotient << n) + symbol_idx\n",
+    "    return result\n",
+    "\n",
+    "def bans_decode_bignum(value: int, zero_probabilities: int | list[int], n=8):\n",
+    "    total_probability = 1 << n\n",
+    "    if type(zero_probabilities) is int:\n",
+    "        assert 0 < zero_probabilities < total_probability\n",
+    "        zero_probabilities = cycle([zero_probabilities])\n",
+    "    else:\n",
+    "        assert all(0 < p < total_probability for p in zero_probabilities)\n",
+    "\n",
+    "    result = bitarray()\n",
+    "    mask = total_probability - 1\n",
+    "    while value > mask:\n",
+    "        symbol_idx = value & mask\n",
+    "        p0 = next(zero_probabilities)\n",
+    "        bit = 0 if symbol_idx < p0 else 1\n",
+    "        result.append(bit)\n",
+    "        p = p0 if bit == 0 else total_probability - p0\n",
+    "        remainder = symbol_idx - (0 if bit == 0 else p0)\n",
+    "        value = (value >> n) * p + remainder\n",
+    "    return result\n",
+    "\n",
+    "t = bans_encode_bignum([1, 0, 0, 0], 11, n=4)\n",
+    "print(t)\n",
+    "bans_decode_bignum(t, 11, n=4)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "bytearray(b'\\x99\\x15\\xfdu\\xc0\\xa1\\xf4!\\xccj\\x80') 11\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "bitarray('100100100100100100100100100100100100100100100100100100100100100100100100100100100100100100')"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def bans_encode(bits: BitSequence, zero_probabilities: int | list[int], n=8):\n",
+    "    if type(zero_probabilities) is int:\n",
+    "        zero_probabilities = [zero_probabilities] * len(bits)\n",
+    "    else:\n",
+    "        assert len(bits) == len(zero_probabilities)\n",
+    "    total_probability = 1 << n\n",
+    "    assert all(0 < p < total_probability for p in zero_probabilities)\n",
+    "\n",
+    "    result = bytearray()\n",
+    "    value = total_probability - 1\n",
+    "    def flush_byte():\n",
+    "        nonlocal result\n",
+    "        nonlocal value\n",
+    "        result.append(value & 0xff)\n",
+    "        value >>= 8\n",
+    "    def encode(x, p0, bit):\n",
+    "        p = p0 if bit == 0 else total_probability - p0\n",
+    "        quotient, remainder = divmod(x, p)\n",
+    "        symbol_idx = remainder + (0 if bit == 0 else p0)\n",
+    "        return (quotient << n) | symbol_idx\n",
+    "\n",
+    "    for bit, p0 in zip(reversed(bits), reversed(zero_probabilities)):\n",
+    "        next_value = encode(value, p0, bit)\n",
+    "        if next_value.bit_length() > n + 8:\n",
+    "            flush_byte()\n",
+    "            next_value = encode(value, p0, bit)\n",
+    "        value = next_value\n",
+    "    while value > 0:\n",
+    "        flush_byte()\n",
+    "    result.reverse()\n",
+    "    return result\n",
+    "\n",
+    "\n",
+    "def bans_decode(data: bytes | bytearray, zero_probabilities: int | list[int], n=8):\n",
+    "    total_probability = 1 << n\n",
+    "    if type(zero_probabilities) is int:\n",
+    "        assert 0 < zero_probabilities < total_probability\n",
+    "        zero_probabilities = cycle([zero_probabilities])\n",
+    "    else:\n",
+    "        assert all(0 < p < total_probability for p in zero_probabilities)\n",
+    "\n",
+    "    result = bitarray()\n",
+    "    value = 0\n",
+    "    data = iter(data)\n",
+    "    def read_byte():\n",
+    "        nonlocal value, data\n",
+    "        try:\n",
+    "            value = (value << 8) | next(data)\n",
+    "            return True\n",
+    "        except StopIteration:\n",
+    "            return False\n",
+    "\n",
+    "    mask = total_probability - 1\n",
+    "    while value.bit_length() <= n:\n",
+    "        if not read_byte():\n",
+    "            break\n",
+    "    while True:\n",
+    "        if value.bit_length() <= n:\n",
+    "            if not read_byte():\n",
+    "                break\n",
+    "\n",
+    "        symbol_idx = value & mask\n",
+    "        p0 = next(zero_probabilities)\n",
+    "        bit = 0 if symbol_idx < p0 else 1\n",
+    "        result.append(bit)\n",
+    "        p = p0 if bit == 0 else total_probability - p0\n",
+    "        remainder = symbol_idx - (0 if bit == 0 else p0)\n",
+    "        value = (value >> n) * p + remainder\n",
+    "    return result\n",
+    "\n",
+    "t = bans_encode([1, 0, 0] * 30, 21, n=5)\n",
+    "print(t, len(t))\n",
+    "bans_decode(t, 21, n=5)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'-0.91, -0.90, -0.89, -0.88, -0.86, -0.83, -0.80, -0.75, -0.67, -0.50, 0.00, 0.50, 0.67, 0.75, 0.80, 0.83, 0.86, 0.88, 0.89, 0.90, 0.91'"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def sigmoid(x):\n",
+    "    return x / (1.0 + abs(x))\n",
+    "\n",
+    "\", \".join(f\"{sigmoid(x):.2f}\" for x in range(-10, 11))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'-40, -36, -32, -28, -24, -20, -16, -12, -8, -4, 0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80'"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def isigmoid(x):\n",
+    "    # Trying to make something that returns an i8. Might not be 100% true.\n",
+    "    # Precision of input is on my mind\n",
+    "    return ((x << 3) // (1 + abs(x >> 4)))\n",
+    "\n",
+    "\", \".join(f\"{isigmoid(x)}\" for x in range(-10, 11))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "What does the context look like for binary-encoded data?\n",
+    "\n",
+    "We can divide it up into discrete pieces, each of which are counted separately, and are combined via the sigmoid function:\n",
+    "\n",
+    "- Partial encoding of the current pixel: for example, if the pixel is \"1011\" we will pass through contexts \"\", \"1\", \"10\", and \"101\".\n",
+    "- Full encoding of NW neighbor (omitted if it doesn't exist)\n",
+    "- Full encoding of N neighbor (ditto for all of these)\n",
+    "- Full encoding of NE neighbor\n",
+    "- Full encoding of W neighbor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "p=PosixPath('../corpus/quantized/floodedcaves_0.png'), bytes=1743 (percent=21.3%, bits per pixel=0.85107421875)\n",
+      "p=PosixPath('../corpus/quantized/age_of_ants-title.png'), bytes=3222 (percent=39.3%, bits per pixel=1.5732421875)\n",
+      "p=PosixPath('../corpus/quantized/build a jetpack_0.png'), bytes=1223 (percent=14.9%, bits per pixel=0.59716796875)\n",
+      "p=PosixPath('../corpus/quantized/shinkansen 1 3_1.png'), bytes=1178 (percent=14.4%, bits per pixel=0.5751953125)\n",
+      "p=PosixPath('../corpus/quantized/pico8_build_a_jetpack-1.png'), bytes=1116 (percent=13.6%, bits per pixel=0.544921875)\n",
+      "p=PosixPath('../corpus/quantized/iso_ray-09_000.png'), bytes=3614 (percent=44.1%, bits per pixel=1.7646484375)\n",
+      "p=PosixPath('../corpus/quantized/age of ants_2.png'), bytes=2190 (percent=26.7%, bits per pixel=1.0693359375)\n",
+      "p=PosixPath('../corpus/quantized/pico8_mot_raymarch3-0.png'), bytes=4233 (percent=51.7%, bits per pixel=2.06689453125)\n",
+      "p=PosixPath('../corpus/quantized/age of ants_5.png'), bytes=1988 (percent=24.3%, bits per pixel=0.970703125)\n",
+      "p=PosixPath('../corpus/quantized/donsol for pico-8  v1 8 4 _2.png'), bytes=1006 (percent=12.3%, bits per pixel=0.4912109375)\n",
+      "p=PosixPath('../corpus/quantized/pico8_witchcrafttd-6.png'), bytes=1272 (percent=15.5%, bits per pixel=0.62109375)\n",
+      "p=PosixPath('../corpus/quantized/pico8_bunnysurvivor-9.png'), bytes=1189 (percent=14.5%, bits per pixel=0.58056640625)\n",
+      "p=PosixPath('../corpus/quantized/pico8_rotslimepires_1_1-0.png'), bytes=3144 (percent=38.4%, bits per pixel=1.53515625)\n",
+      "p=PosixPath('../corpus/quantized/pico8_donsol8_v1-14.png'), bytes=1236 (percent=15.1%, bits per pixel=0.603515625)\n",
+      "p=PosixPath('../corpus/quantized/pico8_ppwr-5.png'), bytes=1742 (percent=21.3%, bits per pixel=0.8505859375)\n",
+      "p=PosixPath('../corpus/quantized/pico8_px9-9.png'), bytes=1064 (percent=13.0%, bits per pixel=0.51953125)\n",
+      "p=PosixPath('../corpus/quantized/build a jetpack_5.png'), bytes=1117 (percent=13.6%, bits per pixel=0.54541015625)\n",
+      "p=PosixPath('../corpus/quantized/pico20068.png'), bytes=1464 (percent=17.9%, bits per pixel=0.71484375)\n",
+      "p=PosixPath('../corpus/quantized/storming the grandmothership_1.png'), bytes=2231 (percent=27.2%, bits per pixel=1.08935546875)\n",
+      "p=PosixPath('../corpus/quantized/build a jetpack_9.png'), bytes=1315 (percent=16.1%, bits per pixel=0.64208984375)\n",
+      "p=PosixPath('../corpus/quantized/hersheys_train_line_0.png'), bytes=2613 (percent=31.9%, bits per pixel=1.27587890625)\n",
+      "Total bytes: 39900\n"
+     ]
+    }
+   ],
+   "source": [
+    "class Predictor:\n",
+    "    def __init__(self):\n",
+    "        self.counts = [Counter() for _ in range(4)]\n",
+    "        gain = 1.27\n",
+    "        weights = [30, 100] * 2\n",
+    "        self.weights = [x/sum(weights) * gain for x in weights]\n",
+    "        self.decay = 0.9\n",
+    "\n",
+    "    def _keys(self, neighbors, partial_pixel):\n",
+    "        assert len(neighbors) == 4\n",
+    "        return tuple(neighbor | (partial_pixel << 4) for neighbor in neighbors)\n",
+    "\n",
+    "    def predict(self, contexts):\n",
+    "        \"\"\"Returns probability that the next bit is zero, out of 256\"\"\"\n",
+    "        total = 0.0\n",
+    "        for counter, key, weight in zip(self.counts, contexts, self.weights):\n",
+    "            total += weight * counter[key]\n",
+    "        prob_float = (sigmoid(total) + 1.0) / 2.0\n",
+    "        prob_int = round(prob_float * 256.0)\n",
+    "        return min(max(1, prob_int), 255)\n",
+    "\n",
+    "    def update(self, contexts, bit):\n",
+    "        \"\"\"Updates the model with the actual bit\"\"\"\n",
+    "        # Note for the future: code might be simpler if we used P(bit=1) everywhere\n",
+    "        delta = -int(bit) * 2 + 1\n",
+    "        for counter, key in zip(self.counts, contexts):\n",
+    "            counter[key] = counter[key] * self.decay + delta\n",
+    "\n",
+    "\n",
+    "def bitwise_encode(img: np.array):\n",
+    "    height, width = img.shape\n",
+    "    assert img.dtype == np.uint8\n",
+    "\n",
+    "    bits = bitarray()\n",
+    "    zero_probabilities = [] # values are 1 through 255\n",
+    "    pred = Predictor()\n",
+    "    num_missed_predictions = 0\n",
+    "    for y, row in enumerate(img):\n",
+    "        for x, val in enumerate(row):\n",
+    "            neighbors = [31] * 4 # 11111 means \"no neighbor\", 0xxxx means \"neighbor is that color\"\n",
+    "            if y > 0:\n",
+    "                if x > 0:\n",
+    "                    neighbors[0] = img[y - 1, x - 1] # NW\n",
+    "                neighbors[1] = img[y - 1, x] # N\n",
+    "                if x < width - 1:\n",
+    "                    neighbors[2] = img[y - 1, x + 1] # NE\n",
+    "            if x > 0:\n",
+    "                neighbors[3] = img[y, x - 1] # W\n",
+    "            # We'll store context as 1(5 neighbor bits)(0-3 partial pixel bits)\n",
+    "            for i in range(4):\n",
+    "                neighbors[i] = int(neighbors[i]) | (1 << 5)\n",
+    "            for i in range(3, -1, -1):\n",
+    "                bit = int((val >> i) & 1)\n",
+    "                bits.append(bit)\n",
+    "                prediction = pred.predict(neighbors)\n",
+    "                missed_prediction = (prediction < 16 and bit == 0) or (prediction > 255 - 16 and bit == 1)\n",
+    "                if missed_prediction:\n",
+    "                    num_missed_predictions += 1\n",
+    "                zero_probabilities.append(prediction)\n",
+    "                pred.update(neighbors, bit)\n",
+    "                for i in range(4):\n",
+    "                    neighbors[i] = (neighbors[i] << 1) | bit\n",
+    "    #print(f\"Missed predictions: {num_missed_predictions} out of {len(bits)}\")\n",
+    "\n",
+    "    return bans_encode(bits, zero_probabilities, n=8)\n",
+    "\n",
+    "\n",
+    "def corpus_bitwise_encode():\n",
+    "    byte_size = 0\n",
+    "    for p in walk_corpus():\n",
+    "        img = np.array(Image.open(p))\n",
+    "        t = len(bitwise_encode(img))\n",
+    "        byte_size += t\n",
+    "        print(f\"{p=}, bytes={t} (percent={t/8192*100:.3}%, bits per pixel={t*8/128**2})\")\n",
+    "    print(\"Total bytes:\", byte_size)\n",
+    "\n",
+    "corpus_bitwise_encode()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Things learned so far:\n",
+    "\n",
+    "- ANS is promising!\n",
+    "- Diagonal neighbors currently count for 30% relative to orthogonal neighbors.\n",
+    "- It's important to bound counts (e.g., with a decay) so they don't go out of control.\n",
+    "    - It looks like we had been confidently wrong (1-15 or 240-255) about 3-4% of the time?\n",
+    "    - Assuming a wrong guess costs 6 bits and 2500 wrong guesses/image: this costs ~2k per image or ~40k for the corpus\n",
+    "    - We shaved off about half that (~20k) by adding decay and fine-tuning it\n",
+    "- If we don't need the precision, we could reduce probability granularities.\n",
+    "- Compression bit-by-bit (no RLE, etc) gets us 23.2% compression on average. This sometimes beats PX9!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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/context-prediction.ipynb

@@ -0,0 +1,449 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "import numpy as np\n",
+    "from collections import Counter\n",
+    "import heapq\n",
+    "from pathlib import Path\n",
+    "import statistics\n",
+    "from bitarray import bitarray\n",
+    "import math"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def transform_context(img_np: np.ndarray):\n",
+    "    h, w = img_np.shape\n",
+    "    front = [0] * (h + 1 + w + 1)\n",
+    "    frequency_ctx = {None: Counter()}\n",
+    "    order_ctx = {None: list(range(16))}\n",
+    "\n",
+    "    result = np.zeros((h, w), dtype=np.uint8)\n",
+    "    for y in range(h):\n",
+    "        for x in range(w):\n",
+    "            # Get current pixel value and context\n",
+    "            corner_idx = h - y + x\n",
+    "            context = tuple(front[corner_idx - 1:corner_idx + 3])\n",
+    "            assert len(context) == 4, f\"bad context {context} ({x=},{y=},{corner_idx=})\"\n",
+    "            pixel = int(img_np[y, x])\n",
+    "            front[corner_idx] = pixel\n",
+    "\n",
+    "            # Get appropriate frequency Counter and order list\n",
+    "            if context in frequency_ctx:\n",
+    "                frequency = frequency_ctx[context]\n",
+    "                order = order_ctx[context]\n",
+    "            else:\n",
+    "                frequency = frequency_ctx[None] # separate context for misses\n",
+    "                order = order_ctx[None]\n",
+    "\n",
+    "            # Eager M2F transform\n",
+    "            idx = order.index(pixel)\n",
+    "            result[y, x] = idx\n",
+    "            frequency[pixel] += 1\n",
+    "            # if idx > 0:\n",
+    "            #     print(f\"encoding {pixel=}, surprised at {x=},{y=}: {context=},{frequency=},{order=}\")\n",
+    "            while idx > 0 and frequency[order[idx - 1]] <= frequency[pixel]:\n",
+    "                order[idx], order[idx - 1] = order[idx - 1], order[idx]\n",
+    "                idx -= 1\n",
+    "\n",
+    "            # Populate context state based on a copy of the default\n",
+    "            if context not in frequency_ctx:\n",
+    "                frequency_ctx[context] = frequency.copy()\n",
+    "                order_ctx[context] = order.copy()\n",
+    "    return result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
+       "       [1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1],\n",
+       "       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
+       "       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
+       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
+       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
+       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8)"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def checkerboard(h, w):\n",
+    "    result = np.zeros((h, w), dtype=np.uint8)\n",
+    "    for y in range(h):\n",
+    "        for x in range(w):\n",
+    "            result[y,x] = (x ^ y) & 1\n",
+    "    return result\n",
+    "\n",
+    "transform_context(checkerboard(7,11))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def huffman(counts: dict):\n",
+    "    h = [(v, (k,)) for k, v in counts.items()]\n",
+    "    prefixes = {k: \"\" for k in counts.keys()}\n",
+    "    heapq.heapify(h)\n",
+    "    while(len(h) > 1):\n",
+    "        (c1, v1), (c2, v2) = heapq.heappop(h), heapq.heappop(h)\n",
+    "        heapq.heappush(h, (c1 + c2, v1 + v2))\n",
+    "        for v in v1:\n",
+    "            prefixes[v] = \"0\" + prefixes[v]\n",
+    "        for v in v2:\n",
+    "            prefixes[v] = \"1\" + prefixes[v]\n",
+    "    keys = sorted(prefixes.keys())\n",
+    "    return {k: prefixes[k] for k in keys}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Average length of 0th code = 1.00 (quartiles [1.0, 1.0, 1.0])\n",
+      "Average length of 1th code = 2.67 (quartiles [2.0, 3.0, 3.0])\n",
+      "Average length of 2th code = 3.52 (quartiles [3.0, 4.0, 4.0])\n",
+      "Average length of 3th code = 4.10 (quartiles [4.0, 4.0, 4.0])\n",
+      "Average length of 4th code = 4.67 (quartiles [4.0, 5.0, 5.0])\n",
+      "Average length of 5th code = 5.00 (quartiles [4.5, 5.0, 5.5])\n",
+      "Average length of 6th code = 5.45 (quartiles [5.0, 5.0, 6.0])\n",
+      "Average length of 7th code = 5.95 (quartiles [5.0, 6.0, 6.0])\n",
+      "Average length of 8th code = 6.40 (quartiles [5.25, 6.0, 7.0])\n",
+      "Average length of 9th code = 6.90 (quartiles [6.0, 6.5, 7.0])\n",
+      "Average length of 10th code = 6.95 (quartiles [6.0, 7.0, 8.0])\n",
+      "Average length of 11th code = 7.53 (quartiles [7.0, 7.0, 8.0])\n",
+      "Average length of 12th code = 8.22 (quartiles [7.0, 8.0, 9.0])\n",
+      "Average length of 13th code = 8.62 (quartiles [8.0, 9.0, 9.0])\n",
+      "Average length of 14th code = 8.69 (quartiles [8.0, 8.0, 10.0])\n",
+      "Average length of 15th code = 8.75 (quartiles [8.0, 8.0, 9.75])\n"
+     ]
+    }
+   ],
+   "source": [
+    "img_code_lengths = []\n",
+    "for p in Path('../corpus/quantized/').glob(\"*.*\"):\n",
+    "    img = Image.open(p)\n",
+    "    #np.array(img)\n",
+    "    img_np = np.array(img)\n",
+    "    transformed = transform_context(img_np)\n",
+    "    c = Counter(transformed.reshape(-1))\n",
+    "    code_lengths = [len(x) for x in huffman(c).values()]\n",
+    "    img_code_lengths.append(code_lengths)\n",
+    "\n",
+    "for i in range(16):\n",
+    "    x = []\n",
+    "    for code_lengths in img_code_lengths:\n",
+    "        if i < len(code_lengths):\n",
+    "            x.append(code_lengths[i])\n",
+    "    mean = statistics.mean(x)\n",
+    "    quarts = statistics.quantiles(x)\n",
+    "    print(f\"Average length of {i}th code = {mean:1.2f} (quartiles {quarts})\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(['a', 'b', 'b', 'c', 'c', 'c', 'd', 'd', 'd', 'e', 'e', 'e'], [0, 1, 2])"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def transform_rle(inputs):\n",
+    "    outputs = []\n",
+    "    for val in inputs:\n",
+    "        last_val = outputs[-1][0] if len(outputs) > 0 else None\n",
+    "        if last_val == val:\n",
+    "            outputs[-1][1] += 1\n",
+    "        else:\n",
+    "            outputs.append([val, 1])\n",
+    "    return [tuple(x) for x in outputs]\n",
+    "\n",
+    "def transform_trig_rle(inputs, trigger_length):\n",
+    "    runs = transform_rle(inputs)\n",
+    "    values = []\n",
+    "    additional_run_lengths = []\n",
+    "    for value, run_length in runs:\n",
+    "        num_literals = min(run_length, trigger_length)\n",
+    "        values.extend(value for _ in range(num_literals))\n",
+    "        if run_length >= trigger_length:\n",
+    "            additional_run_lengths.append(run_length - trigger_length)\n",
+    "    return values, additional_run_lengths\n",
+    "\n",
+    "transform_trig_rle(\"abbcccddddeeeee\", 3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Encoding value 0 as 0 (length 1)\n",
+      "Encoding value 1 as 111 (length 3)\n",
+      "Encoding value 2 as 100 (length 3)\n",
+      "Encoding value 3 as 1100 (length 4)\n",
+      "Encoding value 4 as 1010 (length 4)\n",
+      "Encoding value 5 as 11010 (length 5)\n",
+      "Encoding value 6 as 10110 (length 5)\n",
+      "Encoding value 7 as 110111 (length 6)\n",
+      "Encoding value 8 as 101111 (length 6)\n",
+      "Encoding value 9 as 1101101 (length 7)\n",
+      "Encoding value 10 as 1011101 (length 7)\n",
+      "Encoding value 11 as 1011100 (length 7)\n",
+      "Encoding value 12 as 11011000 (length 8)\n",
+      "Encoding value 13 as 110110011 (length 9)\n",
+      "Encoding value 14 as 1101100101 (length 10)\n",
+      "Encoding value 15 as 1101100100 (length 10)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def analyze_huffman(counts, name):\n",
+    "    huffman_map = huffman(counts)\n",
+    "    for k, v in huffman_map.items():\n",
+    "        print(f\"Encoding {name} {k} as {v} (length {len(v)})\")\n",
+    "\n",
+    "value_counter = Counter()\n",
+    "length_counter = Counter()\n",
+    "for p in Path('../corpus/quantized/').glob(\"*.*\"):\n",
+    "    img = Image.open(p)\n",
+    "    img_np = np.array(img)\n",
+    "    context_coded = transform_context(img_np)\n",
+    "    values, lengths = transform_trig_rle(context_coded.reshape(-1), 3)\n",
+    "    value_counter.update(values)\n",
+    "    length_counter.update(lengths)\n",
+    "\n",
+    "analyze_huffman(value_counter, 'value')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "What did I learn about the distribution of values (post trig-RLE transform)?\n",
+    "\n",
+    "The main takeaway is that the codes look very geometric. For example, running trigger length 2 on the corpus and Huffman coding the values results in codes of length 2, 2, 3, 3, 4, 4, and so on. Note that the number of codes of each length is pretty much constant, much like a Golomb/Rice code. And those are optimal/near-optimal for geometric distributions.\n",
+    "\n",
+    "Running trigger length 3 and above seems to give a slightly different distribution: the Huffman codes have lengths *1*, 3, 3, 4, 4, etc. This is basically just a single prefix bit with a possible Rice code attached: the zero value is encoded as \"0\", while subsequent values are encoded as \"1\" + Rice code.\n",
+    "\n",
+    "Keep in mind, this is specific to the RLE _values_. The RLE _lengths_ are a different story."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[(1, 5988), (2, 4715), (3, 4074), (4, 2536), (5, 1396), (6, 815), (7, 542), (8, 16), (9, 11), (10, 10), (11, 8), (12, 2)]\n",
+      "Encoding bucket 1 as 11 (length 2)\n",
+      "Encoding bucket 2 as 01 (length 2)\n",
+      "Encoding bucket 3 as 00 (length 2)\n",
+      "Encoding bucket 4 as 100 (length 3)\n",
+      "Encoding bucket 5 as 1010 (length 4)\n",
+      "Encoding bucket 6 as 10111 (length 5)\n",
+      "Encoding bucket 7 as 101101 (length 6)\n",
+      "Encoding bucket 8 as 10110011 (length 8)\n",
+      "Encoding bucket 9 as 10110010 (length 8)\n",
+      "Encoding bucket 10 as 10110000 (length 8)\n",
+      "Encoding bucket 11 as 101100011 (length 9)\n",
+      "Encoding bucket 12 as 101100010 (length 9)\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Try buckets based on number of bits\n",
+    "bucket_frequencies = Counter()\n",
+    "for length, count in length_counter.items():\n",
+    "    bucket = (length + 1).bit_length()\n",
+    "    bucket_frequencies[bucket] += count\n",
+    "print([(i, bucket_frequencies[i]) for i in range(1, max(bucket_frequencies.keys()) + 1)])\n",
+    "analyze_huffman(bucket_frequencies, 'bucket')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If we use bucket algorithm `length.bit_length()` and tweak the codes slightly to make them monotonically increasing in size, we get an encoding scheme that looks like this:\n",
+    "\n",
+    "```\n",
+    "00  (run 0)\n",
+    "01  (run 1)\n",
+    "100x    (runs 2-3) (implied leading 1 bit)\n",
+    "101xx   (runs 4-7)\n",
+    "110xxx  (runs 8-15)\n",
+    "1110xxxx    (runs 16-31)\n",
+    "11110xxxxx  (runs 32-63)\n",
+    "111110xxxxxx    (runs 64-128)\n",
+    "```\n",
+    "\n",
+    "If we change the bucket algorithm to `(length + 1).bit_length` so there's always a leading 1 bit, we don't have to reorder any codes, and we get something similar:\n",
+    "\n",
+    "```\n",
+    "00  (run 0)\n",
+    "01x (runs 1-2)\n",
+    "10xx    (runs 3-6)\n",
+    "110xxx  (runs 7-14)\n",
+    "1110xxxx    (runs 15-30)\n",
+    "11110xxxxx  (runs 31-62)\n",
+    "111110xxxxx (runs 63-126)\n",
+    "```\n",
+    "\n",
+    "The code lengths aren't exactly the same, but this is reminiscient of Exp-Golomb coding."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "../corpus/quantized/floodedcaves_0.png compressed to 32.4% (2658 bytes)\n",
+      "../corpus/quantized/age_of_ants-title.png compressed to 77.1% (6318 bytes)\n",
+      "../corpus/quantized/build a jetpack_0.png compressed to 20.4% (1668 bytes)\n",
+      "../corpus/quantized/shinkansen 1 3_1.png compressed to 21.3% (1744 bytes)\n",
+      "../corpus/quantized/pico8_build_a_jetpack-1.png compressed to 25% (2050 bytes)\n",
+      "../corpus/quantized/iso_ray-09_000.png compressed to 54.1% (4435 bytes)\n",
+      "../corpus/quantized/age of ants_2.png compressed to 41.4% (3389 bytes)\n",
+      "../corpus/quantized/pico8_mot_raymarch3-0.png compressed to 62.7% (5136 bytes)\n",
+      "../corpus/quantized/age of ants_5.png compressed to 43.6% (3569 bytes)\n",
+      "../corpus/quantized/donsol for pico-8  v1 8 4 _2.png compressed to 15.2% (1243 bytes)\n",
+      "../corpus/quantized/pico8_witchcrafttd-6.png compressed to 25.1% (2057 bytes)\n",
+      "../corpus/quantized/pico8_bunnysurvivor-9.png compressed to 27.2% (2228 bytes)\n",
+      "../corpus/quantized/pico8_rotslimepires_1_1-0.png compressed to 57% (4670 bytes)\n",
+      "../corpus/quantized/pico8_donsol8_v1-14.png compressed to 18.3% (1499 bytes)\n",
+      "../corpus/quantized/pico8_ppwr-5.png compressed to 41.3% (3381 bytes)\n",
+      "../corpus/quantized/pico8_px9-9.png compressed to 23.3% (1905 bytes)\n",
+      "../corpus/quantized/build a jetpack_5.png compressed to 20.3% (1665 bytes)\n",
+      "../corpus/quantized/pico20068.png compressed to 34.1% (2790 bytes)\n",
+      "../corpus/quantized/storming the grandmothership_1.png compressed to 37.5% (3070 bytes)\n",
+      "../corpus/quantized/build a jetpack_9.png compressed to 20.4% (1675 bytes)\n",
+      "../corpus/quantized/hersheys_train_line_0.png compressed to 52% (4259 bytes)\n"
+     ]
+    }
+   ],
+   "source": [
+    "def append_exp_golomb(dest: bitarray, val: int):\n",
+    "    val += 1\n",
+    "    dest.extend(1 for _ in range(val.bit_length() - 1))\n",
+    "    dest.append(0)\n",
+    "    for i in range(val.bit_length() - 2, -1, -1):\n",
+    "        dest.append((val >> i) & 1)\n",
+    "\n",
+    "def append_rice1(dest: bitarray, val: int, frequentZero=False):\n",
+    "    if frequentZero:\n",
+    "        if val == 0:\n",
+    "            dest.append(0)\n",
+    "            return\n",
+    "        else:\n",
+    "            dest.append(1)\n",
+    "            val -= 1\n",
+    "    dest.extend(1 for _ in range(val >> 1))\n",
+    "    dest.append(0)\n",
+    "    dest.append(val & 1)\n",
+    "\n",
+    "def encode_mixed_trig_rle(values, xlengths, trigger_length):\n",
+    "    result = bitarray()\n",
+    "    runs = transform_rle(values)\n",
+    "    xlengths_iter = iter(xlengths)\n",
+    "    for value, length in runs:\n",
+    "        for _ in range(length):\n",
+    "            append_rice1(result, value, True)\n",
+    "        assert length <= trigger_length\n",
+    "        if length == trigger_length:\n",
+    "            append_exp_golomb(result, next(xlengths_iter))\n",
+    "    assert next(xlengths_iter, None) == None\n",
+    "    return result\n",
+    "\n",
+    "def encode_image(img, trigger_length=3):\n",
+    "    img_np = np.array(img)\n",
+    "    ctx_transformed = transform_context(img_np)\n",
+    "    values, xlengths = transform_trig_rle(ctx_transformed.reshape(-1), trigger_length)\n",
+    "    return encode_mixed_trig_rle(values, xlengths, trigger_length)\n",
+    "\n",
+    "for p in Path('../corpus/quantized/').glob('*.*'):\n",
+    "    img = Image.open(p)\n",
+    "    raw_size_bytes = math.ceil(img.width * img.height * 4 / 8)\n",
+    "    encoded = encode_image(img, trigger_length=2)\n",
+    "    encoded_size_bytes = math.ceil(len(encoded) / 8)\n",
+    "    encoded_percent = encoded_size_bytes * 100 / raw_size_bytes\n",
+    "    print(f\"{p} compressed to {encoded_percent:.3g}% ({encoded_size_bytes} bytes)\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "These percents are pretty much universally worse than PX9. The only image that benefited was the pattern-dithered `iso_ray-09_000`. Possible reasons for the poor performance:\n",
+    "- Too much context causes dilution?\n",
+    "- Maybe frequency-based M2F is suboptimal somehow?"
+   ]
+  }
+ ],
+ "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/palette.ipynb

@@ -0,0 +1,176 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "from pathlib import Path\n",
+    "import numpy as np\n",
+    "import math"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Function to read hex color codes from a file and convert them to RGB tuples\n",
+    "def load_palette(file_path):\n",
+    "    palette = []\n",
+    "    with open(file_path, 'r') as f:\n",
+    "        for line in f:\n",
+    "            hex_color = line.strip()\n",
+    "            # Convert hex to RGB tuple\n",
+    "            rgb_color = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))\n",
+    "            palette.append(rgb_color)\n",
+    "    return palette\n",
+    "\n",
+    "pico8_palette = load_palette('../pico-8-secret-palette.hex')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_image(path: Path) -> Image.Image:\n",
+    "    img = Image.open(path).convert('RGB')\n",
+    "    w, h = img.size\n",
+    "    if w > 128 or h > 128:\n",
+    "        if w > h:\n",
+    "            w, h = 128, int(round(h * (128 / w)))\n",
+    "        else:\n",
+    "            w, h = int(round(w * (128 / h))), 128\n",
+    "        img = img.resize((w, h))\n",
+    "    return img\n",
+    "\n",
+    "img = load_image('../peacock.jpg')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def convert_to_indexed(img_np: np.array, palette: list) -> Image.Image:\n",
+    "    pixels = img_np.reshape(-1, 3).astype(np.float32)\n",
+    "    pal = np.array(palette, dtype=np.float32)\n",
+    "    distances = np.sum((pixels[:, None, :] - pal[None, :, :]) ** 2, axis=2)\n",
+    "    closest_idxs = np.argmin(distances, axis=1).astype(np.uint8).reshape(img_np.shape[1], -1)\n",
+    "    result = Image.fromarray(closest_idxs, mode=\"P\")\n",
+    "    result.putpalette(np.array(palette).reshape(-1).tolist())\n",
+    "    return result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def squared_quantization_error(img_np: np.array, palette: np.array):\n",
+    "    pixels = img_np.reshape(-1, 3).astype(np.float32)\n",
+    "    pal = np.array(palette, dtype=np.float32)\n",
+    "    sq_distances = np.sum((pixels[:, None, :] - pal[None, :, :]) ** 2, axis=2)\n",
+    "    min_errs = np.min(sq_distances, axis=1)\n",
+    "    return np.sum(min_errs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAIAAAACABAMAAAAxEHz4AAAAMFBMVEWiiHkGWrUSU1n/nYFJMztfV08RHTWrUjbCw8f/8eh0LykAh1EdK1ODdpwprf//zKq2QJPYAAAL/0lEQVR4AZ2Yz2skxxXHO9PbjMAR2m5wMOzKqy6vOocZkQnDgnMIA4ukg/BBkdW7YZFJ0th9yRx0EQN7EjGGzkHIR89JF5uATpbIhsDEYDexLpnDImXDHizE/gNBZC7ag+l8v/Wjf8yMdp08mOnqqvc+9d6rH13dVpJ8bFWlvhNuFzVOWJJDJWufL+diAdAu1GUpfA0gN2YBgAkXKoBS/2GoHCgD1ghIJjxYL2omAZ+PAfaTpFXos/QaDyYA2XgMUwE7dOURY6gADq39DICKC04YrtMRJbSj9Pg3DZDsZ1fVcXCgbcwZj5SHO7zQgcMiBbyzkmzcBShOAUgPaGEAa/IGgCQZSyMAD7dzF9gx5GFvB//SRnmgzOnBVEDhAq0hD3tw4RGtdA7LgOpMcJS6RvCOAsCO8kAB1iqAdnkyO9Tv9YwPvIOgAi7QqhqBDKHdbuUh67SjQ1XlSHPHqkvAI2O/RpQU5qBir8fNAOTtOmB1hvAoB3yhM3potav9W6pLuKyd4r0CoBDqFK4dogwaRKsVF6pBHhoAXSDAYW0YygysHX7BsiQUlqrksAVSABwFkLUGoG6mArQeQtjWbAXAvxR4sPJIl+WgjHtg2goAYoCSoxtgb1RUDP8rYHM5DH91rBmM4YcCchfgQLPxCoCeBqUcwJS9ONoIl/h3T+S0klNh3AOtmM9Emq5LpRIgenPb4sxmFicAyoWH+SBMAax2b5Gp1sYkwGFXpQ1Bh1CKYTUWC7TrcSpNAqRmGSCHEXoES9mKv/mSdvWd6QBqqghmxbk3H26sU7tIo+Ws+rdkzXUAECTAcSEefktSXbpgcUiaQqyzqn6NB1z8aD71DcB1qS7TKWflauSf8b6+MzUHbIJsuAAI132bPtyUVcgDr85q7J+1WKoCbFblMsvuAfDl1cvrWbgPgNyFkISiwcbWAtEVsPeO3NtwwE2JWCgULSdqAkDNXhkgzXNEDQB27msXFkuAG1Hg3pQbGbJY1EtAhn9ZdaTsXV8gDZDmzUJzM47TRZsx1MMCYBPQyUZqi3wDZl7AMRTz54jFF/MFwImjpoC+HIa8GhUawKqG6/uu5/mnw1Scwgnh+0+1qt2+EafxN2fjAG7PkJHyYHAME1ckWydxKlIBhF+4cBKLyD2jIiaCxuJCcyko177aOIbV8714ZfO9OD2J4IOcTNwbHCECfwrA1vZMzqx1Cq3nyXcr2IffE+KDPyAdT0GWIjwhGhdQr3qQu9CyZqyZfwXxUvLJ38PVlc0PkufJJ76XJ8Ha8ILoaO56QHtmFwDxUbKHZdn9EIAkcSJPrR+6sCHiqMHxmuJBBxOhlcyG68+imeR00QnjD8Mw2bUeiKZg/PJEuByn/vwUAGPogJvszoaz58/PkYbzEI+B9WTXEUfB0rpaS21nOY7cL0fSg8elUbAsuw3pJHDg6WyylyJTT/BY23ye7J6mR4IAJatRszEPe4QwBZD8MXn2dCb57BkBx3fDECn4ZxwLIVeyJDzAqN5UgPJaUB5cfYaUQfa2ooYQg+Ol8yTZjeM4ir7W/VutG8C9BQfgwQSg0/8M1jMwerbRaAwax/DfogHmkrWuEK3NII48mzOxspzlXHyX9rsAIHEbceMr3uyuduGM8GYcy26RsRoN3CdwQAE48ZSA2YYlDz2Q2XBlS/yapQdBVxyBMAc1qFjOciDE0scoWb1DHnEMAI0dOAyZgdl3cXdJ4OpsRiJuDAhoAdBpt+8tI6fckuCBBKCeYsP+ChPFpk/O072TqLF39354NxZBjPEQ0TuXNhqzdgspFSnWggKMjAs2thMcu3F05mx0Vh5Ejb9Bc4hV1RVMwsLIbllZJ2vdD4CYMwBMSfQIacP2iuaUUecGjAciigPRREmciHQhszJrLssu6dMi5iw9WLMMgA6UZfQAgb6AahR0YQEPlrLLzqUFgO+6qVoLEgCHjQNlADqKmHx2H4kuAO7+qJNZCMHHTvsWA7Xq9AAuyACqDoBlwYdbDfatAXPZXHZpZ7XI/bOYk3YG0OL4VrqXN5ermIIBMs5xEMK/k0HJatfE/NNZKCAJCgBfKArQ7x+oAv9BaBxHQHAshHuRjVBn1bAZUCTgECEoQId1WZ8iS/LvHjMQd9MUF+HeYd3B7IY7jxetMmCUOyDty4QthHACawbiXvQPOln2n/v+QvJcA7AaZUkH8FIDCh8uAxEFiIPzwP0NwjvIfgRnli404LEEMDdSckBB2MAkjBoDN0gH3lv9/mX/svYMj5V96DMHGqDMdQaUF6bq+8FxYzDwb3mR597p9z/9xeUyRjUHqBCMskqhAhzoystjMWiIIxejiBz0D+x3V0XXnd9HczEKWjcrIiBE14oGn5OYy673M1TeWCHqr2z8YYCr498PxOAIM9pnCBbnVuOrfQXAVFZJ5C1EOW/+Vd3lcXQciIEn8Gi60z+wPkwHjcWZl2zEEhgDvDS2uCpzTLyg+4+/CHErwMPxov+p81Gj8cR6KZuxmHhKM5q4FoCDvPYKUygQDY+HlTvZXg1zYOalmu7YDx5XAYqQ5b2TcgmAF0SuhzPPRTb7xuDLfYTKFuRwAsD6Mfk+iAIPBwK44M1lV3v8YKGEKeBR19xfc/15EysBhx3kwNuv6MCB1wPa7VozRf84siEHFXtuQ9jVX+0BVulpUwReA7ugJ9QuYChwgINgAGY5mWZcucbbLdFMPcbges2FUiM3RHlY14DRBECat1v3AOCBjzG802l3DINDoAFKE/+mTffNesuuBUOR+sgBlsIZajRhRAfGAaVGaLIHiDhJA7jguj8RPNwWwuZeyBwUdXmJlkpaIm6epHDA9b3IzhVQYDtOBxOAljI0/7WgKwE49IvfGq/IoYJ8acMhy2aFkpYx1FekAE/XVCAB4vbXY43qtVGd0myFHNewrK2mSIcYBf925K2PNdd3+N6pAGNNxW1TiC4cwDj6+sWtaMNHmdcCnGbahAcYBd8Thakq/RBAbZg21VPNF+WXJkngZ6HpIdTzvnBARhD0AOeTvFYXCEAM+A6wrd9U2YB3ViW8cWju8+3L8ydzWACkBQ2Uvfputo27mhhiCkvxJyIwXyAsfqfLP1vV5Tc71hAH5zXAC9ZZUZG6+oRh7VA9/2yFIt4y1D3ejfCuhXWApXS3YqtuqPqIOSjZMwc7sJfAGuy1eL+dAuDnOQXYLrUSB2HNBjzAZkopXndKqhYAGIXtcpUZBlmJvRA7qQRUX96NBWKAB9dKDYc5aY2/xdY0NQDgwbWCCHwDWPh/AJj+OeBsajdI+CtCcPg0lPLCdc+meoBheAWglgPeBmCqB5gCrwBslELw3i8D5O4lXar3Hl+fRJyosYqUeAswUgxsXlp43yu//isF84+PMMYcz5RvjZW84uEwQqGFWXM9oJbPAaTA4xuSEf1s4aMPM//aEGpqDGH9ggDGoMTYyzOS1XtsXB6/buQAEILoTbzQjMqPv0wdUYp1PA7wFcD9N9Lo4XXBLnqGrRJkoZ5vBPIBg7QoabcMADG43kn3/vsdY1ZcFSBHMERtjpfL0kJwvfibb+3CLi9pgCbYBLRIYKFmdgI44OLpNjzLzYrCFABdsBUgNUsBE6IZD5cKu7xkANoF2uVyqvcCpsD1T+Khf5Hb5YUqwM6NWcBmQFP3bRKGcYwvYbldXqgCKvYtdS6C8YsXCuCe7+eGpmAA28ycPQ6AsRYvPRmieGHs8qsB0H4McE9PI4nw3NtD33XfyQ1N4RUAvRKUC757O0XpprEz1xHGDDNRfYIe8+C07IGv9qb5kbHUVwKwFrZlBGMA0YTPRkzx43anjBjJCPKJrCagzmRLMGgtwn+Bko+NEa0lAid+vbIatTVQv8TpFEZKVDQ+kqAVzAWu64egCqK0Hp1mkQNPzShM7Hljqa7KvhSCweCKlwT0qEUDUAHDVkmLAVRCKNp+PNQATmRYYkvg31nFXB0MpnuwRYBPG4hPYwm4WXSBEvp//KdrPHgAgO/TmqbSWJZho6cNACiXAbjt9eryv7fdxPHW86WN6y2qK/9/yvZt7UcdgMPcA7bkssPzsem3yYVE8YRYpIq2HwshN2ZhB685+Bguzcw5TwJ8tpYAvP0vdYRg9k+FnAgAAAAASUVORK5CYII=",
+      "text/plain": [
+       "<PIL.Image.Image image mode=P size=128x128>"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def pick_palette_indexes(img, n, palette):\n",
+    "    img_np = np.array(img)\n",
+    "    result = list(range(n))\n",
+    "    num_steps = 0\n",
+    "    pal_idx = 0\n",
+    "    while num_steps < n:\n",
+    "        num_steps += 1\n",
+    "        old_val = result[pal_idx]\n",
+    "        result[pal_idx] = None\n",
+    "        best_err = math.inf\n",
+    "        best_val = None\n",
+    "        for val in [i for i in range(len(palette)) if i not in result]:\n",
+    "            result[pal_idx] = val\n",
+    "            err = squared_quantization_error(img_np, [palette[i] for i in result])\n",
+    "            if err < best_err:\n",
+    "                best_err = err\n",
+    "                best_val = val\n",
+    "        result[pal_idx] = best_val\n",
+    "        if best_val != old_val:\n",
+    "            num_steps = 0\n",
+    "        pal_idx = (pal_idx + 1) % len(result)\n",
+    "    return result\n",
+    "\n",
+    "pal_idxs = pick_palette_indexes(img, 16, pico8_palette)\n",
+    "\n",
+    "\n",
+    "pal = [pico8_palette[i] for i in pal_idxs]\n",
+    "convert_to_indexed(np.array(img), pal)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAIAAAACABAMAAAAxEHz4AAAAMFBMVEUAAAAdK1N+JVMAh1GrUjZfV0/Cw8f/8ej/AE3/owD/7CcA5DYprf+Ddpz/d6j/zKruozp8AAAK7UlEQVR4AbWZMW/jOBbHHRyuNyG4S3CfgSAylYzFAnapxQixtvMEg8VoK4+7cbXIZwsEXOUtjNjuFDVzpe3idjyVrfv/H0lJlp1ksMA9TCSK5PvxvcdHUvJ0UpHcy3IwGAyTgX/Ml8PkdelYwMRrEJAkg6h6fl09Sf5PgFFlAcx5XTrplE5MKpOtC1UQfgRQktAGuCgsK8DwBUM607INQNcR7BgMcl68XlXwFe4OF8p9wwKoAEBNJ67fCLUXpZOWYoJ3AWro57XtExVHL1uQShh9FK3KqCa4YUc07JIwDzgPDgBFdrwEeMEEl0gOsDwDEEcBciiF9sUDZB6pbwdiQcRr8aGtK88eIC4sq24jligOMGL5bQB7OQ0WKQ2AK7YwzgKZxaaGlHlx/UdV6SJA9Jfs31CQR//cUqsfxQI7h6IwdE0jqz4YrOq+F0sEWH3J/GpArx+9CRDreVmKjh/FA/Lc17xw75wChr6bA0SXAHe3vhfuLUDbhYsAZV4BDF1bbcGq0dsWVXCLZlfdtsDVjxwADl4AXBHgetYAOwsu9V4FvFdG2sWGBqA5D6eA8akZ76+0bSehCciXOc1gtQfYII5PHYkRBPhmF0qn2sBRRbHV4EByld8iD8ZjaSGYMlgpY1CzFHc71Sm2nko3KrpNMQmNybQ2pgFY8dxUKviAzmJsBViX5SFfT1wgiHkykFCbzAQ3wh4mSZ4AtlKBed44EypAgd0ZJ8SEURC5M0FotJnPH81cQQ0KSEveV6pzXew3MAHICjAlAHJYWv3I0AFlggwmGE3FPB8LBS5cF2W5cQAbRXjghbYNIqiF4WMYzjVCFuob6nMyePmHNuj/Gf1qC3i+eCHgzsx1MDdZkF1nYWZ0T/Tt5edAmX+X5X7TBEy9Nu5wLjLmEf/mWaBgAMRGQQCrWHXULRRqAKJzClg+ieOMQpDnT8YEqm9HT6J8/UlhHouy3DITmEhL/J24kDMCUDZ9ztgqQBADgETW03LRBiDlWoChYf6YfhqjL5S1CYO+I+STrtJXBhYcxQKJ+akLh/ypH5tAB+vJVS+BuYCEag4AM3m80qjKGgDMexuQ5wtl1Dpd9JKkh9L4I1LKW5DDIUULSmcBAIgMnitJJ+uFUuP0k0pUEo/TdfoJXjhAOrkzmJopen9lEDFnBORFpV4in9OF7qUTGHz1z2SMnX8VaCVpTGtj0w0Mu9cA1OYVYJ9iyazEgVWi4ARWSAzACguJZ8gSgQ36DQDyjtwKgE4YbJ0vYHUSI2rpBDOp+mPEEL6mKwBME0B9EApHwCsf9NMJu91CPVmnkwUBqB0xWrnCBJ0DImdC+nu6wjBfkEJIIjiAECw6V8qGQAIZI7P+tC5gMUkQaxP2XxAyyCRWmDmDVYhyVwHQE2VYcCT73TnAmjD7Au01lFYcxxgU09UVc6mH3WCSL6N1GXftNB65p3WwkUkMBhGD4A0QAkKFO6JxhdUUmF9X+e4IwlcYZ27ReWsBXJXiAo2yw+OKSCax+pWlGLNmsDUVuYwxwSwEpj/FcnYArAlI9GwB9JRqC6VukAEpVr/GiBoL+LCM2Oeb6l6psKgtOIoP0Q5te8xAiuskyTGNAf4S3LD2AFD3+2IAtfKgINwPXAzElMHgFzSV+ymvkMMKDiD7vBDwriwONOEQdLudE4D4sKXenhcrC4zCZYelzEmUmUvL5a4s/8LyVAxiZQEBv6DlVOA7NjRsf6KNS788lAcYesDs4gHyVaYRhU2eiwEnhG+YbAjmXARbwBTtS7oAy+pEotJ2u93xfiqIf4DoBUHAeUC5KIvyUJS/IUcZgoYF2/Px2X6I6TqEFxAKVn4uP3J7ZLF2QZ7sZTZ7qJ8O7zEuCYikpgWQ/WRxLQFAeetjwAYrM4p/YM4w4gJAqWDDAzakPj60KMczgOg3CbFoW4gpxLr/YlWkdgs9A3x3gNoGmIAAcN7gzD0AD+VTovVN4SzApiold6kANeFOQxcAhEC/m832s8PTChVT0dhyV24CvAG4+2pmDUUjrYrZ7Mvv+/iuApy50AA8OMJetOFCYAh4WKcxUoqLGdIG1B6QJF3KErbzBYOAe1TGPUZD8pCrkDtSLZcAe0MAF2NgfprNuFUGgTPgDMBxa7HgPZzn6Jpb0uxhjJPS3K6/28bXLPB2HaBOQjdA/s/+WPUC6jv/vg5fcuHB65ewAMId3pifSpwwegx9235ctgDldzpQ+vAJBBZoSQIVIpXxtgP/fYCPSIOTIFbDNgoEwH8NE/QUKymf1o0MwdsA8SATLxq6pBy5Ib0J+A0AHYQIpD2Q6/H5kvYDgCcB4IVNa78LOAZOA4TgLQt4BmBTM9mjtudpbQEOtLcBWyxn5qDJQq3/VSujdMTB9hbguN0uBID9HXNZnAC2OM9eBWCnhsQCoAFKnQCOPJIRgnYM+HtMeRRdXjbQgxi8uuvetGkBDbCAXaOap3Kafq70t5wDAoJ5oG7QUVpwP25pAD1IOs0zYW9/U6sBC9FHGJGPN6jm2Z8/o7DxBhAAsVZLcx7J7Rm1uQ0B9MNHpT9s0AC7pZkFMUA+PNGZjfLajQaIPOX5xhtwjaXwDPWIjRFaea8BrrvUssUKarH/igR4dReAtLQAqLNAp9e4RbEDIIr6dpMvfZsdylngKy/cI6ePbRhvBw2A7fs2YOgA3A30uBGi1wH1r4kLC5gjj4NbeN42MqF0qi991wzDRq4YE3Ct51mmFQx4ATD03Z0W89MWrQeBwXenUpjEOoiNrvjsa8vQ/yD6M8bHjo6v37npSRKed4UL7cr6OQaAWTifZ+qDpFrdZktw9zUAtRH/OQD6+SJA9sQ2tXpmCHCg0ABze/wbgJgACCJgxnu7Eiu6K7zqAvQrQIGdYIN5qMQBEAQfxKg9y0M5U/FqgBiExf4zlzfnktuBcMioAaxaOqy9LbACICFcCDMLAKMWm1feAkGephrfZISAGOj7dFer2j3w84Y2D60LkQM0bMDHIRZxRg9CI/uh3bV2Vh1Xty1KDDygtiGK8I6NH08YAe7IlVqjYI+WywB+/qvreUZ1pAHe8BqKVZGEUxfEE7vpLZTEXwgAcE8/FwCGYoELgdwQB97vmgD8erc7V+fb+mBQASTKFWmDD+dQhsdFI6OKSwB+K3nAM79ZjjZFnrf4rMGZ3ACo+1cBEfOLPZivgvqm6UIoUcR+pHrFGwDRogmQHfp+JCDkNGTcj1pHs4dVLiBs9ogkgPrlAm7PBYC1qFTXppJX9PcaMABg56t5XyAE3Msg9KCr+kWz2ZVrQOQsqDrF9r2GiCxT3a4yedVWFxqAZ+uCb8O7VcCVRAuwI0JM8w3TdzsB7Hwt7werbwFwAUlxKRNeBnxrAmhAcJ3dN0ew5a1PJPz6cPrtumgAsuBRKfwW13rPJKJK5QsARiDLwhBpQMD8MfzzzAKuRpvKg0EriDaGshvgXX/+iM0hnE/bBHrgAQNkENtxPe74I0dAE5DKyAIAEAUzL1oAMaACwIRa8G4kADeHAghC96lWYcSAChDV6tvtfwRgPcDgAnA/4FXqOKvgQO0Co1DJAhaLAzaL8JIHkZ8Qj76P7MknAC4HyJKVFiCJiLWoPaDAfoHW0wNENhQqQapNOXYxCHkmBDqc04BQ9W23+jpEsQmoWjwAQRTA4zUAcGtZ9bCFFwEcEEHAUsBrek+r7FqjSgWjHwS4NwMAMvzXAtLYWIDqtQB8vOSCB8CDMMPW9giGlQuA/wHleVRDWkJkqQAAAABJRU5ErkJggg==",
+      "text/plain": [
+       "<PIL.Image.Image image mode=P size=128x128>"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "convert_to_indexed(np.array(img), pico8_palette[:16])"
+   ]
+  }
+ ],
+ "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/pyproject.toml

@@ -0,0 +1,20 @@
+[tool.poetry]
+name = "color-cell"
+version = "0.1.0"
+description = ""
+authors = ["Chris Mounce <christophermounce@gmail.com>"]
+readme = "README.md"
+package-mode = false
+
+[tool.poetry.dependencies]
+python = "^3.12"
+jupyter = "^1.1.1"
+matplotlib = "^3.9.2"
+pillow = "^11.0.0"
+numpy = "^2.1.2"
+bitarray = "^3.0.0"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

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,83 @@
+from typing import Generator
+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)
+
+def walk_corpus(quantized=True) -> Generator[Path, None, None]:
+    if quantized:
+        return Path('../corpus/quantized/').glob('*.png')
+    else:
+        return Path('../corpus/').glob('*.*')

rust/Cargo.lock

@@ -241,7 +241,11 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "clap",
+ "empfindung",
  "image",
+ "lab",
+ "ordered-float",
+ "rand",
 ]
 
 [[package]]
@@ -302,6 +306,16 @@ version = "1.13.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0"
 
+[[package]]
+name = "empfindung"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8af095ec32b4a626b56a2e188515e4e5480d82b02e5d8a830fdb61934226b0df"
+dependencies = [
+ "lab",
+ "rgb",
+]
+
 [[package]]
 name = "equivalent"
 version = "1.0.1"
@@ -485,6 +499,12 @@ version = "0.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f5d4a7da358eff58addd2877a45865158f0d78c911d43a5784ceb7bbf52833b0"
 
+[[package]]
+name = "lab"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf36173d4167ed999940f804952e6b08197cae5ad5d572eb4db150ce8ad5d58f"
+
 [[package]]
 name = "lebe"
 version = "0.5.2"
@@ -642,6 +662,15 @@ version = "1.19.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92"
 
+[[package]]
+name = "ordered-float"
+version = "4.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4a91171844676f8c7990ce64959210cd2eaef32c2612c50f9fae9f8aaa6065a6"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "paste"
 version = "1.0.15"

rust/Cargo.toml

@@ -3,10 +3,21 @@ name = "color-cell"
 version = "0.1.0"
 edition = "2021"
 
+[profile.dev.package.zune-jpeg]
+opt-level = 3
+
+[profile.dev.package.lab]
+opt-level = 3
+
+[profile.dev.package.empfindung]
+opt-level = 3
+
 [dependencies]
 anyhow = "1.0.86"
 clap = { version = "4.5.15", features = ["derive"] }
+empfindung = "0.2.6"
 image = "0.25.2"
+lab = "0.11.0"
+ordered-float = "4.2.2"
+rand = "0.8.5"
 
-[profile.dev.package.zune-jpeg]
-opt-level = 3

rust/src/color_dist.rs

@@ -0,0 +1,87 @@
+use empfindung::cie00;
+use lab::Lab;
+
+use crate::indexed::redmean_sq;
+
+
+fn ciede2000(a: [u8;3], b: [u8;3]) -> f32 {
+    cie00::diff(Lab::from_rgb(&a), Lab::from_rgb(&b)) // / 100.0
+}
+
+fn ciede2000_sq(a: [u8;3], b: [u8;3]) -> f32 {
+    let diff = ciede2000(a, b);
+    diff * diff
+}
+
+fn redmean(a: [u8;3], b: [u8;3]) -> f32 {
+    redmean_sq(a, b).sqrt()
+}
+
+#[cfg(test)]
+mod test {
+    use core::f32;
+    use std::f32::INFINITY;
+
+    use rand::prelude::*;
+
+    use super::ciede2000;
+
+    fn cost(params: [f32; 3]) -> f32 {
+        let decode = |x: u32| -> [u8; 3] {
+            let r = (x >> 8) as u8;
+            let g = ((x >> 4) & 15) as u8;
+            let b = (x & 15) as u8;
+            [r | r<<4, g | g<<4, b | b<<4]
+        };
+        let mut result = 0.0;
+        // TODO: precompute
+        for color1 in 0..4096 {
+            if color1 % 73 != 0 {
+                continue;
+            }
+            let rgb1 = decode(color1);
+            for color2 in (color1 + 1)..4096 {
+                if color2 % 73 != 0 {
+                    continue;
+                }
+                let rgb2 = decode(color2);
+                let true_diff = ciede2000(rgb1, rgb2);
+                let mut approx_diff = 0.0;
+                for i in 0..3 {
+                    let delta = rgb1[i] as f32 - rgb2[i] as f32;
+                    approx_diff += params[i] * delta * delta;
+                }
+                approx_diff = approx_diff.sqrt();
+                let err = true_diff - approx_diff;
+                result += err * err;
+            }
+        }
+        result
+    }
+
+
+    #[test]
+    fn tinker() {
+        let mut params = [1.0; 3];
+        let mut best_params = [1.0; 3];
+        let mut best_cost = f32::INFINITY;
+        let mut scale = 1.0;
+        for i in 0..5000 {
+            if i > 4990 {
+                dbg!((i, best_params, best_cost));
+            }
+            for c in 0..3 {
+                params[c] += scale * thread_rng().gen_range(-0.01..0.01f32);
+            }
+            let new_cost = cost(params);
+            if new_cost < best_cost {
+                best_cost = new_cost;
+                best_params = params;
+            } else {
+                params = best_params;
+                scale *= 0.99999;
+            }
+        }
+        dbg!((best_params, best_cost));
+    }
+}

rust/src/encode.rs

@@ -0,0 +1,148 @@
+use image::{Rgb, RgbImage};
+
+use crate::indexed::redmean_sq;
+
+pub type Color = [u8; 3]; // TODO: Convert to indexed
+
+#[derive(Debug)]
+pub struct EncodeTree {
+    pub error: f32,
+    pub bytes: f32,
+    pub strategy: EncodeStrategy,
+}
+
+#[derive(Debug)]
+pub enum EncodeStrategy {
+    Solid(Color),
+    // TODO: add Bitmap([Color; 2]),
+    QuadSplit(Box<[EncodeTree; 4]>),
+}
+
+impl EncodeTree {
+    pub fn apply(&self, img: &mut RgbImage, region: Region) {
+        match &self.strategy {
+            EncodeStrategy::Solid(color) => {
+                for y in region.y..region.y2() {
+                    for x in region.x..region.x2() {
+                        img.put_pixel(x, y, Rgb(*color))
+                    }
+                }
+            }
+            EncodeStrategy::QuadSplit(subtrees) => {
+                let width = region.width / 2;
+                let height = region.height / 2;
+                let coords = [
+                    [region.x, region.y],
+                    [region.x + width, region.y],
+                    [region.x, region.y + height],
+                    [region.x + width, region.y + height],
+                ];
+                for (tree, &[x, y]) in subtrees.iter().zip(coords.iter()) {
+                    tree.apply(
+                        img,
+                        Region {
+                            x,
+                            y,
+                            width,
+                            height,
+                        },
+                    );
+                }
+            }
+        }
+    }
+}
+
+#[derive(Clone, Copy)]
+pub struct Region {
+    pub x: u32,
+    pub y: u32,
+    pub width: u32,
+    pub height: u32,
+}
+
+impl Region {
+    pub fn x2(&self) -> u32 {
+        self.x + self.width
+    }
+
+    pub fn y2(&self) -> u32 {
+        self.y + self.height
+    }
+
+    pub fn area(&self) -> u32 {
+        self.width * self.height
+    }
+}
+
+pub fn choose_encoding(img: &RgbImage, region: Region, max_error: f32) -> EncodeTree {
+    // Attempt to quantize as solid color
+    let info = analyze_solid(img, region);
+    if info.error <= max_error {
+        return EncodeTree {
+            error: info.error,
+            bytes: 1.0,
+            strategy: EncodeStrategy::Solid(info.color),
+        };
+    }
+
+    // Error is too high, so we have to subdivide
+    let mut subinfo = vec![];
+    let width = region.width / 2;
+    let height = region.height / 2;
+    for y in [region.y, region.y + height] {
+        for x in [region.x, region.x + width] {
+            subinfo.push(choose_encoding(
+                img,
+                Region {
+                    x,
+                    y,
+                    width,
+                    height,
+                },
+                max_error / 4.0,
+            ));
+        }
+    }
+    EncodeTree {
+        error: subinfo.iter().fold(0.0, |acc, info| acc + info.error),
+        bytes: 1.0 + subinfo.iter().fold(0.0, |acc, info| acc + info.bytes),
+        strategy: EncodeStrategy::QuadSplit(Box::new(subinfo.try_into().unwrap())),
+    }
+}
+
+struct SolidInfo {
+    color: Color,
+    error: f32,
+}
+
+fn analyze_solid(img: &RgbImage, region: Region) -> SolidInfo {
+    // Calculate average color
+    let mut avg = [0.0, 0.0, 0.0];
+    for y in region.y..region.y2() {
+        for x in region.x..region.x2() {
+            let pixel = img.get_pixel(x, y);
+            for (i, component) in avg.iter_mut().enumerate() {
+                *component += (pixel.0[i] as f32).powf(2.2);
+            }
+        }
+    }
+    for component in &mut avg {
+        *component /= region.area() as f32;
+    }
+    let avg = avg.map(|x| x.powf(0.4545).round() as u8);
+
+    // Measure squared error
+    let mut err = 0.0;
+    for y in region.y..region.y2() {
+        for x in region.x..region.x2() {
+            let pixel = img.get_pixel(x, y);
+            err += redmean_sq(avg, pixel.0);
+        }
+    }
+
+    SolidInfo {
+        color: avg,
+        error: err,
+    }
+}

rust/src/indexed.rs

@@ -0,0 +1,227 @@
+use std::collections::BinaryHeap;
+
+use empfindung::cie00;
+use image::RgbImage;
+use lab::Lab;
+use ordered_float::NotNan;
+use rand::Rng;
+
+pub fn euclid_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
+    let mut result = 0.0;
+    let weights = [1.641, 5.764, 2.088];
+    for (i, weight) in weights.iter().enumerate() {
+        let delta = x[i] as f32 - y[i] as f32;
+        result += weight * delta * delta;
+    }
+    result
+}
+
+pub fn redmean_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
+    let mut ds: [f32; 3] = [0.0; 3];
+    for (i, d) in ds.iter_mut().enumerate() {
+        let delta = x[i] as f32 - y[i] as f32;
+        *d = delta * delta;
+    }
+    let rm = (x[0] as f32 + y[0] as f32) / 2.0;
+    let result = (2.0 + rm / 256.0) * ds[0] + 4.0 * ds[1] + (2.0 + (255.0 - rm) / 256.0) * ds[2];
+    // normalize value to the range 0.0 - 1.0
+    result / 584971.0
+}
+
+fn rgb_to_ycocg(x: [u8; 3]) -> [f32; 3] {
+    let [r, g, b] = x.map(|x| x as f32);
+    let co = r - b;
+    let tmp = b + co/2.0;
+    let cg = g - tmp;
+    let y = tmp + cg/2.0;
+    [y, co, cg]
+}
+
+pub fn ycocg_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
+    let x = rgb_to_ycocg(x);
+    let y = rgb_to_ycocg(y);
+    let mut result = 0.0;
+    let weights = [15.03, 1.982, 7.702];
+    for (i, weight) in weights.iter().enumerate() {
+        let delta = (x[i] as f32 - y[i] as f32)/255.0;
+        result += weight * delta * delta;
+    }
+    result
+}
+
+type Rgb = [u8; 3];
+
+const PICO8_PALETTE: [Rgb; 32] = [
+    // 0-15: standard palette
+    [0x00, 0x00, 0x00],
+    [0x1d, 0x2b, 0x53],
+    [0x7e, 0x25, 0x53],
+    [0x00, 0x87, 0x51],
+    [0xab, 0x52, 0x36],
+    [0x5f, 0x57, 0x4f],
+    [0xc2, 0xc3, 0xc7],
+    [0xff, 0xf1, 0xe8],
+    [0xff, 0x00, 0x4d],
+    [0xff, 0xa3, 0x00],
+    [0xff, 0xec, 0x27],
+    [0x00, 0xe4, 0x36],
+    [0x29, 0xad, 0xff],
+    [0x83, 0x76, 0x9c],
+    [0xff, 0x77, 0xa8],
+    [0xff, 0xcc, 0xaa],
+    // 16-31: secret palette
+    // In PICO-8, these indexes start at 128
+    [0x29, 0x18, 0x14],
+    [0x11, 0x1d, 0x35],
+    [0x42, 0x21, 0x36],
+    [0x12, 0x53, 0x59],
+    [0x74, 0x2f, 0x29],
+    [0x49, 0x33, 0x3b],
+    [0xa2, 0x88, 0x79],
+    [0xf3, 0xef, 0x7d],
+    [0xbe, 0x12, 0x50],
+    [0xff, 0x6c, 0x24],
+    [0xa8, 0xe7, 0x2e],
+    [0x00, 0xb5, 0x43],
+    [0x06, 0x5a, 0xb5],
+    [0x75, 0x46, 0x65],
+    [0xff, 0x6e, 0x59],
+    [0xff, 0x9d, 0x81],
+];
+
+#[derive(Clone, Copy)]
+pub struct PicoColor(u8);
+
+impl PicoColor {
+    pub fn to_rgb(self) -> Rgb {
+        let idx = ((self.0 & 0x80) >> 3) | (self.0 & 0x0F);
+        PICO8_PALETTE[idx as usize]
+    }
+}
+
+pub struct PicoImage {
+    pub palette: [PicoColor; 16],
+    pub width: u32,
+    pub height: u32,
+    pub data: Vec<u8>,
+}
+
+impl PicoImage {
+    pub fn new(width: u32, height: u32) -> Self {
+        let mut palette = [PicoColor(0); 16];
+        for (i, color) in palette.iter_mut().enumerate() {
+            *color = PicoColor(i as u8);
+        }
+        Self {
+            palette,
+            width,
+            height,
+            data: vec![0; (width * height) as usize],
+        }
+    }
+
+    pub fn to_indexed(img: &RgbImage) -> Self {
+        // Brute-force nearest-color conversion
+        // TODO begin using this everywhere
+        let mut result = Self::new(img.width(), img.height());
+        let mut discrepancies: Vec<(f32, (Rgb, Rgb), f32)> = vec![];
+        for y in 0..img.height() {
+            let offset = (y * img.width()) as usize;
+            for x in 0..img.width() {
+                let original = img.get_pixel(x, y).0;
+                let best = (0..16)
+                    .min_by_key(|&i| {
+                        let candidate = PicoColor(i).to_rgb();
+                        // let diff = redmean_sq(original, candidate);
+                        // let diff = euclid_sq(original, candidate);
+                        let diff = ycocg_sq(original, candidate);
+                        // let diff =
+                        //     cie00::diff(Lab::from_rgb(&original), Lab::from_rgb(&candidate));
+                        NotNan::new(diff).unwrap()
+                    })
+                    .unwrap();
+
+                // Log places where our color formula disagrees with CIEDE2000
+                let best_rgb = PicoColor(best).to_rgb();
+                let ciede2000 = cie00::diff(Lab::from_rgb(&original), Lab::from_rgb(&best_rgb));
+                let diff = 24.36 * ycocg_sq(original, best_rgb).sqrt();
+                let discrepancy = (
+                    ciede2000,
+                    (original, best_rgb),
+                    diff
+                );
+                if discrepancies.len() < 10 {
+                    discrepancies.push(discrepancy);
+                } else {
+                    let idx = rand::thread_rng().gen_range(0..discrepancies.len());
+                    if discrepancies[idx].0 < discrepancy.0 {
+                        discrepancies[idx] = discrepancy;
+                    }
+                }
+
+                // if x % 43 == 22 && y % 43 == 22 {
+                //     let best = PicoColor(best).to_rgb();
+                //     let ciede2000 = cie00::diff(Lab::from_rgb(&original), Lab::from_rgb(&best));
+                //     let redmean = redmean_sq(original, best);
+                //     dbg!((original, best));
+                //     dbg!((redmean, ciede2000 * ciede2000 / 100_00.0));
+                // }
+                result.data[offset + x as usize] = best;
+            }
+        }
+        for d in discrepancies {
+            dbg!(d);
+        }
+        result
+    }
+
+    pub fn get(&self, x: u32, y: u32) -> u8 {
+        self.data[(y * self.width + x) as usize]
+    }
+
+    pub fn to_rgb(&self) -> RgbImage {
+        let mut result = RgbImage::new(self.width, self.height);
+        for y in 0..self.height {
+            for x in 0..self.width {
+                let index = self.get(x, y);
+                let color = self.palette[index as usize];
+                result.put_pixel(x, y, image::Rgb(color.to_rgb()));
+            }
+        }
+        result
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_to_rgb() {
+        let to_rgb = |i| PicoColor(i).to_rgb();
+
+        // Regular palette
+        assert_eq!(to_rgb(0), [0, 0, 0]);
+        assert_eq!(to_rgb(1), PICO8_PALETTE[1]);
+        assert_eq!(to_rgb(2), to_rgb(0x12));
+        assert_eq!(to_rgb(3), to_rgb(0x73));
+
+        // Secret palette
+        assert_eq!(to_rgb(128), PICO8_PALETTE[16]);
+        assert_eq!(to_rgb(128 + 15), PICO8_PALETTE[31]);
+        assert_eq!(to_rgb(128 + 16), PICO8_PALETTE[16]);
+    }
+
+    #[test]
+    fn test_quantization_round_trip() {
+        for i in 0..16 {
+            let mut img = RgbImage::new(1, 1);
+            let input_color = PICO8_PALETTE[i];
+            img.put_pixel(0, 0, image::Rgb(input_color));
+            let indexed = PicoImage::to_indexed(&img);
+            let img = indexed.to_rgb();
+            let output_color = img.get_pixel(0, 0).0;
+            assert_eq!(input_color, output_color);
+        }
+    }
+}

rust/src/main.rs

@@ -1,6 +1,12 @@
+mod color_dist;
+mod encode;
+mod indexed;
+
 use anyhow::{bail, Result};
 use clap::Parser;
-use image::{DynamicImage, Rgb, RgbImage};
+use encode::{choose_encoding, Region};
+use image::{DynamicImage, RgbImage};
+use indexed::PicoImage;
 
 #[derive(Debug, Parser)]
 #[command()]
@@ -29,64 +35,33 @@ fn main() -> Result<()> {
             output.put_pixel(x, y, *pixel);
         }
     }
-    quadtree_quant(
+
-        &mut output,
+    // Lossy compression
-        args.max_err.unwrap_or(2000.0) * 128.0 * 128.0,
+    let error_scale_factor = 0.2; // Adjusted by feel. Q=0 is unrecognizable but not blank.
-        [0, 0],
+    let error_per_pixel = (1.0 - args.max_err.unwrap_or(80.0) / 100.0) * error_scale_factor;
-        128,
+    let sq_error_per_pixel = error_per_pixel * error_per_pixel;
-    );
+    quadtree_quant(&mut output, sq_error_per_pixel * 128.0 * 128.0);
+
+    // Convert to indexed
+    let indexed = PicoImage::to_indexed(&output);
+    output = indexed.to_rgb();
+
     output.save(&args.output)?;
     println!("Wrote {}", &args.output);
     Ok(())
 }
 
-fn quadtree_quant(img: &mut RgbImage, max_err: f32, pos: [u32; 2], size: u32) {
+fn quadtree_quant(img: &mut RgbImage, max_err: f32) {
-    // Calculate average color
+    let region = Region {
-    let mut avg = [0.0, 0.0, 0.0];
+        x: 0,
-    for y in pos[1]..pos[1] + size {
+        y: 0,
-        for x in pos[0]..pos[0] + size {
+        width: img.width(),
-            let pixel = img.get_pixel(x, y);
+        height: img.height(),
-            for (i, component) in avg.iter_mut().enumerate() {
+    };
-                *component += (pixel.0[i] as f32).powf(2.2);
+    let tree = choose_encoding(img, region, max_err);
-            }
+    println!(
-        }
+        "error = {}, predicted size = {} bytes",
-    }
+        tree.error, tree.bytes
-    for component in &mut avg {
+    );
-        *component /= (size * size) as f32;
+    tree.apply(img, region);
-    }
-    let avg = avg.map(|x| x.powf(0.4545).round() as u8);
-
-    // Measure squared error
-    let mut err = 0.0;
-    for y in pos[1]..pos[1] + size {
-        for x in pos[0]..pos[0] + size {
-            let pixel = img.get_pixel(x, y);
-            err += redmean_sq(avg, pixel.0);
-        }
-    }
-    if err > max_err {
-        let size = size / 2;
-        let max_err = max_err / 4.0;
-        let [x, y] = pos;
-        quadtree_quant(img, max_err, [x, y], size);
-        quadtree_quant(img, max_err, [x + size, y], size);
-        quadtree_quant(img, max_err, [x, y + size], size);
-        quadtree_quant(img, max_err, [x + size, y + size], size);
-    } else {
-        for y in pos[1]..pos[1] + size {
-            for x in pos[0]..pos[0] + size {
-                img.put_pixel(x, y, Rgb(avg));
-            }
-        }
-    }
-}
-
-fn redmean_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
-    let mut ds: [f32; 3] = [0.0, 0.0, 0.0];
-    for i in 0..2 {
-        ds[i] = x[i] as f32 - y[i] as f32;
-        ds[i] *= ds[i];
-    }
-    let rm = (x[0] as f32 + y[0] as f32) / 2.0;
-    (2.0 + rm / 256.0) * ds[0] + 4.0 * ds[1] + (2.0 + (255.0 - rm) / 256.0) * ds[2]
 }