diff --git a/python/ans-coding.ipynb b/python/ans-coding.ipynb new file mode 100644 index 0000000..7c626c5 --- /dev/null +++ b/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": 123, + "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": 123, + "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": 128, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "p=PosixPath('../corpus/quantized/floodedcaves_0.png'), bytes=2315 (bits=18520, bits per pixel=1.13037109375)\n", + "p=PosixPath('../corpus/quantized/age_of_ants-title.png'), bytes=4301 (bits=34408, bits per pixel=2.10009765625)\n", + "p=PosixPath('../corpus/quantized/build a jetpack_0.png'), bytes=1804 (bits=14432, bits per pixel=0.880859375)\n", + "p=PosixPath('../corpus/quantized/shinkansen 1 3_1.png'), bytes=1977 (bits=15816, bits per pixel=0.96533203125)\n", + "p=PosixPath('../corpus/quantized/pico8_build_a_jetpack-1.png'), bytes=1722 (bits=13776, bits per pixel=0.8408203125)\n", + "p=PosixPath('../corpus/quantized/iso_ray-09_000.png'), bytes=4857 (bits=38856, bits per pixel=2.37158203125)\n", + "p=PosixPath('../corpus/quantized/age of ants_2.png'), bytes=3170 (bits=25360, bits per pixel=1.5478515625)\n", + "p=PosixPath('../corpus/quantized/pico8_mot_raymarch3-0.png'), bytes=5249 (bits=41992, bits per pixel=2.56298828125)\n", + "p=PosixPath('../corpus/quantized/age of ants_5.png'), bytes=2831 (bits=22648, bits per pixel=1.38232421875)\n", + "p=PosixPath('../corpus/quantized/donsol for pico-8 v1 8 4 _2.png'), bytes=1684 (bits=13472, bits per pixel=0.822265625)\n", + "p=PosixPath('../corpus/quantized/pico8_witchcrafttd-6.png'), bytes=1759 (bits=14072, bits per pixel=0.85888671875)\n", + "p=PosixPath('../corpus/quantized/pico8_bunnysurvivor-9.png'), bytes=1738 (bits=13904, bits per pixel=0.8486328125)\n", + "p=PosixPath('../corpus/quantized/pico8_rotslimepires_1_1-0.png'), bytes=5030 (bits=40240, bits per pixel=2.4560546875)\n", + "p=PosixPath('../corpus/quantized/pico8_donsol8_v1-14.png'), bytes=2474 (bits=19792, bits per pixel=1.2080078125)\n", + "p=PosixPath('../corpus/quantized/pico8_ppwr-5.png'), bytes=3169 (bits=25352, bits per pixel=1.54736328125)\n", + "p=PosixPath('../corpus/quantized/pico8_px9-9.png'), bytes=1699 (bits=13592, bits per pixel=0.82958984375)\n", + "p=PosixPath('../corpus/quantized/build a jetpack_5.png'), bytes=1677 (bits=13416, bits per pixel=0.81884765625)\n", + "p=PosixPath('../corpus/quantized/pico20068.png'), bytes=2601 (bits=20808, bits per pixel=1.27001953125)\n", + "p=PosixPath('../corpus/quantized/storming the grandmothership_1.png'), bytes=3602 (bits=28816, bits per pixel=1.7587890625)\n", + "p=PosixPath('../corpus/quantized/build a jetpack_9.png'), bytes=1773 (bits=14184, bits per pixel=0.86572265625)\n", + "p=PosixPath('../corpus/quantized/hersheys_train_line_0.png'), bytes=4459 (bits=35672, bits per pixel=2.17724609375)\n", + "Total bytes: 59891\n" + ] + } + ], + "source": [ + "class Predictor:\n", + " def __init__(self):\n", + " self.counts = [Counter() for _ in range(4)]\n", + " gain = 0.020\n", + " weights = [2, 100, 2, 100]\n", + " self.weights = [x/sum(weights) * gain for x in weights]\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 i in range(4):\n", + " # total += self.weights[i] * self.counts[i][contexts[i]]\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] += 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} (bits={t*8}, bits per pixel={t*8/len(img.flatten())})\")\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 aren't super useful as context, compared to orthogonal neighbors\n", + "- Compression bit-by-bit (no RLE, etc) gets us 34.8% compression: definitely doing something, but not great on its own\n", + "- From eyeballing failed predictions, we might benefit from clamping the probability values\n", + " - It looks like we're 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", + " - If we don't need the precision, we could also just reduce probability granularities" + ] + }, + { + "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 +} diff --git a/python/util.py b/python/util.py index e31c2da..fa1765f 100644 --- a/python/util.py +++ b/python/util.py @@ -1,3 +1,4 @@ +from typing import Generator from PIL import Image from pathlib import Path import numpy as np @@ -74,3 +75,9 @@ def _squared_quantization_error(img_np: np.ndarray, palette: np.ndarray): 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('*.*')