Pokédex 1–1024 v2

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
all_colors_with_pokemon.py — "better" version focused on sprite fidelity + exact 24-bit coverage

Goals
- Output: 4096×4096 RGBA PNG.
- RGB: every 24-bit RGB value (0..16,777,215) appears **exactly once** in the whole image.
- Pokémon: a 32×32 grid of sprites (IDs 1..1024), drawn from the *same* 24-bit pool.
- Background alpha: 128; Sprite alpha: 255.
- Performance: GPU used for the big shuffle when available; CPU/Numba for the rest.

Key idea (why this works and why sprites don't turn into noise)
- The canvas starts as a full permutation of all 24-bit colors.
- For each sprite pixel, we swap its desired color (or a very close unused substitute if that RGB was already used
  by a previous sprite pixel) into that pixel location.
- For duplicates, instead of a slow global nearest-neighbor search (which often devolves into far-away colors
  and produces "noise"), we allocate substitute colors from a **quantized RGB bin pool**.
  This keeps substitutes near the original color and stays fast.

Dependencies
    pip install numpy pillow requests tqdm psutil numba cupy-cuda12x

Colab usage
    from google.colab import drive
    drive.mount('/content/drive')
    !pip install numpy pillow requests tqdm psutil numba cupy-cuda12x
    !python all_colors_with_pokemon.py
"""

import os
import logging
import psutil
from io import BytesIO

import numpy as np
from PIL import Image
import requests
from tqdm import tqdm
from numba import njit

# -----------------------------
# GPU shuffle (CuPy)
# -----------------------------
USE_GPU = True
try:
    import cupy as cp
except Exception:
    USE_GPU = False

# -----------------------------
# Configuration
# -----------------------------
IMAGE_SIZE = 4096
TOTAL_PIXELS = IMAGE_SIZE * IMAGE_SIZE  # 16,777,216
SEED = 42
POKE_COUNT = 1024
GRID = 32
TILE_SIZE = IMAGE_SIZE // GRID  # 128
SPRITE_CACHE_DIR = "/content/drive/MyDrive/pokemon_sprites"
OUTPUT_FILE = "all_colors_with_pokemon.png"

# Duplicate handling: quantized bins
# 32 bins per channel -> 32^3 = 32768 bins
BIN_SHIFT = 3  # 8 values per bin
BINS_PER_CH = 256 >> BIN_SHIFT  # 32
NBINS = BINS_PER_CH * BINS_PER_CH * BINS_PER_CH  # 32768
MAX_BIN_RING = 6  # expand search in bin-space (0..6) before global fallback

# Logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)


def log_mem(stage: str):
    vm = psutil.virtual_memory()
    logger.info(f"[{stage}] RAM {vm.used//1024//1024}MB/{vm.total//1024//1024}MB")


# -----------------------------
# Sprite fetch & cache
# -----------------------------

def fetch_sprites() -> tuple[np.ndarray, np.ndarray]:
    """Return (positions, orig_colors) for fully opaque pixels only.

    - Sprites are cached to Google Drive.
    - We use NEAREST scaling for pixel-art crispness.

    Note: `.convert('RGBA')` is only done if the source image isn't already RGBA.
    It's negligible compared to network + resize, and it guarantees we have an alpha channel.
    """
    os.makedirs(SPRITE_CACHE_DIR, exist_ok=True)
    pos_list: list[int] = []
    col_list: list[int] = []

    for pid in tqdm(range(1, POKE_COUNT + 1), desc="Sprites"):
        cache_path = os.path.join(SPRITE_CACHE_DIR, f"{pid}.png")
        img = None

        if os.path.exists(cache_path):
            img = Image.open(cache_path)
        else:
            try:
                r = requests.get(f"https://pokeapi.co/api/v2/pokemon/{pid}/", timeout=10)
                r.raise_for_status()
                data = r.json()
                url = data["sprites"]["front_default"]
                if not url:
                    continue
                img = Image.open(BytesIO(requests.get(url, timeout=15).content))
                img.save(cache_path)
            except Exception as e:
                logger.warning(f"Sprite {pid} fetch failed: {e}")
                continue

        if img is None:
            continue

        if img.mode != "RGBA":
            img = img.convert("RGBA")

        img = img.resize((TILE_SIZE, TILE_SIZE), Image.NEAREST)
        arr = np.asarray(img, dtype=np.uint8)

        # fully opaque only (keeps edges crisp; no anti-aliased semi-alpha pixels)
        mask = arr[..., 3] == 255
        ys, xs = np.nonzero(mask)

        ids = (
            (arr[..., 0].astype(np.uint32) << 16)
            | (arr[..., 1].astype(np.uint32) << 8)
            | (arr[..., 2].astype(np.uint32))
        )

        row, col = divmod(pid - 1, GRID)
        y0, x0 = row * TILE_SIZE, col * TILE_SIZE

        for y, x in zip(ys, xs):
            pos_list.append(int((y0 + y) * IMAGE_SIZE + (x0 + x)))
            col_list.append(int(ids[y, x]))

    log_mem("After sprite fetch")
    return np.asarray(pos_list, dtype=np.int32), np.asarray(col_list, dtype=np.uint32)


# -----------------------------
# Bin pool construction
# -----------------------------

@njit
def _bin_of_color(c: np.uint32) -> np.int32:
    r = (c >> 16) & 0xFF
    g = (c >> 8) & 0xFF
    b = c & 0xFF
    br = r >> BIN_SHIFT
    bg = g >> BIN_SHIFT
    bb = b >> BIN_SHIFT
    return (br << (2 * 5)) + (bg << 5) + bb  # 5 bits per component (0..31)


@njit
def build_bin_pool_from_perm(perm: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Group a color permutation into contiguous per-bin segments.

    Returns:
        grouped: uint32[N] colors, grouped by bin
        start:   int32[NBINS] inclusive start offset
        end:     int32[NBINS] exclusive end offset

    Order inside each bin inherits the randomness of `perm`.
    """
    counts = np.zeros(NBINS, dtype=np.int32)
    n = perm.shape[0]

    for i in range(n):
        b = _bin_of_color(perm[i])
        counts[b] += 1

    start = np.empty(NBINS, dtype=np.int32)
    end = np.empty(NBINS, dtype=np.int32)
    s = 0
    for b in range(NBINS):
        start[b] = s
        s += counts[b]
        end[b] = s

    # next write pointer per bin
    nxt = start.copy()
    grouped = np.empty(n, dtype=np.uint32)

    for i in range(n):
        c = perm[i]
        b = _bin_of_color(c)
        j = nxt[b]
        grouped[j] = c
        nxt[b] = j + 1

    return grouped, start, end


# -----------------------------
# Embedding with bin-based duplicate remap
# -----------------------------

@njit
def embed_swaps_binpool(
    perm: np.ndarray,
    pos_of_color: np.ndarray,
    positions: np.ndarray,
    orig_colors: np.ndarray,
    grouped: np.ndarray,
    start: np.ndarray,
    end: np.ndarray,
) -> None:
    """Swap desired colors into sprite positions.

    - First occurrence of a sprite RGB uses the exact color.
    - Duplicates allocate a new (unused-for-sprite) color from nearby quantized bins.

    This preserves exact global coverage because `perm` stays a permutation; we only swap.
    """
    used = np.zeros(TOTAL_PIXELS, dtype=np.bool_)

    # cursor per bin into `grouped`
    cur = start.copy()

    # global cursor for last-resort fallback
    global_cur = 0

    n = positions.shape[0]
    for i in range(n):
        p = positions[i]
        orig = orig_colors[i]

        # Choose c: exact if not used yet by sprite
        if not used[orig]:
            c = orig
            used[c] = True
        else:
            # Allocate near-by unused color from bin pool
            # Determine target bin from original
            b0 = _bin_of_color(orig)
            br0 = (b0 >> 10) & 31
            bg0 = (b0 >> 5) & 31
            bb0 = b0 & 31

            found = False

            # Expand ring in bin-space
            for ring in range(MAX_BIN_RING + 1):
                # iterate cube shell in bin coords (Chebyshev ring)
                rmin = br0 - ring
                rmax = br0 + ring
                gmin = bg0 - ring
                gmax = bg0 + ring
                bmin = bb0 - ring
                bmax = bb0 + ring

                if rmin < 0:
                    rmin = 0
                if gmin < 0:
                    gmin = 0
                if bmin < 0:
                    bmin = 0
                if rmax > 31:
                    rmax = 31
                if gmax > 31:
                    gmax = 31
                if bmax > 31:
                    bmax = 31

                for rr in range(rmin, rmax + 1):
                    for gg in range(gmin, gmax + 1):
                        for bb in range(bmin, bmax + 1):
                            if ring != 0:
                                # keep only the shell
                                if (abs(rr - br0) < ring and abs(gg - bg0) < ring and abs(bb - bb0) < ring):
                                    continue

                            b = (rr << 10) + (gg << 5) + bb
                            # pop next unused from this bin
                            while cur[b] < end[b]:
                                cand = grouped[cur[b]]
                                cur[b] += 1
                                if not used[cand]:
                                    c = cand
                                    used[c] = True
                                    found = True
                                    break
                            if found:
                                break
                        if found:
                            break
                    if found:
                        break
                if found:
                    break

            if not found:
                # Global fallback: advance through grouped until we find unused
                while global_cur < grouped.shape[0] and used[grouped[global_cur]]:
                    global_cur += 1
                if global_cur >= grouped.shape[0]:
                    # should never happen
                    c = 0
                else:
                    c = grouped[global_cur]
                    used[c] = True
                    global_cur += 1

        # Swap c into p
        src = pos_of_color[c]
        tmp = perm[p]
        perm[p] = c
        perm[src] = tmp
        pos_of_color[c] = p
        pos_of_color[tmp] = src


# -----------------------------
# Main
# -----------------------------

def main():
    log_mem("Start")

    # 1) Shuffle all colors (GPU if possible)
    if USE_GPU:
        try:
            rng = cp.random.RandomState(SEED)
            perm = cp.asnumpy(rng.permutation(cp.arange(TOTAL_PIXELS, dtype=cp.uint32)))
            log_mem("After GPU shuffle")
        except Exception as e:
            logger.warning(f"GPU shuffle failed, CPU fallback: {e}")
            perm = np.random.default_rng(SEED).permutation(TOTAL_PIXELS).astype(np.uint32)
            log_mem("After CPU shuffle")
    else:
        perm = np.random.default_rng(SEED).permutation(TOTAL_PIXELS).astype(np.uint32)
        log_mem("After CPU shuffle")

    # 2) Reverse lookup: position of each color in perm
    pos_of_color = np.empty(TOTAL_PIXELS, dtype=np.int32)
    for idx in range(TOTAL_PIXELS):
        pos_of_color[int(perm[idx])] = idx
    log_mem("After lookup build")

    # 3) Sprites
    positions, orig_colors = fetch_sprites()
    logger.info(f"Sprite pixels: {positions.size}")

    # 4) Build bin pools from the *random* permutation to get randomized candidates per bin
    logger.info("Building quantized bin pool from perm (Numba)…")
    grouped, start, end = build_bin_pool_from_perm(perm)
    log_mem("After bin pool")

    # 5) Embed (swap) with bin-based duplicate remap
    logger.info("Embedding sprites with bin-pool remap (Numba)…")
    embed_swaps_binpool(perm, pos_of_color, positions, orig_colors, grouped, start, end)
    log_mem("After embedding")

    # 6) Build RGBA buffer (RGB from perm, alpha background 128; sprite pixels 255)
    buf = np.empty((TOTAL_PIXELS, 4), dtype=np.uint8)
    buf[:, 0] = (perm >> 16) & 0xFF
    buf[:, 1] = (perm >> 8) & 0xFF
    buf[:, 2] = perm & 0xFF
    buf[:, 3] = 128
    buf[positions, 3] = 255
    log_mem("After buffer build")

    # 7) Optional: verify unique color count (expensive; enable if you want)
    # flat_ids = (buf[:,0].astype(np.uint32)<<16) | (buf[:,1].astype(np.uint32)<<8) | buf[:,2].astype(np.uint32)
    # u = np.unique(flat_ids).size
    # logger.info(f"Unique RGB count: {u}")

    img = Image.fromarray(buf.reshape((IMAGE_SIZE, IMAGE_SIZE, 4)), "RGBA")
    img.save(OUTPUT_FILE)
    log_mem("Done")


if __name__ == "__main__":
    main()
Alpha channel
ACJ
79 entries
Date: January 8, 2026
Colors: 16,777,216
Pixels: 16,777,216
Dimensions: 4,096 × 4,096
Bytes: 55,525,129