cc-stuff/cc-pic.py

#!/usr/bin/env python3

from typing import BinaryIO, TextIO
from PIL import Image, ImageColor
from argparse import ArgumentParser, RawTextHelpFormatter
from textwrap import dedent
from functools import lru_cache

try:
    PALETTE_ADAPTIVE = Image.Palette.ADAPTIVE
except Exception:
    PALETTE_ADAPTIVE = Image.ADAPTIVE


class Converter:
    CC_COLORS = [
        ("0", "colors.white"),
        ("1", "colors.orange"),
        ("2", "colors.magenta"),
        ("3", "colors.lightBlue"),
        ("4", "colors.yellow"),
        ("5", "colors.lime"),
        ("6", "colors.pink"),
        ("7", "colors.gray"),
        ("8", "colors.lightGray"),
        ("9", "colors.cyan"),
        ("a", "colors.purple"),
        ("b", "colors.blue"),
        ("c", "colors.brown"),
        ("d", "colors.green"),
        ("e", "colors.red"),
        ("f", "colors.black"),
    ]

    DEFAULT_PALETTE = [
        0xf0, 0xf0, 0xf0,
        0xf2, 0xb2, 0x33,
        0xe5, 0x7f, 0xd8,
        0x99, 0xb2, 0xf2,

        0xde, 0xde, 0x6c,
        0x7f, 0xcc, 0x19,
        0xf2, 0xb2, 0xcc,
        0x4c, 0x4c, 0x4c,

        0x99, 0x99, 0x99,
        0x4c, 0x99, 0xb2,
        0xb2, 0x66, 0xe5,
        0x33, 0x66, 0xcc,

        0x7f, 0x66, 0x4c,
        0x57, 0xa6, 0x4e,
        0xcc, 0x4c, 0x4c,
        0x11, 0x11, 0x11
    ]

    DEFAULT_GRAYSCALE_PALETTE = [
        0xf0, 0xf0, 0xf0,
        0x9d, 0x9d, 0x9d,
        0xbe, 0xbe, 0xbe,
        0xbf, 0xbf, 0xbf,
        0xb8, 0xb8, 0xb8,
        0x76, 0x76, 0x76,
        0xd0, 0xd0, 0xd0,
        0x4c, 0x4c, 0x4c,
        0x99, 0x99, 0x99,
        0x87, 0x87, 0x87,
        0xa9, 0xa9, 0xa9,
        0x77, 0x77, 0x77,
        0x65, 0x65, 0x65,
        0x6e, 0x6e, 0x6e,
        0x76, 0x76, 0x76,
        0x11, 0x11, 0x11
    ]

    PIX_BITS = [[1, 2], [4, 8], [16, 0]]

    MAX_DIFF = 3 * 255

    def __init__(self, image: Image.Image, palette: list[int] | int = PALETTE_ADAPTIVE, dither: bool = True):
        dither_mode = Image.Dither.FLOYDSTEINBERG if dither else Image.Dither.NONE
        if isinstance(palette, list):
            img_pal = Image.new("P", (1, 1))
            img_pal.putpalette(palette)
            self._img = image.quantize(len(palette) // 3, palette=img_pal, dither=dither_mode)
        else:
            self._img = image.convert("P", palette=palette, colors=16, dither=dither_mode)

        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]
        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 in range(len(line)):
                pix = self._imgdata[x + ox, y + oy]
                assert pix < 16, f"{pix} is too big at {x+ox}:{y+oy}"
                brightness = self._brightness(pix)
                if brightness > brightest_l:
                    brightest_l, brightest_i = brightness, pix
                if brightness < darkest_l:
                    darkest_l, darkest_i = brightness, pix
        return darkest_i, brightest_i

    @lru_cache()
    def _is_darker(self, bg: int, fg: int, c: int) -> bool:
        return self._distance(bg, c) < self._distance(fg, c)

    def _get_block(self, x: int, y: int) -> tuple[int, int, int]:
        dark_i, bri_i = self._get_colors(x, y)
        assert dark_i < 16, f"{dark_i} is too big"
        assert bri_i < 16, f"{bri_i} is too big"
        out: int = 0
        for oy, line in enumerate(self.PIX_BITS):
            for ox, bit in enumerate(line):
                if not self._is_darker(
                    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._imgdata[x + 1, y + 2]):
            out ^= 31
            dark_i, bri_i = bri_i, dark_i
        return out, dark_i, bri_i

    @staticmethod
    def _write_varint(fp: BinaryIO, value: int):
        value &= 0xFFFFFFFF
        mask: int = 0xFFFFFF80
        while True:
            if (value & mask) == 0:
                fp.write(bytes([value & 0xFF]))
                return
            fp.write(bytes([(value & 0x7F) | 0x80]))
            value >>= 7

    def export_binary(self, io: BinaryIO, version: int = -1):
        if version == -2:
            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])
                io.write(line)
            return
        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]))
        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]))
        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])
            io.write(line)

    def export(self, io: TextIO):
        io.write("local m = peripheral.find('monitor')\n")
        io.write("m.setTextScale(0.5)\n")
        io.write(f"-- image: {self._img.width}x{self._img.height}\n")
        io.write("\n")
        io.write("-- configuring palette\n")
        for i in range(16):
            r, g, b = self._palette[i * 3 : (i + 1) * 3]
            io.write(
                f"m.setPaletteColor({self.CC_COLORS[i][1]}, 0x{r:02x}{g:02x}{b:02x})\n"
            )
        io.write("\n")
        io.write("-- writing pixels\n")

        for i, y in enumerate(range(0, self._img.height - 2, 3), 1):
            s = []
            bgs = ""
            fgs = ""
            io.write(f"m.setCursorPos(1, {i}); ")
            for x in range(0, self._img.width - 1, 2):
                ch, bg, fg = self._get_block(x, y)
                s.append(ch + 0x80)
                bgs += self.CC_COLORS[bg][0]
                fgs += self.CC_COLORS[fg][0]
            io.write(
                "m.blit(string.char(%s), '%s', '%s')\n"
                % (str.join(", ", map(str, s)), fgs, bgs)
            )


