pokemon-color-search/anim_ingest.py

import io

from PIL import Image
from bs4 import BeautifulSoup
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/"

# removing all forms of a pokemon, and also pokestars
start_with_filters = [
  # no significant visual changes
  "arceus-", "silvally-", "genesect-", "pumpkaboo-", "gourgeist-", "unown-", "giratina-",
  # cannot start the battle in alternate form
  "castform-", "cherrim-", "aegislash-", "xerneas-", "wishiwashi-", "eiscue-", "mimikyu-",
  "cramorant-", "morepeko-",
  # weird event thing
  "greninja-", "eevee-", "pikachu-", "zarude-", "magearna-",
  # pokestars
  "pokestar",
]

# removing all forms of a type
end_with_filters = [
  "-mega.gif", "-megax.gif", "-megay.gif", "-primal.gif", "-ultra.gif",
  "-gmax.gif", "-eternamax.gif", "-totem.gif", "-f.gif", "-b.gif",
]

# removing pokemon entirely
full_filters = [
  # darmanitan zen forms (cannot start in zen)
  "darmanitan-galarzen.gif", "darmanitan-zen.gif",
  # minior core forms (cannot start in anything but -meteor, renamed below)
  "minior.gif", "minior-blue.gif", "minior-green.gif", "minior-indigo.gif", 
  "minior-orange.gif", "minior-violet.gif", "minior-yellow.gif",
  # because it is a create-a-pokemon
  "astrolotl.gif", "aurumoth.gif", "caribolt.gif", "cawmodore.gif", "chromera.gif", "crucibelle.gif",
  "equilibra.gif", "fidgit.gif", "jumbao.gif", "justyke.gif", "kerfluffle.gif", "kitsunoh.gif",
  "krilowatt.gif", "malaconda.gif", "miasmaw.gif", "mollux.gif", "naviathan.gif", "necturna.gif",
  "pajantom.gif", "plasmanta.gif", "pluffle.gif", "protowatt.gif", "scratchet.gif", "smogecko.gif",
  "smoguana.gif", "smokomodo.gif", "snaelstrom.gif", "stratagem.gif", "tomohawk.gif", "volkraken.gif", "voodoom.gif",
  # typos/duplicates
  "buffalant.gif", "klinklang-back.gif", "krikretot.gif", "pumpkabo-super.gif", "magcargo%20.gif",
  "ratatta-a.gif", "ratatta-alola.gif", "raticate-a.gif",
  "rotom-h.gif", "rotom-m.gif", "rotom-s.gif", "rotom-w.gif",
  # not a pokemon
  "substitute.gif",
]

# force certain pokemon to stay
force_keep = [ "meowstic-f.gif" ]

# rename certain pokemon after the fact
rename = {
  # dash consistency
  "nidoranm": "nidoran-m",
  "nidoranf": "nidoran-f",
  "porygonz": "porygon-z",
  "tapubulu": "tapu-bulu",
  "tapufini": "tapu-fini",
  "tapukoko": "tapu-koko",
  "tapulele": "tapu-lele",
  "hooh": "ho-oh",
  "mimejr": "mime-jr",
  "mrmime": "mr-mime",
  "mrmime-galar": "mr-mime-galar",
  "mrrime": "mr-rime",
  "jangmoo": "jangmo-o",
  "hakamoo": "hakamo-o",
  "kommoo": "kommo-o",
  "typenull": "type-null",
  "oricorio-pompom": "oricorio-pom-pom",
  "necrozma-duskmane": "necrozma-dusk-mane",
  "necrozma-dawnwings": "necrozma-dawn-wings",
  "toxtricity-lowkey": "toxtricity-low-key",

  # rename forms
  "shellos": "shellos-west",
  "shaymin": "shaymin-land",
  "meloetta": "meloetta-aria",
  "keldeo": "keldeo-ordinary",
  "hoopa": "hoopa-confined",
  "burmy": "burmy-plant",
  "wormadam": "wormadam-plant",
  "deerling": "deerling-spring",
  "sawsbuck": "sawsbuck-spring",
  "vivillon": "vivillon-meadow",
  "basculin": "basculin-redstriped",
  "meowstic": "meowstic-male",
  "meowstic-f": "meowstic-female",
  "flabebe": "flabebe-red",
  "floette": "floette-red",
  "florges": "florges-red",
  "minior-meteor": "minior",
  "sinistea": "sinistea-phony",
  "polteageist": "polteageist-phony",
  "gastrodon": "gastrodon-west",
  "furfrou": "furfrou-natural",
  "wishiwashi": "wishiwashi-school",
  "tornadus": "tornadus-incarnate",
  "landorus": "landorus-incarnate",
  "thundurus": "thundurus-incarnate",
  "calyrex-ice": "calyrex-ice-rider",
  "calyrex-shadow": "calyrex-shadow-rider",
  "urshifu-rapidstrike": "urshifu-rapid-strike",
  "zacian": "zacian-hero",
  "zamazenta": "zamazenta-hero",
}


def get_all_pokemon() -> list[str]:
  soup = BeautifulSoup(requests.get(back_base).text, "html.parser")
  gifs = [href for a in soup.find_all("a") if (href := a.get("href")).endswith("gif")]
  return [
    g[:-4] 
    for g in gifs 
    if g in force_keep or (
      g not in full_filters 
      and not any(g.startswith(f) for f in start_with_filters) 
      and not any(g.endswith(f) for f in end_with_filters)
    )
  ]


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


if __name__ == "__main__":
  pkmn = get_all_pokemon()
  print("Found", len(pkmn), "sprites...")
  errors = []
  with open("database-anim.js", "w") as outfile:
    outfile.write("const databaseV2 = [\n")
    for name in pkmn:
      print("Ingesting", name, "...")
      try:
        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'  [ "{rename.get(name, name)}", {", ".join(str(n) for n in stats)} ],\n')
      except Exception as e:
        print(e)
        errors.append(name)
    outfile.write("];\n")
  print("Errors:", errors)