cmounce/color-cell
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Working compression example with binary ANS

Browse Source
This commit is contained in:
Chris Mounce
2024-12-01 17:46:32 -08:00
parent 211c4a511b
commit d789388f55
2 changed files with 412 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ans-coding.ipynb
+405
View File
@@ -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
}