pokemon-color-search/ingest.py

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

from PIL import Image

from convert import rgb_to_cieluv


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: list[tuple[float, float, float]]) -> 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
  return sum(comp * comp for pix in pixels for comp in pix) / len(pixels)


def y_metric(pixels: list[tuple[float, float, float]]) -> tuple[float, float, float]:
  # Y metric - the mean pixel of the image
  return tuple(sum(p[i] for p in pixels) / len(pixels) for i in range(3))


ImageInfo = namedtuple("ImageInfo", ["name", "xrgb", "xluv", "yr", "yg", "yb", "yl", "yu", "yv"])


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 = [(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 CIELUV values
  luv_pixels = [rgb_to_cieluv(*p) for p in rgb_pixels]

  # compute and return metrics
  xrgb = x_metric(rgb_pixels)
  xluv = x_metric(luv_pixels)
  yr, yg, yb = y_metric(rgb_pixels)
  yl, yu, yv = y_metric(luv_pixels)
  return ImageInfo(
    name=name,
    xrgb=xrgb,
    xluv=xluv,
    yr=yr,
    yg=yg,
    yb=yb,
    yl=yl,
    yu=yu,
    yv=yv,
  )


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

  data = [ingest_png("pngs/" + fn) for f in os.listdir("pngs") 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.yr, d.yg, d.yb] for d in data)

  with open("database-luv.csv", "w") as outfile:
    writer = csv.writer(outfile, delimiter=",", quotechar="'")
    writer.writerows([d.name, d.xluv, d.yl, d.yu, d.yv] for d in data)

  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"xLUV: {info.xluv}",
        f"yRGB: [ {info.yr}, {info.yg}, {info.yb} ]",
        f"yLUV: [ {info.yl}, {info.yu}, {info.yv} ]",
      ))
      outfile.write(f"  {{ {fields} }},\n")
    outfile.write("];\n")