Buckets

A question from our most active contributor Samuel35 prompted me to revisit an idea I had very early in the project (but had abandoned for technical reasons at the time), and I, in turn, prompted Chatgpt to turn it into Python code to run on a Google Colab notebook:

#!/usr/bin/env python3
"""
allRGB image-guided bucket completion

Creates a 4096 x 4096 PNG containing every 24-bit RGB color exactly once.
The input image is used as a visual guide. Unique colors from the source image
are placed first, then every remaining RGB color is grouped into 16 palette
buckets and assigned to remaining pixels whose guide color belongs to the same
or nearest available bucket.

Dependencies:
    pip install pillow numpy

CLI:
    python allrgb_image_guided_buckets.py input.png output.png

Google Colab:
    Run this cell/file without command-line arguments. It will open an upload
    dialog and then download the generated PNG.

Local web server:
    Import process_image() from this file and call it from your route/worker.
    Do not run this directly inside a blocking request handler for production;
    it is CPU/RAM intensive and is better handled by a job queue.
"""

from __future__ import annotations

import argparse
import os
import sys
import time
from pathlib import Path
from typing import Iterable, Literal, Optional, Tuple

import numpy as np
from PIL import Image, ImageOps

# Pillow sometimes protects against very large images. For this use case, large
# image files are expected.
Image.MAX_IMAGE_PIXELS = None

CANVAS_SIZE = 4096
TOTAL_COLORS = CANVAS_SIZE * CANVAS_SIZE  # 16,777,216 == 2^24

# Integer Rec.709-ish luminance coefficients scaled to 256.
# 54 + 183 + 19 = 256.
LUMA_R = 54
LUMA_G = 183
LUMA_B = 19

# Used when the adaptive palette produces fewer than 16 distinct colors.
# The adaptive colors are kept first; these only complete the bucket set.
FALLBACK_16 = np.array(
    [
        (0, 0, 0),
        (255, 255, 255),
        (128, 128, 128),
        (255, 0, 0),
        (0, 255, 0),
        (0, 0, 255),
        (255, 255, 0),
        (0, 255, 255),
        (255, 0, 255),
        (255, 128, 0),
        (128, 0, 255),
        (255, 0, 128),
        (128, 64, 0),
        (0, 128, 128),
        (0, 64, 128),
        (128, 128, 0),
    ],
    dtype=np.uint8,
)


try:
    RESAMPLING_LANCZOS = Image.Resampling.LANCZOS
    RESAMPLING_NEAREST = Image.Resampling.NEAREST
    QUANTIZE_MEDIANCUT = Image.Quantize.MEDIANCUT
except AttributeError:  # Older Pillow fallback.
    RESAMPLING_LANCZOS = Image.LANCZOS
    RESAMPLING_NEAREST = Image.NEAREST
    QUANTIZE_MEDIANCUT = Image.MEDIANCUT


FitMode = Literal["cover", "contain", "stretch"]
PaletteMode = Literal["adaptive", "fixed"]


def log(message: str, quiet: bool = False) -> None:
    if not quiet:
        print(message, flush=True)


def parse_hex_rgb(value: str) -> Tuple[int, int, int]:
    value = value.strip()
    if value.startswith("#"):
        value = value[1:]
    if len(value) != 6:
        raise ValueError("Expected a hex color like #ffffff")
    return (int(value[0:2], 16), int(value[2:4], 16), int(value[4:6], 16))


def normalize_to_rgb(path: str | Path, alpha_background: Tuple[int, int, int]) -> Image.Image:
    """Load an image, respect EXIF orientation, and flatten transparency to RGB."""
    img = Image.open(path)
    img = ImageOps.exif_transpose(img)

    has_alpha = img.mode in {"RGBA", "LA"} or "transparency" in img.info
    if has_alpha:
        rgba = img.convert("RGBA")
        bg = Image.new("RGBA", rgba.size, (*alpha_background, 255))
        bg.alpha_composite(rgba)
        return bg.convert("RGB")

    return img.convert("RGB")


