pokemon-color-search/ingest.py

#!/usr/bin/env python3
from collections import namedtuple

import numpy as np
from scipy.cluster import vq
from PIL import Image
from colorspacious import cspace_convert

verbose = False
seed = 20220211


def is_outline(r: int, g: int, b: int, a: int) -> bool:
    # returns true if a pixel is transparent or pure black
    return a == 0 or (r, g, b) == (0, 0, 0)


def x_metric(pixels: np.array) -> float:
    # X metric - the mean squared Euclidean norm
    # computed as the sum of the squares of the components of the pixels,
    # normalized by the number of pixels
    if verbose:
        print("  Computing X...")
    return sum(sum(pixels ** 2)) / len(pixels)


def y_metric(pixels: np.array) -> np.array:
    # Y metric - the mean pixel of the image
    if verbose:
        print("  Computing Y...")
    return pixels.mean(0)


def z_metric(pixels: np.array) -> tuple[np.array, np.array]:
    # Z metric - run k-means, and return the mean of each cluster
    # k chosen somewhat arbitrarily to be 3
    if verbose:
        print("  Computing Z...")
    means, labels = vq.kmeans2(pixels.astype(float), 3, minit="++", seed=seed)
    return means, np.unique(labels, return_counts=True)[1] / len(pixels)


ImageInfo = namedtuple(
    "ImageInfo", ["name", "xrgb", "xjab", "yrgb", "yjab", "zrgb", "zjab"]
)


def ingest_png(file_name: str) -> ImageInfo:
    print(f"Ingesting {file_name}")

    # image name - strip leading path and trailing extension
    name = file_name.rsplit("/", maxsplit=1)[1].split(".", maxsplit=1)[0]

    # read non-outline pixels of image
    rgb_pixels = np.array([
        (r, g, b)
        for r, g, b, a in Image.open(file_name).convert("RGBA").getdata()
        if not is_outline(r, g, b, a)
    ])

    # convert RGB pixels to CAM02 values
    jab_pixels = cspace_convert(rgb_pixels, "sRGB255", "CAM02-UCS")

    # compute and return metrics
    return ImageInfo(
        name=name,
        xrgb=x_metric(rgb_pixels),
        xjab=x_metric(jab_pixels),
        yrgb=y_metric(rgb_pixels),
        yjab=y_metric(jab_pixels),
        zrgb=z_metric(rgb_pixels),
        zjab=z_metric(jab_pixels),
    )


if __name__ == "__main__":
    import csv
    import os
    import sys

    dir = "pngs" if len(sys.argv) < 2 else sys.argv[1]

    data = [
        ingest_png(dir + "/" + fn)
        for f in os.listdir(dir)
        if (fn := os.fsdecode(f)).endswith(".png")
    ]

    with open("database.csv", "w") as outfile:
        writer = csv.writer(outfile, delimiter=",", quotechar="'")
        writer.writerows([d.name, d.xrgb, *d.yrgb] for d in data)

    with open("database-cam02.csv", "w") as outfile:
        writer = csv.writer(outfile, delimiter=",", quotechar="'")
        writer.writerows([d.name, d.xjab, *d.yjab] for d in data)


    def write_array(contents):
        return f"[ {', '.join(str(c) for c in contents)} ]"


    with open("database.js", "w") as outfile:
        outfile.write("const database = [\n")
        for info in data:
            fields = ", ".join(
                (
                    f'name: "{info.name}"',
                    f"xRGB: {info.xrgb}",
                    f"xJAB: {info.xjab}",
                    f"yRGB: {write_array(info.yrgb)}",
                    f"yJAB: {write_array(info.yjab)}",
                    f"zRGB: {write_array([write_array(z) for z in info.zrgb[0]] + [write_array(info.zjab[1])])}",
                    f"zJAB: {write_array([write_array(z) for z in info.zjab[0]] + [write_array(info.zjab[1])])}",
                )
            )
            outfile.write(f"  {{ {fields} }},\n")
        outfile.write("];\n")