Pokédex 1–1024 v1

Which Pokémon is this?

Made with Chatgpt, Google Gemini, Google Colab and Google Drive. (Not sponsored.)

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Comprehensive script for generating an image with all 16,777,216 unique 24-bit RGB colors,
while embedding Pokémon sprites that closely resemble their originals.

Key Features:
- **All Colors Exactly Once:** Ensures every possible 24-bit RGB color is present exactly once.
- **GPU-Accelerated Shuffle:** Utilizes CuPy for fast initial color permutation if a GPU is available,
  with an automatic fallback to NumPy on CPU if not.
- **Optimized Sprite Embedding:** A Numba JIT-accelerated core loop embeds sprite pixels:
    - Prioritizes using the sprite's exact original color if available.
    - If the original color is taken, it searches for the *closest available* color in RGB space
      (using Euclidean distance) within a specified radius (MAX_RAD=3 by default).
      This greatly improves sprite visual fidelity compared to a simple "first available" search.
    - Performs an efficient 1:1 color swap, ensuring the total color count remains exact.
- **Sprite Caching:** Fetches Pokémon sprites from PokeAPI, resizes them to 128x128,
  and caches them to Google Drive to speed up subsequent runs.
- **Memory Management Logging:** Provides insights into RAM usage at different stages.
- **Robust Google Drive Integration:** Mounts Google Drive to the standard `/content/drive` path
  and correctly sets up persistent caches within 'My Drive'.

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

Usage in Google Colab:
1. Save this script as a .py file (e.g., `pokemon_color_image.py`) in your Google Drive.
   A good location might be `My Drive/Colab_Scripts/pokemon_color_image.py`.
2. In a Colab notebook cell, first mount your Google Drive:
   This step will prompt for authorization.
   Example code for mounting:
   from google.colab import drive
   drive.mount('/content/drive')
   print("Google Drive mounted successfully!")
3. In a *separate* Colab notebook cell, run your script:
   Example code for running the script:
   import os
   # Ensure NUMBA_CACHE_DIR is set *before* your script starts to compile Numba functions.
   # It points to a directory *within* your mounted Google Drive.
   os.environ['NUMBA_CACHE_DIR'] = "/content/drive/MyDrive/numba_cache"

   # Define the path to your script on Google Drive
   script_path = "/content/drive/MyDrive/Colab_Scripts/pokemon_color_image.py"

   # Run the script using the Python interpreter.
   !python "$script_path"
   The script will then create necessary directories (`numba_cache`, `pokemon_sprites`) inside your Google Drive if they don't exist.
