Files
color-cell/python/ans-coding.ipynb
T

406 lines
16 KiB
Plaintext

{
"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
}