pokemon-color-search/gen9_ingest.py

import json

import numpy as np
from colorspacious import cspace_convert

from ingest import is_outline
from anim_ingest import (
    get_all_pokemon,
    flatten_stats,
    compute_stats,
    base as old_base,
    load_image,
    score_clustering_jab,
    merge_dist_jab,
    score_clustering_rgb,
    merge_dist_rgb,
)

new_base = "https://play.pokemonshowdown.com/sprites/gen5/"
png = ".png"

# TODO this could probably be simplified but it works for a first pass
new_pokemon = sorted(
    set(get_all_pokemon(new_base, png)) - set(get_all_pokemon(old_base))
)

# hacked together from what's in anim_ingest and ingest
for name in new_pokemon:
    print(f"Ingesting {name}...")
    image = load_image(new_base, name, png)

    # read non-outline pixels of image
    rgb_pixels = np.array(
        [
            (r, g, b)
            for r, g, b, a in image.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 stats
    jab_stats = flatten_stats(
        compute_stats(
            jab_pixels,
            score_clustering_jab,
            merge_dist_jab,
        )
    )[1:]
    rgb_stats = flatten_stats(
        compute_stats(
            rgb_pixels,
            score_clustering_rgb,
            merge_dist_rgb,
        )
    )[1:]

    # and fuck it, just print it to a new file and we'll copy it in manually
    with open("database-gen9.json", "a") as outfile:
        outfile.write(json.dumps([name, len(rgb_pixels), *jab_stats, *rgb_stats]) + ",\n")