def fit_image(
    img: Image.Image,
    size: Tuple[int, int],
    mode: FitMode,
    resample,
    background: Tuple[int, int, int],
) -> Image.Image:
    """Return an RGB image fitted to size."""
    if mode == "stretch":
        return img.resize(size, resample=resample)

    if mode == "cover":
        return ImageOps.fit(img, size, method=resample, centering=(0.5, 0.5))

    if mode == "contain":
        contained = ImageOps.contain(img, size, method=resample)
        canvas = Image.new("RGB", size, background)
        x = (size[0] - contained.size[0]) // 2
        y = (size[1] - contained.size[1]) // 2
        canvas.paste(contained, (x, y))
        return canvas

    raise ValueError(f"Unknown fit mode: {mode}")


def rgb_flat_to_ids(rgb_flat: np.ndarray) -> np.ndarray:
    """Convert flat RGB rows to 24-bit integer ids: 0xRRGGBB."""
    rgb = rgb_flat.astype(np.uint32, copy=False)
    return (rgb[:, 0] << 16) | (rgb[:, 1] << 8) | rgb[:, 2]


def rgb_image_to_ids(img: Image.Image) -> np.ndarray:
    arr = np.asarray(img, dtype=np.uint8)
    return rgb_flat_to_ids(arr.reshape(-1, 3))


def ids_to_rgb_array(ids: np.ndarray, size: int = CANVAS_SIZE) -> np.ndarray:
    rgb = np.empty((ids.size, 3), dtype=np.uint8)
    rgb[:, 0] = ((ids >> 16) & 255).astype(np.uint8)
    rgb[:, 1] = ((ids >> 8) & 255).astype(np.uint8)
    rgb[:, 2] = (ids & 255).astype(np.uint8)
    return rgb.reshape((size, size, 3))


def luma_from_rgb_flat(rgb_flat: np.ndarray) -> np.ndarray:
    rgb = rgb_flat.astype(np.uint16, copy=False)
    return (
        LUMA_R * rgb[:, 0]
        + LUMA_G * rgb[:, 1]
        + LUMA_B * rgb[:, 2]
    ).astype(np.uint16)


def luma_from_ids(ids: np.ndarray) -> np.ndarray:
    r = ((ids >> 16) & 255).astype(np.uint32)
    g = ((ids >> 8) & 255).astype(np.uint32)
    b = (ids & 255).astype(np.uint32)
    return (LUMA_R * r + LUMA_G * g + LUMA_B * b).astype(np.uint16)


def unique_rows_preserve_order(rows: Iterable[Tuple[int, int, int]]) -> list[Tuple[int, int, int]]:
    seen: set[Tuple[int, int, int]] = set()
    out: list[Tuple[int, int, int]] = []
    for row in rows:
        t = (int(row[0]), int(row[1]), int(row[2]))
        if t not in seen:
            seen.add(t)
            out.append(t)
    return out


def build_palette_16(guide_img: Image.Image, mode: PaletteMode, quiet: bool = False) -> np.ndarray:
    """
    Build the 16 bucket colors.

    adaptive: median-cut palette from a small guide image, completed with fallback
    colors if the guide has fewer than 16 distinct colors.

    fixed: deterministic 16-color reference palette.
    """
    if mode == "fixed":
        log("Using fixed 16-color bucket palette.", quiet)
        return FALLBACK_16.copy()

    small = guide_img.copy()
    small.thumbnail((512, 512), RESAMPLING_LANCZOS)
    q = small.quantize(colors=16, method=QUANTIZE_MEDIANCUT)
    raw_palette = q.getpalette() or []
    counts = q.getcolors(maxcolors=512) or []

    # Use most frequent quantized colors first.
    colors: list[Tuple[int, int, int]] = []
    for _count, index in sorted(counts, reverse=True):
        offset = int(index) * 3
        if offset + 2 < len(raw_palette):
            colors.append(
                (
                    int(raw_palette[offset]),
                    int(raw_palette[offset + 1]),
                    int(raw_palette[offset + 2]),
                )
            )

    # Complete with fixed reference colors if necessary.
    colors.extend(tuple(map(int, row)) for row in FALLBACK_16)
    unique = unique_rows_preserve_order(colors)[:16]

    if len(unique) != 16:
        raise RuntimeError("Could not construct a 16-color palette.")

    palette = np.array(unique, dtype=np.uint8)
    log("16 bucket palette:", quiet)
    for i, (r, g, b) in enumerate(palette):
        log(f"  {i:02d}: #{int(r):02x}{int(g):02x}{int(b):02x}", quiet)
    return palette


