diff --git a/scripts/update-palette b/scripts/update-palette new file mode 100755 index 00000000..1e665715 --- /dev/null +++ b/scripts/update-palette @@ -0,0 +1,311 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: BSD-3-Clause +# +# update-palette - Update the palettes of all PNG graphics +# +# All of the PNGs in Freedoom are paletted, and the palettes of each PNG match +# the colours in the PLAYPAL lump. If a user wants to make changes to the +# palette, they would have to update the palette in all of Freedoom's graphics +# for consistency. +# +# This script takes a new PLAYPAL as an argument, compares the old and new +# palettes, and modifies every paletted PNG file in the repo so that the new +# colours are used. +# +# Dependencies: +# - pypng: https://gitlab.com/drj11/pypng (retrieved Sept 10 2023) + +import argparse +import array +from functools import reduce +from itertools import cycle, zip_longest +import struct +from operator import eq +import os +from os.path import dirname, relpath, realpath, join, normpath +import png +from sys import argv +import zlib + +# Misc. utility functions +# https://docs.python.org/3.8/library/itertools.html#itertools-recipes +def grouper(iterable, n, fillvalue=None): + "Collect data into fixed-length chunks or blocks" + # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx" + args = [iter(iterable)] * n + return zip_longest(*args, fillvalue=fillvalue) + + +# Parse the command line arguments, and return a dict with the arguments +def parse_args(): + parser = argparse.ArgumentParser( + "update-palette", + description="This script takes a new palette, compares the new " + "palette with the old one, and scans and updates images in the repo " + "which use the colours that were replaced in the new palette." + ) + parser.add_argument("palette", help="The new palette to use") + # This is a potential vulnerability, and besides, it doesn't make + # sense to update the palette for images which aren't in the repo. + # parser.add_argument( + # "--dir", "-d", help=( + # "The directory to recursively process. " + # "Defaults to repository root directory."), + # default=normpath(join(dirname(realpath(argv[0])), ".."))) + parser.add_argument( + "-d", "--dry-run", + help="Do not modify any PNGs, just show which ones would be modified", + action='store_true') + args = parser.parse_args() + return args + + +# Compare the old palette and the new palette, and return a dict with the +# differences. +def compare_palettes(directory, new_palette): + old_palette = join(directory, "lumps", "playpal", "playpal-base.lmp") + + # The "new" palette is the old palette? + old_pal_full_path = relpath(normpath(old_palette), directory) + new_pal_full_path = relpath(normpath(new_palette), directory) + if old_pal_full_path == new_pal_full_path: + raise ValueError("You're trying to replace the old palette with " + "itself! Try another palette.") + + # Read both palettes into a more usable format + with open(old_palette, "rb") as handle: + old_palette = handle.read(768) + if len(old_palette) < 768: + raise ValueError("Old palette is too short!") + old_palette = list(grouper(old_palette[:768], 3)) + + with open(new_palette, "rb") as handle: + new_palette = handle.read(768) + if len(new_palette) < 768: + raise ValueError("New palette is too short!") + new_palette = list(grouper(new_palette[:768], 3)) + + # Given a colour palette and a dict, return a dict with the indexes of + # each colour. This function is meant to be used with functools.reduce. + def get_duplicate_colours(value, index_colour): + index, colour = index_colour + value.setdefault(colour, []).append(index) + return value + + # Scan the old palette for duplicate colours + old_palette_duplicates = reduce( + get_duplicate_colours, enumerate(old_palette), {}) + # Eliminate unique colours + old_palette_duplicates = dict(filter( + # kv[1] is the value, and it's unique if it's length is 1 + lambda kv: None if len(kv[1]) == 1 else kv, + old_palette_duplicates.items() + )) + + # Map old indices to new colours for now. + old_to_new = {} + for index, zipper in enumerate(zip(old_palette, new_palette)): + old_colour, new_colour = zipper + if old_colour != new_colour: + old_to_new[index] = new_colour + + # Does the new palette replace ALL instances of a duplicate colour in the + # old palette? If so, map the old colour to the first replacement in the + # new palette, but if not, leave the old colour unchanged. + replaced = {} + for colour, indices in old_palette_duplicates.items(): + if all(map(lambda i: i in old_to_new.keys(), indices)): + replaced[colour] = old_to_new[min(indices)] + print("replaced", replaced) + + # Find the closest colour in the old palette for each grayscale colour + def is_grayscale(colour): + return all(map(eq, colour[1], cycle((colour[1][0],)))) + # def to_grayscale(colour): + # return colour[0] if is_grayscale(colour) else colour + grayscale_colours_in_palette = dict( + map(lambda g: (g[0], g[1][0]), + filter(is_grayscale, enumerate(new_palette))) + ) + def closest_grayscale(grayscale): + nonlocal grayscale_colours_in_palette + distances = sorted(map( + lambda kv: (kv[0], abs(kv[1] - grayscale)), + grayscale_colours_in_palette.items() + ), key=lambda kv: kv[1]) + closest = distances[0][1] + distances = sorted( + filter(lambda kv: kv[1] == closest, distances), + key=lambda kv: kv[0]) + return distances[0][0] + # This is a map from grayscale to palette index + gray_map = array.array('B', map(closest_grayscale, range(256))) + # Turn it into a grayscale to colour map, in case any grayscale colours + # were changed in the new palette + gray_map = dict( + map(lambda g: (g[0], old_to_new[g[1]]), + filter( + lambda g: g[1] in old_to_new, + enumerate(gray_map)))) + print(len(gray_map), gray_map) + + # Replace the keys in old_to_new, which are indices, with the old palette + # colours they correspond to. This way, we have a colour-to-colour dict. + old_to_new = dict(filter(lambda kv: kv[1], map( + lambda iv: (old_palette[iv[0]], + replaced.get(old_palette[iv[0]], iv[1])), + old_to_new.items() + ))) + return old_to_new, gray_map + + +# "Stolen" from the map-color-index script +# Process a directory recursively for PNG files. +def process_dir(colour_map, gray_map, dry_run, directory, palette): + pngs_changed_count = 0 + pngs_examined_count = 0 + + for dirpath, dirnames, filenames in os.walk(directory): + for png_base in filenames: + if not png_base.lower().endswith(".png"): + continue + png_path = os.path.join(dirpath, png_base) + pngs_examined_count += 1 + if process_png(colour_map, gray_map, png_path, dry_run, directory): + pngs_changed_count += 1 + + +# Process a PNG file in place. +def process_png(colour_map, gray_map, png_path, dry, directory): + # Read the PNG file + png_reader = png.Reader(filename=png_path) + + # Change the old colours to the new colours + modified_palette_colours = set() + def maybe_modify_plte(plte_data): + nonlocal colour_map + nonlocal modified_palette_colours + modified = False + + # A pypng palette colour is 3 elements long, or 4 if the image has + # transparent colours. + if len(plte_data[0]) > 3: + alphas = list(map(lambda co: co[3], plte_data)) + else: + alphas = [None] * len(plte_data) + + colours = list(map(lambda co: co[0:3], plte_data)) + # Scan the palette for colours that should be changed, and change them + # according to colour_map + for index, colour in enumerate(colours): + if colour in colour_map: + modified = True + modified_palette_colours.add(index) + colours[index] = colour_map[colour] + # pypng needs the palette as an iterable of 3-tuples for opaque + # paletted images, or 4-tuples for transparent paletted images + colours = list(map(lambda c, a: (*c,) if a is None else (*c, a), + colours, alphas)) + return modified, colours + + def is_paletted_colour_changed(row): + nonlocal modified_palette_colours + return any(map(lambda colour_index: \ + colour_index in modified_palette_colours, row)) + + def maybe_modify_truecolour_image(rows, channels, transparent_colour): + nonlocal colour_map + modified = False + + def maybe_modify_row(row): + nonlocal colour_map + nonlocal channels + nonlocal transparent_colour + nonlocal modified + def maybe_modify_colour(co): + nonlocal colour_map + nonlocal transparent_colour + nonlocal modified + # Seperate RGB and alpha + co_rgb = co[0:3] + co_alpha = co[3] if len(co) > 3 else None + # Get the new colour from the colour map + new_colour = ( + colour_map.get(co_rgb, co_rgb) + if co_rgb != transparent_colour else co_rgb + ) + if co_alpha is not None: + new_colour = (*new_colour, co_alpha) + modified = modified or new_colour != co + return new_colour + pixels = map(maybe_modify_colour, + grouper(row, channels)) + pixels = bytearray(b"".join(map(bytes, pixels))) + return pixels + + new_rows = list(map(maybe_modify_row, rows)) + + return modified, new_rows + + def maybe_modify_grayscale_image(rows, channels): + pass + + plte_modified = False + idat_modified = False + width, height, rows, info = png_reader.read() + # "rows" is an iterator, so it needs to be a list so that it can be + # modified if the image is true-colour, or scanned if the image is + # paletted + rows = list(rows) + grayscale = info.get("greyscale") + is_paletted = info.get("palette") + bit_depth = info.get("bitdepth") + transparent_colour = info.get("transparent") + has_alpha = info.get("alpha", False) + channels = info.get("planes") + if is_paletted: + # Modify the PLTE chunk if necessary, and check if the PLTE + # modifications affect the colours used in the IDAT chunk. + plte_modified, new_palette = maybe_modify_plte(info["palette"]) + if plte_modified: + idat_modified = any(map(is_paletted_colour_changed, rows)) + elif not grayscale: + new_palette = None + idat_modified, rows = maybe_modify_truecolour_image( + rows, channels, transparent_colour) + else: + idat_modified, rows = maybe_modify_grayscale_image(rows, channels) + + # Write the modified PNG file + if idat_modified: + print("{} was changed".format(relpath(png_path, start=directory))) + if not dry: + with open(png_path, "wb") as png_file: + png_writer = png.Writer( + width, height, bitdepth=bit_depth, palette=new_palette, + alpha=has_alpha, transparent=transparent_colour, + greyscale=grayscale) + png_writer.write(png_file, rows) + + return idat_modified + + +if __name__ == "__main__": + args = parse_args() + directory = normpath(join(dirname(realpath(argv[0])), "..")) + colour_map, gray_map = compare_palettes(directory, args.palette) + print("The rest of the script is disabled for development right now") + exit(0) + + process_dir(colour_map, gray_map, directory=directory, **vars(args)) + # Replace old playpal-base.lmp + if not args.dry_run: + playpal_base_path = ( + join(directory, "lumps", "playpal", "playpal-base.lmp")) + # The old/new palette being the same is checked in compare_palettes + with open(args.palette, "rb") as new_palfile, \ + open(playpal_base_path, "wb") as old_palfile: + # Only copy the first 768 bytes of the new palette to playpal-base + new_pal = new_palfile.read(768) + # Palette length is checked in compare_palettes + old_palfile.write(new_pal)