#!/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()