def nearest_palette_labels_for_rgb_flat(
    rgb_flat: np.ndarray,
    palette: np.ndarray,
    chunk_pixels: int,
    quiet: bool = False,
    label: str = "guide pixels",
) -> np.ndarray:
    n = rgb_flat.shape[0]
    labels = np.empty(n, dtype=np.uint8)

    pal = palette.astype(np.int32)
    pal_norm = np.einsum("ij,ij->i", pal, pal).astype(np.int32)

    t0 = time.time()
    for start in range(0, n, chunk_pixels):
        end = min(start + chunk_pixels, n)
        chunk = rgb_flat[start:end].astype(np.int32, copy=False)
        chunk_norm = np.einsum("ij,ij->i", chunk, chunk).astype(np.int32)
        dots = chunk @ pal.T
        dist = chunk_norm[:, None] + pal_norm[None, :] - 2 * dots
        labels[start:end] = np.argmin(dist, axis=1).astype(np.uint8)

        if not quiet and (start == 0 or end == n or (time.time() - t0) > 10):
            pct = 100.0 * end / n
            print(f"  Bucketed {label}: {pct:5.1f}%", flush=True)
            t0 = time.time()

    return labels


def nearest_palette_labels_for_ids(
    ids: np.ndarray,
    palette: np.ndarray,
    chunk_colors: int,
    quiet: bool = False,
    label: str = "colors",
) -> np.ndarray:
    n = ids.size
    labels = np.empty(n, dtype=np.uint8)

    pal = palette.astype(np.int32)
    pr = pal[:, 0]
    pg = pal[:, 1]
    pb = pal[:, 2]
    pal_norm = (pr * pr + pg * pg + pb * pb).astype(np.int32)

    t0 = time.time()
    for start in range(0, n, chunk_colors):
        end = min(start + chunk_colors, n)
        chunk = ids[start:end]

        r = ((chunk >> 16) & 255).astype(np.int32)
        g = ((chunk >> 8) & 255).astype(np.int32)
        b = (chunk & 255).astype(np.int32)

        chunk_norm = r * r + g * g + b * b
        dots = r[:, None] * pr[None, :] + g[:, None] * pg[None, :] + b[:, None] * pb[None, :]
        dist = chunk_norm[:, None] + pal_norm[None, :] - 2 * dots
        labels[start:end] = np.argmin(dist, axis=1).astype(np.uint8)

        if not quiet and n and (start == 0 or end == n or (time.time() - t0) > 10):
            pct = 100.0 * end / n
            print(f"  Bucketed {label}: {pct:5.1f}%", flush=True)
            t0 = time.time()

    return labels


def evenly_spaced_indices(length: int, take: int) -> np.ndarray:
    """Return 'take' indices spread through range(length)."""
    if take <= 0:
        return np.empty(0, dtype=np.int64)
    if take >= length:
        return np.arange(length, dtype=np.int64)
    return np.linspace(0, length - 1, take, dtype=np.int64)


def assign_sorted_by_luma(
    out_ids: np.ndarray,
    assigned_pos: np.ndarray,
    positions: np.ndarray,
    colors: np.ndarray,
    guide_luma_flat: np.ndarray,
    color_luma: Optional[np.ndarray] = None,
) -> None:
    """Assign colors to positions, matching dark-to-light within the given set."""
    if positions.size != colors.size:
        raise ValueError("positions and colors must have the same length")
    if colors.size == 0:
        return

    pos_order = np.argsort(guide_luma_flat[positions], kind="stable")
    if color_luma is None:
        color_luma = luma_from_ids(colors)
    color_order = np.argsort(color_luma, kind="stable")

    sorted_positions = positions[pos_order]
    sorted_colors = colors[color_order]
    out_ids[sorted_positions] = sorted_colors
    assigned_pos[sorted_positions] = True


