ii-vision/transcoder/make_data_tables.py

import functools
import os
import sys
from typing import Iterable, Type

import colormath.color_conversions
import colormath.color_diff
import colormath.color_objects
import numpy as np
import weighted_levenshtein
from etaprogress.progress import ProgressBar

import colours
import palette
import screen


PIXEL_CHARS = "0123456789ABCDEF"
DATA_DIR = "transcoder/data"

def pixel_char(i: int) -> str:
    return PIXEL_CHARS[i]


@functools.lru_cache(None)
def pixel_string(pixels: Iterable[int]) -> str:
    return "".join(pixel_char(p) for p in pixels)


class EditDistanceParams:
    """Data class for parameters to Damerau-Levenshtein edit distance."""

    # Don't even consider insertions and deletions into the string, they don't
    # make sense for comparing pixel strings
    insert_costs = np.ones(128, dtype=np.float64) * 100000
    delete_costs = np.ones(128, dtype=np.float64) * 100000

    # Smallest substitution value is ~20 from palette.diff_matrices, i.e.
    # we always prefer to transpose 2 pixels rather than substituting colours.
    # TODO: is quality really better allowing transposes?
    transpose_costs = np.ones((128, 128), dtype=np.float64)

    # These will be filled in later
    substitute_costs = np.zeros((128, 128), dtype=np.float64)

    # Substitution costs to use when evaluating other potential offsets at which
    # to store a content byte.  We penalize more harshly for introducing
    # errors that alter pixel colours, since these tend to be very
    # noticeable as visual noise.
    #
    # TODO: currently unused
    error_substitute_costs = np.zeros((128, 128), dtype=np.float64)


def compute_diff_matrix(pal: Type[palette.BasePalette]):
    """Compute matrix of perceptual distance between colour pairs.

    Specifically CIE2000 delta values for this palette.
    """
    dm = np.ndarray(shape=(16, 16), dtype=np.int32)

    for colour1, a in pal.RGB.items():
        alab = colormath.color_conversions.convert_color(
            a, colormath.color_objects.LabColor)
        for colour2, b in pal.RGB.items():
            blab = colormath.color_conversions.convert_color(
                b, colormath.color_objects.LabColor)
            dm[colour1.value, colour2.value] = int(
                colormath.color_diff.delta_e_cie2000(alab, blab))
    return dm


def compute_substitute_costs(pal: Type[palette.BasePalette]):
    """Compute costs for substituting one colour pixel for another."""

    edp = EditDistanceParams()

    diff_matrix = compute_diff_matrix(pal)

    # Penalty for changing colour
    for i, c in enumerate(PIXEL_CHARS):
        for j, d in enumerate(PIXEL_CHARS):
            cost = diff_matrix[i, j]
            edp.substitute_costs[(ord(c), ord(d))] = cost
            edp.substitute_costs[(ord(d), ord(c))] = cost
            edp.error_substitute_costs[(ord(c), ord(d))] = 5 * cost
            edp.error_substitute_costs[(ord(d), ord(c))] = 5 * cost

    return edp


def edit_distance(
        edp: EditDistanceParams,
        a: str,
        b: str,
        error: bool) -> np.float64:
    """Damerau-Levenshtein edit distance between two pixel strings."""
    res = weighted_levenshtein.dam_lev(
        a, b,
        insert_costs=edp.insert_costs,
        delete_costs=edp.delete_costs,
        substitute_costs=(
            edp.error_substitute_costs if error else edp.substitute_costs),
    )

    # Make sure result can fit in a uint16
    assert (0 <= res < 2 ** 16), res
    return res


def compute_edit_distance(
        edp: EditDistanceParams,
        bitmap_cls: Type[screen.Bitmap],
        nominal_colours: Type[colours.NominalColours]
) -> np.ndarray:
    """Computes edit distance matrix between all pairs of pixel strings.

    Enumerates all possible values of the masked bit representation from
    bitmap_cls (assuming it is contiguous, i.e. we enumerate all
    2**bitmap_cls.MASKED_BITS values).  These are mapped to the dot
    representation, turned into coloured pixel strings, and we compute the
    edit distance.

    The effect of this is that we precompute the effect of storing all possible
    byte values against all possible screen backgrounds (e.g. as
    influencing/influenced by neighbouring bytes).
    """

    bits = bitmap_cls.MASKED_BITS

    bitrange = np.uint64(2 ** bits)

    edit = np.zeros(
        shape=(len(bitmap_cls.BYTE_MASKS), np.uint64(bitrange * bitrange)),
        dtype=np.uint16)

    bar = ProgressBar(
        bitrange * (bitrange - 1) / 2 * len(bitmap_cls.PHASES), max_width=80)

    num_dots = bitmap_cls.MASKED_DOTS

    cnt = 0
    for i in range(np.uint64(bitrange)):
        pair_base = np.uint64(i) << bits
        for o, ph in enumerate(bitmap_cls.PHASES):
            # Compute this in the outer loop since it's invariant under j
            first_dots = bitmap_cls.to_dots(i, byte_offset=o)
            first_pixels = pixel_string(
                colours.dots_to_nominal_colour_pixel_values(
                    num_dots, first_dots, nominal_colours,
                    init_phase=ph)
            )

            # Matrix is symmetrical with zero diagonal so only need to compute
            # upper triangle
            for j in range(i):
                cnt += 1
                if cnt % 100000 == 0:
                    bar.numerator = cnt
                    print(bar, end='\r')
                    sys.stdout.flush()

                pair = pair_base + np.uint64(j)

                second_dots = bitmap_cls.to_dots(j, byte_offset=o)
                second_pixels = pixel_string(
                    colours.dots_to_nominal_colour_pixel_values(
                        num_dots, second_dots, nominal_colours,
                        init_phase=ph)
                )
                edit[o, pair] = edit_distance(
                    edp, first_pixels, second_pixels, error=False)

    return edit


def make_edit_distance(
        pal: Type[palette.BasePalette],
        edp: EditDistanceParams,
        bitmap_cls: Type[screen.Bitmap],
        nominal_colours: Type[colours.NominalColours]
):
    """Write file containing (D)HGR edit distance matrix for a palette."""

    dist = compute_edit_distance(edp, bitmap_cls, nominal_colours)
    data = "%s/%s_palette_%d_edit_distance.npz" % (
        DATA_DIR, bitmap_cls.NAME, pal.ID.value)
    np.savez_compressed(data, edit_distance=dist)


def main():
    try:
        os.mkdir(DATA_DIR, mode=0o755)
    except FileExistsError:
        pass

    for p in palette.PALETTES.values():
        print("Processing palette %s" % p)
        edp = compute_substitute_costs(p)

        # TODO: still worth using error distance matrices?

        make_edit_distance(p, edp, screen.HGRBitmap, colours.HGRColours)
        make_edit_distance(p, edp, screen.DHGRBitmap, colours.DHGRColours)


if __name__ == "__main__":
    main()