"""
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, prange # prange allows for parallel compilation within Numba if configured

# --- Configuration Section ---
# Standard Google Drive mount point in Colab
GOOGLE_DRIVE_MOUNT_POINT = "/content/drive"

# Numba Caching Setup: Specify where Numba should store compiled functions.
# This path is *relative to your mounted Google Drive's "My Drive"*
# Example: /content/drive/MyDrive/numba_cache
NUMBA_CACHE_ROOT = os.path.join(GOOGLE_DRIVE_MOUNT_POINT, "MyDrive", "numba_cache")

# Image Generation Parameters
IMAGE_SIZE = 4096 # The final image will be IMAGE_SIZE x IMAGE_SIZE pixels
TOTAL_PIXELS = IMAGE_SIZE * IMAGE_SIZE # Total possible 24-bit colors: 2^24 = 16,777,216
SEED = 42 # Random seed for initial color shuffling (for reproducibility)

# Sprite Parameters
POKE_COUNT = 1024 # Number of Pokémon sprites to fetch (from 1 to POKE_COUNT)
TILE_SIZE = IMAGE_SIZE // 32 # Size (width/height) of each individual sprite tile in pixels (e.g., 4096/32 = 128)
# Directory within your mounted Google Drive to cache fetched Pokémon sprites.
# Example: /content/drive/MyDrive/pokemon_sprites
SPRITE_CACHE_DIR = os.path.join(GOOGLE_DRIVE_MOUNT_POINT, "MyDrive", "pokemon_sprites")

# Output File
OUTPUT_FILE = "all_colors_with_pokemon.png" # Name of the final output image file.
# This will be saved in the Colab instance's /content/ directory, which you can download.

# Color Remapping Parameters
MAX_RAD = 3 # Maximum Chebyshev search radius for finding nearest free colors.
            # A larger radius means potentially better color matches but more computation.

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

# --- GPU (CuPy) Detection ---
USE_GPU = True
try:
    import cupy as cp
    # Test if CuPy can allocate memory to ensure a CUDA device is actually available.
    _ = cp.array([1])
except (ImportError, cp.cuda.runtime.CudaError):
    USE_GPU = False
    print("CuPy not available or CUDA device not found. Falling back to CPU for shuffling.")

# --- Helper Functions ---

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


def precompute_euclidean_sorted_offsets(max_radius: int) -> np.ndarray:
    """
    Precomputes all (dr,dg,db) offsets within a Chebyshev cube of 'max_radius'
    and sorts them by their squared Euclidean distance from the origin.
    This ensures that when iterating through these offsets, colors closer in RGB space
    are checked first, leading to better visual matches for the sprites.
    """
    offsets_set = set() # Use a set to store unique (dr, dg, db) tuples
    for dr in range(-max_radius, max_radius + 1):
        for dg in range(-max_radius, max_radius + 1):
            for db in range(-max_radius, max_radius + 1):
                if dr == 0 and dg == 0 and db == 0:
                    continue # The (0,0,0) offset corresponds to the original color, handled separately.
                offsets_set.add((dr, dg, db))

    # Convert to a list of tuples and sort them based on their squared Euclidean distance.
    # Sorting ensures we check closest neighbors first.
    sorted_offsets = sorted(list(offsets_set), key=lambda x: x[0]**2 + x[1]**2 + x[2]**2)
    return np.array(sorted_offsets, dtype=np.int16)


def fetch_and_process_sprites(
    poke_count: int, tile_size: int, image_size: int, cache_dir: str
) -> tuple[np.ndarray, np.ndarray]:
    """
    Fetches Pokémon sprites from PokeAPI, caches them, resizes them, and
    extracts opaque pixel positions and their original 24-bit colors.
    
    Args:
        poke_count (int): The number of Pokémon sprites to fetch (by ID from 1 to poke_count).
        tile_size (int): The target size (width and height) for each sprite.
        image_size (int): The overall dimension of the final output image.
        cache_dir (str): Directory path to cache the downloaded sprites.

    Returns:
        tuple[np.ndarray, np.ndarray]: A tuple containing two NumPy arrays:
            - sprite_pixel_positions (np.int32): Global 1D indices in the final image
                                                  where sprite pixels will be placed.
            - original_colors (np.uint32): The original 24-bit integer color of each
                                           corresponding sprite pixel.
    """
    # Ensure the sprite cache directory exists on Google Drive
    os.makedirs(cache_dir, exist_ok=True)
    
    sprite_data_collector = [] # Temporarily store (global_pixel_index, color_id) tuples

    for pid in tqdm(range(1, poke_count + 1), desc="Fetching and processing sprites"):
        path = os.path.join(cache_dir, f"{pid}.png")
        img = None
        source_log = "" # To store the source of the sprite

        # Try loading from cache first
        if os.path.exists(path):
            try:
                img = Image.open(path)
                source_log = "from cache (Drive)"
            except IOError:
                logger.warning(f"Could not open cached sprite {path}. Attempting to re-fetch.")
                img = None # Force re-fetch if cache file is corrupt

        # If not loaded from cache or cache failed, fetch from API
        if img is None:
            try:
                # Fetch Pokémon data from PokeAPI
                r = requests.get(f"https://pokeapi.co/api/v2/pokemon/{pid}/", timeout=15)
                r.raise_for_status() # Raise an exception for HTTP errors (4xx or 5xx)
                data = r.json()
                
                # Get the URL for the front_default sprite
                url = data.get('sprites', {}).get('front_default')

                if not url:
                    logger.debug(f"No 'front_default' sprite URL found for Pokémon ID {pid}. Skipping.")
                    continue
                
                # Fetch sprite image content
                img_content = requests.get(url, timeout=15).content
                img = Image.open(BytesIO(img_content))
                img.save(path) # Save fetched image to cache
                source_log = "from PokeAPI"

            except requests.exceptions.RequestException as e:
                logger.warning(f"Network error fetching Pokémon ID {pid}: {e}. Skipping.")
                continue
            except KeyError:
                logger.warning(f"Unexpected data structure from PokeAPI for Pokémon ID {pid}. Skipping.")
                continue
            except Exception as e: # Catch any other unexpected errors during fetch/save
                logger.warning(f"An unexpected error occurred while fetching/saving Pokémon ID {pid}: {e}. Skipping.")
                continue

        if img is None: # If after all attempts, image is still not loaded
            continue

        # Log the source of the sprite, especially for the first one
        if pid == 1:
            logger.info(f"Pokémon ID {pid} sprite loaded {source_log}.")

        # Convert to RGBA and resize
        if img.mode != 'RGBA':
            img = img.convert('RGBA')
        
        # Resize using NEAREST to maintain pixel art quality
        img = img.resize((tile_size, tile_size), Image.NEAREST)
        arr = np.array(img)

        # Identify opaque pixels (alpha channel == 255)
        mask = arr[..., 3] == 255
        ys, xs = np.nonzero(mask) # Get local (y, x) coordinates of opaque pixels

        # Calculate the top-left corner of the current Pokémon's tile
        row, col = divmod(pid - 1, image_size // tile_size)
        y0, x0 = row * tile_size, col * tile_size

        # Collect global pixel positions and their original colors
        for y_local, x_local in zip(ys, xs):
            global_pixel_idx = (y0 + y_local) * image_size + (x0 + x_local)
            # Extract RGB components and combine into a single uint32 color ID
            r, g, b = arr[y_local, x_local, :3].astype(np.uint32)
            color_id = (r << 16) | (g << 8) | b
            sprite_data_collector.append((global_pixel_idx, color_id))
    
    if not sprite_data_collector:
        logger.warning("No opaque sprite pixels were collected. The output image might not contain sprites.")
        return np.array([], dtype=np.int32), np.array([], dtype=np.uint32)

    # Convert the collected list of tuples into two separate NumPy arrays
    positions = np.array([item[0] for item in sprite_data_collector], dtype=np.int32)
    orig_colors = np.array([item[1] for item in sprite_data_collector], dtype=np.uint32)

    log_mem("After sprite fetch and processing")
    return positions, orig_colors

@njit(cache=True) # Numba caching enabled. NUMBA_CACHE_DIR is set in the environment.
def embed_swaps_numba(
    perm: np.ndarray,
    pos_of_color: np.ndarray,
    positions: np.ndarray,
    orig_colors: np.ndarray,
    euclidean_sorted_offsets: np.ndarray
) -> np.ndarray:
    """
    Numba JIT-accelerated function to embed sprite colors into the permutation
    array using a nearest-neighbor search for occupied colors.
    Modifies perm and pos_of_color in-place.
    
    This version prioritizes the closest available color in RGB space for better
    sprite resemblance.
    """
    # Create a boolean array to track which of the 16,777,216 colors have been "used"
    # This array is directly indexed by the 24-bit color integer.
    used_colors = np.zeros(TOTAL_PIXELS, dtype=np.bool_)

    n_sprite_pixels = positions.shape[0]

    # Iterate over each sprite pixel. prange enables Numba's parallel execution for this loop
    # if Numba is configured to use multiple threads (e.g., via OMP_NUM_THREADS).
    for i in prange(n_sprite_pixels):
        p = positions[i]             # The global pixel index in the final image where this sprite pixel belongs
        orig_rgb_int = orig_colors[i] # The original 24-bit integer color of the sprite pixel

        chosen_color = -1 # Initialize with an invalid value

        # Attempt 1: Try to use the exact original color of the sprite pixel.
        if not used_colors[orig_rgb_int]:
            chosen_color = orig_rgb_int
            used_colors[chosen_color] = True
        else:
            # Attempt 2: Search for the closest available color in RGB space.
            # We convert the original 24-bit integer color back to its R, G, B components.
            r0 = (orig_rgb_int >> 16) & 0xFF
            g0 = (orig_rgb_int >> 8) & 0xFF
            b0 = orig_rgb_int & 0xFF
            
            # Iterate through pre-computed offsets. Since they are sorted by Euclidean distance,
            # the first available color found will be the closest.
            for j in range(euclidean_sorted_offsets.shape[0]):
                dr, dg, db = euclidean_sorted_offsets[j]
                
                r = r0 + dr # Candidate Red component
                g = g0 + dg # Candidate Green component
                b = b0 + db # Candidate Blue component
                
                # Check if the candidate color's components are within valid RGB bounds (0-255).
                if 0 <= r < 256 and 0 <= g < 256 and 0 <= b < 256:
                    # Reconstruct the 24-bit integer color ID for the candidate.
                    candidate_color = (r << 16) | (g << 8) | b
                    # Check if this candidate color has not been used by another sprite pixel yet.
                    if not used_colors[candidate_color]:
                        chosen_color = candidate_color # This is our chosen color
                        used_colors[chosen_color] = True # Mark it as used
                        break # Found the closest available neighbor, so we can stop searching.
            
            if chosen_color == -1: # If chosen_color is still -1, it means no neighbor within MAX_RAD was found.
                # Fallback: Search linearly for *any* unused color.
                # This is a last resort to guarantee the 1:1 mapping constraint is met.
                # This should be very rare with a reasonable MAX_RAD and sufficient free colors.
                for cand_linear in range(TOTAL_PIXELS): # Iterate through all possible colors (0 to 2^24-1)
                    if not used_colors[cand_linear]:
                        chosen_color = cand_linear
                        used_colors[chosen_color] = True
                        break
                # If after this loop chosen_color is still -1, it's a critical error implying
                # that all 16M colors are already used, which should not be possible unless
                # the number of sprite pixels exceeds TOTAL_PIXELS (which it does not).
                if chosen_color == -1:
                    pass # Placeholder: In a production system, consider raising an error or logging a critical failure.

        # Now, perform the crucial 1:1 swap in the permutation array (`perm`) and update
        # the reverse lookup table (`pos_of_color`).
        
        # 1. Get the current position (index in `perm`) where the `chosen_color` currently resides.
        src_pos_of_chosen_color = pos_of_color[chosen_color]
        
        # 2. Get the color that is currently at the target sprite pixel position `p`.
        color_at_sprite_target_pos = perm[p]

        # 3. Place the `chosen_color` (the one we want for the sprite pixel) at `p`.
        perm[p] = chosen_color
        
        # 4. Move the `color_at_sprite_target_pos` (the one that was displaced from `p`)
        #    to the position where `chosen_color` used to be.
        perm[src_pos_of_chosen_color] = color_at_sprite_target_pos

        # 5. Update the reverse lookup table (`pos_of_color`) to reflect these swaps.
        #    The chosen_color is now at 'p'.
        pos_of_color[chosen_color] = p
        #    The displaced color (color_at_sprite_target_pos) is now at src_pos_of_chosen_color.
        pos_of_color[color_at_sprite_target_pos] = src_pos_of_chosen_color
    
    return used_colors # Returning the mask can be useful for debugging or further analysis


def initialize_permutation_and_lookup(total_pixels: int, seed: int) -> tuple[np.ndarray, np.ndarray]:
    """
    Initializes the 24-bit color permutation array and its reverse lookup table.
    Prioritizes GPU shuffling with CuPy for speed, falling back to CPU NumPy if CuPy is unavailable or fails.
    """
    if USE_GPU:
        try:
            logger.info("Attempting initial color shuffle using CuPy on GPU...")
            rng = cp.random.RandomState(seed)
            # Generate the permutation on GPU, then transfer to CPU for Numba processing.
            perm_gpu = rng.permutation(cp.arange(total_pixels, dtype=cp.uint32))
            perm = cp.asnumpy(perm_gpu) 
            del perm_gpu # Explicitly release CuPy array to free GPU memory
            log_mem("After GPU shuffle (transferred to CPU)")
        except cp.cuda.runtime.CudaError as e:
            logger.warning(f"CuPy CUDA error during shuffle ({e}). Falling back to CPU shuffle.")
            perm = np.random.default_rng(seed).permutation(total_pixels).astype(np.uint32)
            log_mem("After CPU shuffle (fallback)")
        except Exception as e:
            logger.warning(f"Unexpected error with CuPy shuffle ({e}). Falling back to CPU shuffle.")
            perm = np.random.default_rng(seed).permutation(total_pixels).astype(np.uint32)
            log_mem("After CPU shuffle (fallback)")
    else:
        logger.info("Using CPU for initial color shuffle (CuPy not enabled or failed).")
        perm = np.random.default_rng(seed).permutation(total_pixels).astype(np.uint32)
        log_mem("After CPU shuffle")

    # Build reverse lookup: pos_of_color[color_id] = index_in_perm_array
    # This array allows O(1) lookup to find where a specific color is currently located in the `perm` array.
    pos_of_color = np.empty(total_pixels, dtype=np.int32)
    # Using prange for this loop as well, as it can be parallelized by Numba.
    for idx in prange(total_pixels):
        pos_of_color[perm[idx]] = idx
    
    log_mem("After lookup table build")
    return perm, pos_of_color


def build_and_save_image(perm: np.ndarray, positions: np.ndarray, image_size: int, output_file: str):
    """
    Constructs the final RGBA image buffer from the modified permutation array
    and saves it to a PNG file.
    """
    logger.info("Building final RGBA image buffer...")
    
    # Initialize the alpha channel: 128 for background pixels (semi-transparent),
    # and 255 for sprite pixels (fully opaque).
    alpha_channel = np.full(TOTAL_PIXELS, 128, dtype=np.uint8)
    alpha_channel[positions] = 255 # Set alpha for all sprite pixel positions to 255

    # Create the final RGBA buffer. The `perm` array contains the 24-bit color IDs.
    # We extract the R, G, B components using bitwise operations.
    buf = np.empty((TOTAL_PIXELS, 4), dtype=np.uint8)
    buf[:, 0] = (perm >> 16) & 0xFF  # Red channel (most significant 8 bits)
    buf[:, 1] = (perm >> 8) & 0xFF   # Green channel (middle 8 bits)
    buf[:, 2] = perm & 0xFF          # Blue channel (least significant 8 bits)
    buf[:, 3] = alpha_channel        # Alpha channel

    log_mem("After RGBA buffer construction")

    # Reshape the 1D buffer into a 2D image array (IMAGE_SIZE x IMAGE_SIZE x 4)
    # and save it using PIL (Pillow).
    logger.info(f"Saving final image to {output_file}...")
    img = Image.fromarray(buf.reshape((image_size, image_size, 4)), 'RGBA')
    img.save(output_file)
    logger.info(f"Image saved successfully to {output_file}.")
    log_mem("Done saving image")


def main():
    log_mem("Script Start")

    # IMPORTANT: The Google Drive mounting step `drive.mount('/content/drive')`
    # MUST be executed in a separate Colab notebook cell before running this script
    # via `!python`. This script assumes Drive is already mounted at /content/drive.

    # Ensure the Numba cache and sprite cache directories exist on Google Drive.
    # os.makedirs will create them if they don't, but will not delete existing contents.
    # These paths are within the *mounted* Google Drive, so they operate on your Drive.
    os.makedirs(NUMBA_CACHE_ROOT, exist_ok=True)
    os.makedirs(SPRITE_CACHE_DIR, exist_ok=True)
    logger.info(f"Numba cache directory ensured: {NUMBA_CACHE_ROOT}")
    logger.info(f"Sprite cache directory ensured: {SPRITE_CACHE_DIR}")

    # Step 1: Precompute color offsets for optimized nearest-neighbor search
    euclidean_sorted_offsets = precompute_euclidean_sorted_offsets(MAX_RAD)
    logger.info(f"Precomputed {euclidean_sorted_offsets.shape[0]} Euclidean-sorted offsets for color search (radius up to {MAX_RAD}).")

    # Step 2: Initialize the main color permutation array and its reverse lookup table.
    perm, pos_of_color = initialize_permutation_and_lookup(TOTAL_PIXELS, SEED)

    # Step 3: Fetch, process, and collect data for Pokémon sprites.
    positions, orig_colors = fetch_and_process_sprites(POKE_COUNT, TILE_SIZE, IMAGE_SIZE, SPRITE_CACHE_DIR)
    
    if positions.size == 0:
        logger.error("No sprite pixels were collected. Cannot proceed with embedding. Exiting.")
        return

    # Step 4: Embed sprite colors into the permutation array.
    logger.info(f"Embedding {positions.shape[0]} sprite pixels with optimized nearest-neighbor remap...")
    _ = embed_swaps_numba(perm, pos_of_color, positions, orig_colors, euclidean_sorted_offsets)
    log_mem("After embedding sprites into permutation")

    # Step 5: Construct the final RGBA image buffer and save it to a file.
    build_and_save_image(perm, positions, IMAGE_SIZE, OUTPUT_FILE)

    logger.info("Script execution completed successfully!")

if __name__ == '__main__':
    main()

Pokémon, Pokémon names, and Pokémon sprites are trademarks of The Pokémon Company.