def place_color_ids_by_buckets(
    *,
    out_ids: np.ndarray,
    assigned_pos: np.ndarray,
    guide_labels: np.ndarray,
    guide_luma_flat: np.ndarray,
    color_ids: np.ndarray,
    palette: np.ndarray,
    chunk_colors: int,
    quiet: bool,
    stage_name: str,
) -> None:
    """
    Place a unique set of color ids into still-unassigned pixels.

    First pass: exact bucket-to-bucket assignment.
    Fallback: any remaining colors are assigned to any remaining pixels by
    luminance. This guarantees completion even when bucket counts do not match.
    """
    if color_ids.size == 0:
        log(f"{stage_name}: no colors to place.", quiet)
        return

    available_before = int(np.count_nonzero(~assigned_pos))
    if color_ids.size > available_before:
        raise RuntimeError(
            f"{stage_name}: trying to place {color_ids.size:,} colors into "
            f"only {available_before:,} open pixels."
        )

    log(f"{stage_name}: bucket-matching {color_ids.size:,} colors.", quiet)
    color_labels = nearest_palette_labels_for_ids(
        color_ids,
        palette,
        chunk_colors=chunk_colors,
        quiet=quiet,
        label=stage_name,
    )
    color_luma_all = luma_from_ids(color_ids)
    color_was_assigned = np.zeros(color_ids.size, dtype=bool)

    for bucket in range(16):
        color_local = np.flatnonzero(color_labels == bucket)
        if color_local.size == 0:
            continue

        # Current open positions whose guide pixel wants this bucket.
        positions = np.flatnonzero((guide_labels == bucket) & (~assigned_pos))
        if positions.size == 0:
            continue

        n = min(color_local.size, positions.size)
        if n == 0:
            continue

        # Sort both sides by luminance, then take evenly if one side is bigger.
        # This avoids always taking the darkest or brightest colors from an
        # over-supplied bucket.
        colors_for_bucket = color_ids[color_local]
        luma_for_bucket = color_luma_all[color_local]

        color_order = np.argsort(luma_for_bucket, kind="stable")
        pos_order = np.argsort(guide_luma_flat[positions], kind="stable")

        color_take = evenly_spaced_indices(color_order.size, n)
        pos_take = evenly_spaced_indices(pos_order.size, n)

        chosen_color_local = color_local[color_order[color_take]]
        chosen_positions = positions[pos_order[pos_take]]
        chosen_colors = color_ids[chosen_color_local]

        # The chosen arrays are already dark-to-light because they are selected
        # from sorted orders.
        out_ids[chosen_positions] = chosen_colors
        assigned_pos[chosen_positions] = True
        color_was_assigned[chosen_color_local] = True

        log(
            f"  Bucket {bucket:02d}: placed {n:,} / {color_local.size:,} colors "
            f"into {positions.size:,} matching pixels.",
            quiet,
        )

    leftover = color_ids[~color_was_assigned]
    if leftover.size:
        open_positions = np.flatnonzero(~assigned_pos)
        if leftover.size > open_positions.size:
            raise RuntimeError(
                f"{stage_name}: fallback has {leftover.size:,} colors but only "
                f"{open_positions.size:,} pixels."
            )

        # Use all leftover colors. If there are more open positions than leftover
        # colors, choose positions spread over the remaining luminance range.
        pos_order = np.argsort(guide_luma_flat[open_positions], kind="stable")
        pos_take = evenly_spaced_indices(pos_order.size, leftover.size)
        chosen_positions = open_positions[pos_order[pos_take]]
        assign_sorted_by_luma(
            out_ids=out_ids,
            assigned_pos=assigned_pos,
            positions=chosen_positions,
            colors=leftover,
            guide_luma_flat=guide_luma_flat,
        )
        log(f"{stage_name}: fallback-placed {leftover.size:,} colors.", quiet)
    else:
        log(f"{stage_name}: no fallback needed.", quiet)