def main():
    parser = ArgumentParser(
        description="ComputerCraft Palette Image converter",
        formatter_class=RawTextHelpFormatter,
    )
    parser.add_argument(
        "-t",
        dest="textmode",
        action="store_true",
        help="Output a Lua script instead of binary image",
    )
    parser.add_argument(
        "-D",
        dest="nodither",
        action="store_true",
        help="Disable dithering"
    )
    parser.add_argument(
        "-W",
        dest="width",
        default=4 * 8 - 1,
        type=int,
        help="Width in characters",
    )
    parser.add_argument(
        "-H",
        dest="height",
        default=3 * 6 - 2,
        type=int,
        help="Height in characters",
    )
    parser.add_argument(
        "-V",
        dest="cpi_version",
        type=int,
        default=-1,
        choices=(-2, -1, 0, 1),
        help=dedent(
            """\
            Force specific CPI version to be used.
            Only applies to binary format.
            Valid versions:
                -V -2       Uses raw format. No headers, default palette.
                            Used for OBCB-CC project.
                -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",
        choices=("center", "cover", "tile", "full", "extend", "fill"),
        default="full",
        help=dedent(
            """\
            Image placement mode (same as in hsetroot)
                -p center       Render image centered on screen
                -p cover        Centered on screen, scaled to fill fully
                -p tile         Render image tiles
                -p full         Maximum aspect ratio
                -p extend       Same as "full" but filling borders
                -p fill         Stretch to fill"""
        ),
    )
    parser.add_argument(
        "-P",
        dest="palette",
        default="auto",
        help=dedent(
            """\
            Palette to be used for that conversion.
            Should be 16 colors or less
            Valid options are:
                -P auto         Determine palette automatically
                -P default      Use default CC:Tweaked color palette
                -P defaultgray  Use default CC:Tweaked grayscale palette
                -P "list:#RRGGBB,#RRGGBB,..."   Use a set list of colors
                -P "cpi:path"   Load palette from a CCPI file
                -P "gpl:path"   Parse GIMP palette file and use first 16 colors
                -P "txt:path"   Load palette from a list of hex values
            """
        )
    )
    parser.add_argument("image_path")
    parser.add_argument("output_path")

    args = parser.parse_args()

    with Image.new("RGB", (args.width * 2, args.height * 3)) as canv:
        with Image.open(args.image_path).convert("RGB") as img:
            if args.placement == "fill":
                canv.paste(img.resize(canv.size), (0, 0))
            elif args.placement in ("full", "extend", "cover"):
                aspect = canv.width / img.width
                if (img.height * aspect > canv.height) != (
                    args.placement == "cover"
                ):
                    aspect = canv.height / img.height
                new_w, new_h = int(img.width * aspect), int(
                    img.height * aspect
                )
                top = int((canv.height - new_h) / 2)
                left = int((canv.width - new_w) / 2)
                resized_img = img.resize((new_w, new_h))
                canv.paste(resized_img, (left, top))
                if args.placement == "extend":
                    if left > 0:
                        right = left - 1 + new_w
                        w = 1
                        while right + w < canv.width:
                            canv.paste(
                                canv.crop(
                                    (left + 1 - w, 0, left + 1, canv.height)
                                ),
                                (left + 1 - w * 2, 0),
                            )
                            canv.paste(
                                canv.crop((right, 0, right + w, canv.height)),
                                (right + w, 0),
                            )
                            w *= 2
                    if top > 0:
                        bottom = top - 1 + new_h
                        h = 1
                        while bottom + h < canv.height:
                            canv.paste(
                                canv.crop(
                                    (0, top + 1 - h, canv.width, top + 1)
                                ),
                                (top + 1 - h * 2, 0),
                            )
                            canv.paste(
                                canv.crop((0, bottom, canv.width, bottom + h)),
                                (0, bottom + h),
                            )
                            h *= 2

            elif args.placement in ("center", "tile"):
                left = int((canv.width - img.width) / 2)
                top = int((canv.height - img.height) / 2)
                if args.placement == "tile":
                    while left > 0:
                        left -= img.width
                    while top > 0:
                        top -= img.height
                    x = left
                    while x < canv.width:
                        y = top
                        while y < canv.height:
                            canv.paste(img, (x, y))
                            y += img.height
                        x += img.width
                else:
                    canv.paste(img, (left, top))
            else:
                pass

            palette = PALETTE_ADAPTIVE
            if args.cpi_version == -2:
                args.palette = "default"

            if args.palette == "auto":
                palette = PALETTE_ADAPTIVE
            elif args.palette == "default":
                palette = Converter.DEFAULT_PALETTE
            elif args.palette == "defaultgray":
                palette = Converter.DEFAULT_GRAYSCALE_PALETTE
            elif args.palette.startswith("txt:"):
                with open(args.palette[4:], "r") as fp:
                    palette = []
                    for line in fp:
                        palette += ImageColor.getcolor(line.strip(), "RGB")  # type: ignore
                        assert len(palette) <= 16 * 3
            elif args.palette.startswith("list:"):
                palette = []
                for c in args.palette[5:].split(","):
                    palette += ImageColor.getcolor(c, "RGB")  # type: ignore
                    assert len(palette) <= 16 * 3
            elif args.palette.startswith("cpi:"):
                raise ValueError("not implemented yet")
            elif args.palette.startswith("gpl:"):
                raise ValueError("not implemented yet")
            else:
                raise ValueError(f"invalid palette identifier: {args.palette!r}")

            converter = Converter(canv, palette, dither=not args.nodither)
            # converter._img.save("/tmp/_ccpictmp.png")
            if args.textmode:
                with open(args.output_path, "w") as fp:
                    converter.export(fp)
            else:
                with open(args.output_path, "wb") as fp:
                    converter.export_binary(fp, args.cpi_version)


if __name__ == "__main__":
    main()