pokemon-color-search/anim-ingest.py

import io
import os

from PIL import Image
from colorspacious import cspace_convert
import requests
import numpy as np

import ingest

base = "https://play.pokemonshowdown.com/sprites/ani/"
back_base = "https://play.pokemonshowdown.com/sprites/ani-back/"


def get_all_pokemon() -> list[str]:
  return ["abomasnow", "malamar", "porygon", "luxray", "eternatus", "roselia"]


def load_image(base: str, name: str) -> Image:
  return Image.open(io.BytesIO(requests.get(base + name + ".gif").content))


def get_all_pixels(im: Image) -> list[tuple[int, int, int]]:
  rgb_pixels = []
  for fr in range(getattr(im, "n_frames", 1)):
    im.seek(fr)
    rgb_pixels += [
      (r, g, b) 
      for r, g, b, a in im.convert("RGBA").getdata() 
      if not ingest.is_outline(r, g, b, a)
    ]
  return rgb_pixels


with open("database-anim.js", "w") as outfile:
  outfile.write("const databaseV2 = [\n")
  for name in get_all_pokemon():
    print("Ingesting", name, "...")
    front = get_all_pixels(load_image(base, name))
    back = get_all_pixels(load_image(back_base, name))
    rgb_pixels = np.array(front + back)
    jab_pixels = cspace_convert(rgb_pixels, "sRGB255", "CAM02-UCS")
    stats = [len(rgb_pixels), *ingest.all_stats(jab_pixels), *ingest.all_stats(rgb_pixels)]
    outfile.write(f'  [ "{name}", {", ".join(str(n) for n in stats)} ],\n')
  outfile.write("];\n")