def place_first_visible_source_occurrences(
    *,
    out_ids: np.ndarray,
    assigned_pos: np.ndarray,
    source_mask: np.ndarray,
    placed_source_mask: np.ndarray,
    stamp_ids: np.ndarray,
    quiet: bool,
) -> int:
    """
    Place source colors where they first occur in a nearest-neighbor fitted guide.

    This gives the source palette a spatial anchor before bucket filling starts.
    """
    log("Placing first visible occurrence of each source color.", quiet)
    values, first_indices = np.unique(stamp_ids, return_index=True)
    is_source_color = source_mask[values]
    values = values[is_source_color].astype(np.uint32, copy=False)
    first_indices = first_indices[is_source_color]

    # If two colors somehow map to the same first index, this would be a problem,
    # but np.unique returns one first index per distinct value, so they are unique.
    out_ids[first_indices] = values
    assigned_pos[first_indices] = True
    placed_source_mask[values] = True

    placed = int(values.size)
    log(f"Placed {placed:,} source colors at visible source locations.", quiet)
    return placed


def verify_allrgb(out_ids: np.ndarray, quiet: bool = False) -> None:
    log("Verifying exact allRGB coverage.", quiet)

    if out_ids.size != TOTAL_COLORS:
        raise RuntimeError(f"Output has {out_ids.size:,} pixels, expected {TOTAL_COLORS:,}.")

    # Strong practical proof: with exactly 16,777,216 pixels, seeing all possible
    # color ids means every color appears exactly once.
    seen = np.zeros(TOTAL_COLORS, dtype=bool)
    seen[out_ids] = True
    seen_count = int(np.count_nonzero(seen))
    if seen_count != TOTAL_COLORS:
        raise RuntimeError(
            f"Verification failed: saw {seen_count:,} unique RGB colors, "
            f"expected {TOTAL_COLORS:,}."
        )

    # Additional cheap invariants. Useful when debugging.
    expected_sum = TOTAL_COLORS * (TOTAL_COLORS - 1) // 2
    actual_sum = int(out_ids.astype(np.uint64).sum())
    if actual_sum != expected_sum:
        raise RuntimeError(f"Verification failed: sum mismatch {actual_sum} != {expected_sum}.")

    actual_xor = int(np.bitwise_xor.reduce(out_ids).item())
    if actual_xor != 0:
        raise RuntimeError(f"Verification failed: xor mismatch {actual_xor} != 0.")

    log("Verification passed: every 24-bit RGB color is present exactly once.", quiet)


