209 lines
6.7 KiB
Python
209 lines
6.7 KiB
Python
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()
|