From e61774d4dd022889fefaa2ac96cfc228aef65fda Mon Sep 17 00:00:00 2001
From: hkc <hatkidchan@gmail.com>
Date: Thu, 18 Jan 2024 18:26:08 +0300
Subject: [PATCH] Performance improvements, version selection

---
 cc-pic.py | 60 ++++++++++++++++++++++++++++++++++++++++---------------
 1 file changed, 44 insertions(+), 16 deletions(-)

diff --git a/cc-pic.py b/cc-pic.py
index 764ebef..ba9dcf2 100644
--- a/cc-pic.py
+++ b/cc-pic.py
@@ -4,6 +4,7 @@ from typing import BinaryIO, TextIO
 from PIL import Image
 from argparse import ArgumentParser, RawTextHelpFormatter
 from textwrap import dedent
+from functools import lru_cache
 
 
 class Converter:
@@ -28,31 +29,34 @@ class Converter:
 
     PIX_BITS = [[1, 2], [4, 8], [16, 0]]
 
-    MAX_DIFF = (3**0.5) * 255
+    MAX_DIFF = 3 * 255
 
     def __init__(self, image: Image.Image):
         self._img = image.convert("P", palette=Image.ADAPTIVE, colors=16)
+        self._imgdata = self._img.load()
         self._palette: list[int] = self._img.getpalette() or []
         if len(self._palette) < 16 * 3:
             self._palette += [0] * ((16 * 3) - len(self._palette))
 
+    @lru_cache
     def _brightness(self, i: int) -> float:
         r, g, b = self._palette[i * 3 : (i + 1) * 3]
         return (r + g + b) / 768
 
+    @lru_cache
     def _distance(self, a: int, b: int) -> float:
         r1, g1, b1 = self._palette[a * 3 : (a + 1) * 3]
         r2, g2, b2 = self._palette[b * 3 : (b + 1) * 3]
-        return (
-            (r1 - r2) ** 2 + (g1 - g2) ** 2 + (b1 - b2) ** 2
-        ) ** 0.5 / self.MAX_DIFF
+        rd, gd, bd = r1 - r2, g1 - g2, b1 - b2
+        return (rd * rd + gd * gd + bd * bd) / self.MAX_DIFF
 
+    @lru_cache
     def _get_colors(self, x: int, y: int) -> tuple[int, int]:
         brightest_i, brightest_l = 0, 0
         darkest_i, darkest_l = 0, 768
         for oy, line in enumerate(self.PIX_BITS):
             for ox, bit in enumerate(line):
-                pix = self._img.getpixel((x + ox, y + oy))
+                pix = self._imgdata[x + ox, y + oy]
                 brightness = self._brightness(pix)
                 if brightness > brightest_l:
                     brightest_l, brightest_i = brightness, pix
@@ -69,11 +73,11 @@ class Converter:
         for oy, line in enumerate(self.PIX_BITS):
             for ox, bit in enumerate(line):
                 if not self._is_darker(
-                    dark_i, bri_i, self._img.getpixel((x + ox, y + oy))
+                    dark_i, bri_i, self._imgdata[x + ox, y + oy]
                 ):
                     out |= bit
         # bottom right pixel fix?
-        if not self._is_darker(dark_i, bri_i, self._img.getpixel((x + 1, y + 2))):
+        if not self._is_darker(dark_i, bri_i, self._imgdata[x + 1, y + 2]):
             out ^= 31
             dark_i, bri_i = bri_i, dark_i
         return out, dark_i, bri_i
@@ -89,24 +93,31 @@ class Converter:
             fp.write(bytes([(value & 0x7F) | 0x80]))
             value >>= 7
 
-    def export_binary(self, io: BinaryIO):
-        if self._img.width <= 510 and self._img.height <= 765:
-            io.write(b"CCPI") # old format
+    def export_binary(self, io: BinaryIO, version: int = -1):
+        if version == -1:
+            if self._img.width <= 255 * 2 and self._img.height < 255 * 3:
+                version = 0
+            else:
+                version = 1
+
+        if version == 0:
+            io.write(b"CCPI")  # old format
             io.write(bytes([self._img.width // 2, self._img.height // 3, 0]))
             io.write(bytes(self._palette[: 16 * 3]))
-        else:
-            io.write(b"CPI\x01") # CPIv1
+        elif version == 1:
+            io.write(b"CPI\x01")  # CPIv1
             self._write_varint(io, self._img.width // 2)
             self._write_varint(io, self._img.height // 3)
             io.write(bytes(self._palette[: 16 * 3]))
-        written = 0
+        else:
+            raise ValueError(f"invalid version {version}")
+
         for y in range(0, self._img.height - 2, 3):
             line: bytearray = bytearray()
             for x in range(0, self._img.width - 1, 2):
                 ch, bg, fg = self._get_block(x, y)
                 line.extend([(ch + 0x80) & 0xFF, fg << 4 | bg])
-            written += io.write(line)
-        assert written == (self._img.width // 2) * (self._img.height // 3) * 2
+            io.write(line)
 
     def export(self, io: TextIO):
         io.write("local m = peripheral.find('monitor')\n")
@@ -163,6 +174,23 @@ def main():
         type=int,
         help="Height in characters",
     )
+    parser.add_argument(
+        "-V",
+        dest="cpi_version",
+        type=int,
+        default=-1,
+        choices=(-1, 0, 1),
+        help=dedent(
+            """\
+            Force specific CPI version to be used.
+            Only applies to binary format.
+            Valid versions:
+                -V -1       Choose any fitting one
+                            For images smaller than 255x255, uses CPIv0
+                -V 0        OG CPI, 255x255 maximum, uncompressed
+                -V 1        CPIv1, huge images, uncompressed"""
+        ),
+    )
     parser.add_argument(
         "-p",
         dest="placement",
@@ -260,7 +288,7 @@ def main():
                     converter.export(fp)
             else:
                 with open(args.output_path, "wb") as fp:
-                    converter.export_binary(fp)
+                    converter.export_binary(fp, args.cpi_version)
 
 
 if __name__ == "__main__":