def process_image(
    input_path: str | Path,
    output_path: str | Path,
    *,
    fit: FitMode = "cover",
    palette_mode: PaletteMode = "adaptive",
    alpha_background: Tuple[int, int, int] = (255, 255, 255),
    chunk_pixels: int = 262_144,
    chunk_colors: int = 262_144,
    verify: bool = True,
    png_compress_level: int = 6,
    quiet: bool = False,
) -> Path:
    """
    Main importable function.

    Args:
        input_path: Source image path.
        output_path: PNG output path.
        fit: cover, contain, or stretch.
        palette_mode: adaptive or fixed.
        alpha_background: RGB background used for transparent pixels/padding.
        chunk_pixels/chunk_colors: Lower these if your machine has limited RAM.
        verify: If True, prove all 24-bit RGB colors are present once.
        png_compress_level: 0..9; lower is faster/larger, higher is smaller/slower.
        quiet: Suppress progress messages.
    """
    input_path = Path(input_path)
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    log(f"Input:  {input_path}", quiet)
    log(f"Output: {output_path}", quiet)

    src = normalize_to_rgb(input_path, alpha_background)
    log(f"Loaded source image: {src.size[0]:,} x {src.size[1]:,}", quiet)

    log("Extracting unique source RGB colors.", quiet)
    source_ids_all = rgb_image_to_ids(src)
    source_unique_ids = np.unique(source_ids_all).astype(np.uint32, copy=False)
    del source_ids_all

    if source_unique_ids.size > TOTAL_COLORS:
        # This cannot happen for RGB, but keep the guard for clarity.
        raise RuntimeError("Source image has more unique RGB colors than the output can contain.")

    log(f"Source unique colors: {source_unique_ids.size:,}", quiet)

    source_mask = np.zeros(TOTAL_COLORS, dtype=bool)
    source_mask[source_unique_ids] = True
    placed_source_mask = np.zeros(TOTAL_COLORS, dtype=bool)

    log(f"Creating {CANVAS_SIZE:,} x {CANVAS_SIZE:,} guide images.", quiet)
    guide_img = fit_image(
        src,
        (CANVAS_SIZE, CANVAS_SIZE),
        mode=fit,
        resample=RESAMPLING_LANCZOS,
        background=alpha_background,
    )
    stamp_img = fit_image(
        src,
        (CANVAS_SIZE, CANVAS_SIZE),
        mode=fit,
        resample=RESAMPLING_NEAREST,
        background=alpha_background,
    )

    palette = build_palette_16(guide_img, mode=palette_mode, quiet=quiet)

    log("Bucket-labeling guide pixels.", quiet)
    guide_arr = np.asarray(guide_img, dtype=np.uint8)
    guide_flat = guide_arr.reshape(-1, 3)
    guide_labels = nearest_palette_labels_for_rgb_flat(
        guide_flat,
        palette,
        chunk_pixels=chunk_pixels,
        quiet=quiet,
        label="guide pixels",
    )
    guide_luma_flat = luma_from_rgb_flat(guide_flat)
    del guide_arr, guide_flat, guide_img

    out_ids = np.empty(TOTAL_COLORS, dtype=np.uint32)
    assigned_pos = np.zeros(TOTAL_COLORS, dtype=bool)

    stamp_ids = rgb_image_to_ids(stamp_img)
    del stamp_img

    place_first_visible_source_occurrences(
        out_ids=out_ids,
        assigned_pos=assigned_pos,
        source_mask=source_mask,
        placed_source_mask=placed_source_mask,
        stamp_ids=stamp_ids,
        quiet=quiet,
    )
    del stamp_ids

    unplaced_source_ids = source_unique_ids[~placed_source_mask[source_unique_ids]]
    log(f"Unplaced source colors after visible pass: {unplaced_source_ids.size:,}", quiet)

    place_color_ids_by_buckets(
        out_ids=out_ids,
        assigned_pos=assigned_pos,
        guide_labels=guide_labels,
        guide_luma_flat=guide_luma_flat,
        color_ids=unplaced_source_ids,
        palette=palette,
        chunk_colors=chunk_colors,
        quiet=quiet,
        stage_name="Source-color completion",
    )
    del unplaced_source_ids, placed_source_mask

    log("Building remaining RGB color set.", quiet)
    remaining_ids = np.flatnonzero(~source_mask).astype(np.uint32, copy=False)
    del source_mask, source_unique_ids
    log(f"Remaining non-source colors: {remaining_ids.size:,}", quiet)

    place_color_ids_by_buckets(
        out_ids=out_ids,
        assigned_pos=assigned_pos,
        guide_labels=guide_labels,
        guide_luma_flat=guide_luma_flat,
        color_ids=remaining_ids,
        palette=palette,
        chunk_colors=chunk_colors,
        quiet=quiet,
        stage_name="Remaining allRGB completion",
    )
    del remaining_ids, guide_labels, guide_luma_flat

    open_count = int(np.count_nonzero(~assigned_pos))
    if open_count != 0:
        raise RuntimeError(f"Internal error: {open_count:,} pixels were left unassigned.")
    del assigned_pos

    if verify:
        verify_allrgb(out_ids, quiet=quiet)

    log("Converting color ids to RGB image.", quiet)
    out_rgb = ids_to_rgb_array(out_ids, size=CANVAS_SIZE)
    del out_ids

    log("Saving PNG.", quiet)
    Image.fromarray(out_rgb, mode="RGB").save(
        output_path,
        format="PNG",
        optimize=False,
        compress_level=int(png_compress_level),
    )
    log(f"Done: {output_path}", quiet)
    return output_path


def default_output_path(input_path: str | Path) -> Path:
    input_path = Path(input_path)
    stamp = time.strftime("%Y%m%d-%H%M%S")
    return input_path.with_name(f"{input_path.stem}-allrgb-{stamp}.png")


def in_google_colab() -> bool:
    try:
        import google.colab  # type: ignore  # noqa: F401
        return True
    except Exception:
        return False


def run_colab_sketch() -> None:
    from google.colab import files  # type: ignore

    print("Upload a source image.")
    uploaded = files.upload()
    if not uploaded:
        print("No file uploaded.")
        return

    input_name = next(iter(uploaded.keys()))
    output_name = f"{Path(input_name).stem}-allrgb.png"

    result = process_image(
        input_name,
        output_name,
        fit="cover",
        palette_mode="adaptive",
        alpha_background=(255, 255, 255),
        verify=True,
        png_compress_level=6,
        quiet=False,
    )
    files.download(str(result))


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Create a 4096x4096 allRGB image guided by an input image."
    )
    parser.add_argument("input", nargs="?", help="Input image file")
    parser.add_argument("output", nargs="?", help="Output PNG file")
    parser.add_argument(
        "--fit",
        choices=["cover", "contain", "stretch"],
        default="cover",
        help="How to fit the source image into the square output. Default: cover.",
    )
    parser.add_argument(
        "--palette-mode",
        choices=["adaptive", "fixed"],
        default="adaptive",
        help="Use source-adaptive or fixed 16-color buckets. Default: adaptive.",
    )
    parser.add_argument(
        "--alpha-bg",
        default="#ffffff",
        help="Background for transparency and contain-padding. Default: #ffffff.",
    )
    parser.add_argument(
        "--chunk-pixels",
        type=int,
        default=262_144,
        help="Guide-pixel chunk size. Lower this to reduce RAM. Default: 262144.",
    )
    parser.add_argument(
        "--chunk-colors",
        type=int,
        default=262_144,
        help="RGB-color chunk size. Lower this to reduce RAM. Default: 262144.",
    )
    parser.add_argument(
        "--no-verify",
        action="store_true",
        help="Skip exact allRGB verification. Not recommended.",
    )
    parser.add_argument(
        "--png-compress-level",
        type=int,
        default=6,
        choices=range(0, 10),
        metavar="0..9",
        help="PNG compression level. Lower is faster/larger. Default: 6.",
    )
    parser.add_argument("--quiet", action="store_true", help="Suppress progress messages.")
    return parser


def strip_notebook_kernel_args(argv: list[str]) -> list[str]:
    """
    Remove arguments injected by Jupyter/Colab when code is run from a notebook.

    In Colab, sys.argv often looks like:
        ["-f", "/root/.local/share/jupyter/runtime/kernel-xxxx.json"]

    Those arguments belong to the notebook kernel, not to this script. Without this
    cleanup, argparse exits with "unrecognized arguments: -f".
    """
    cleaned: list[str] = []
    i = 0
    while i < len(argv):
        arg = argv[i]

        # Jupyter/IPython kernel connection file. It can appear as either
        # "-f path/to/kernel.json" or "--f=path/to/kernel.json".
        if arg == "-f" and i + 1 < len(argv):
            i += 2
            continue
        if arg.startswith("--f=") or arg.startswith("-f="):
            i += 1
            continue

        # Be conservative: only drop obvious kernel json files, not arbitrary
        # positional paths the user may have supplied.
        if arg.endswith(".json") and "kernel-" in Path(arg).name:
            i += 1
            continue

        cleaned.append(arg)
        i += 1

    return cleaned


def main(argv: Optional[list[str]] = None) -> int:
    if argv is None:
        argv = sys.argv[1:]

    if in_google_colab():
        argv = strip_notebook_kernel_args(list(argv))
        if not argv:
            run_colab_sketch()
            return 0

    parser = build_arg_parser()
    args = parser.parse_args(argv)

    if not args.input:
        parser.print_help()
        return 2

    output = args.output or str(default_output_path(args.input))
    alpha_background = parse_hex_rgb(args.alpha_bg)

    process_image(
        args.input,
        output,
        fit=args.fit,
        palette_mode=args.palette_mode,
        alpha_background=alpha_background,
        chunk_pixels=args.chunk_pixels,
        chunk_colors=args.chunk_colors,
        verify=not args.no_verify,
        png_compress_level=args.png_compress_level,
        quiet=args.quiet,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())