Publish ember-colorized source code

src/ember_colorized/__init__.py

@@ -0,0 +1,3 @@
+"""Ember Colorizer — Apply Ember palettes to images."""
+
+__version__ = "0.1.0"

src/ember_colorized/cli.py

@@ -0,0 +1,170 @@
+# Copyright JesusChapman <jesuschapman@openlat.dev>
+# 2026
+
+from __future__ import annotations
+
+import argparse
+import sys
+import threading
+import time
+from pathlib import Path
+
+from PIL import Image
+
+from .colorizer import colorize, colorize_fast, gpu_info
+from .palettes import PALETTES
+
+
+SPINNERS = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
+
+
+class Spinner:
+    """Animated terminal spinner shown during processing."""
+
+    def __init__(self, message: str):
+        self.message = message
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._run, daemon=True)
+        self._idx = 0
+
+    def _run(self):
+        while not self._stop.is_set():
+            sys.stdout.write(f"\r  {SPINNERS[self._idx % len(SPINNERS)]} {
+                             self.message}  ")
+            sys.stdout.flush()
+            self._idx += 1
+            self._stop.wait(0.08)
+
+    def start(self):
+        self._thread.start()
+
+    def update(self, message: str):
+        self.message = message
+
+    def stop(self):
+        self._stop.set()
+        self._thread.join(timeout=0.2)
+        sys.stdout.write(f"\r{'':60}\r")
+        sys.stdout.flush()
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="ember-colorizer",
+        description="Apply Ember color palettes to images with high precision.",
+    )
+    p.add_argument("input", type=Path,
+                   help="Input image path (PNG, JPG, JPEG, WEBP, BMP)")
+    p.add_argument(
+        "-c", "--colors",
+        type=str,
+        required=True,
+        choices=list(PALETTES.keys()),
+        help="Ember palette to apply",
+    )
+    p.add_argument(
+        "-o", "--output",
+        type=Path,
+        default=None,
+        help="Output file path (overrides auto-generated name)",
+    )
+    p.add_argument(
+        "-s", "--strength",
+        type=float,
+        default=1.0,
+        help="Recolor strength: 0.0 = original, 1.0 = full recolor (default: 1.0)",
+    )
+    p.add_argument(
+        "-m", "--mode",
+        type=str,
+        default="aggressive",
+        choices=["aggressive", "fast"],
+        help="Algorithm: aggressive (K-means) or fast (direct mapping, default: aggressive)",
+    )
+    p.add_argument(
+        "-k", "--clusters",
+        type=int,
+        default=12,
+        help="K-means clusters for aggressive mode (default: 12)",
+    )
+    p.add_argument(
+        "--use-gpu",
+        action="store_true",
+        default=False,
+        help="Use GPU acceleration (Apple Silicon MPS / NVIDIA CUDA / cuML)",
+    )
+    return p
+
+
+def _output_path(input_path: Path, palette_name: str) -> Path:
+    return input_path.parent / f"{input_path.stem}-{palette_name}-colorized{input_path.suffix}"
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+
+    input_path: Path = args.input
+    if not input_path.exists():
+        print(f"Error: file not found: {input_path}", file=sys.stderr)
+        sys.exit(1)
+
+    if input_path.suffix.lower() not in {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff"}:
+        print(f"Error: unsupported file type: {
+              input_path.suffix}", file=sys.stderr)
+        sys.exit(1)
+
+    has_gpu, gpu_name = gpu_info()
+
+    if args.use_gpu and not has_gpu:
+        print("Warning: no GPU backend available, falling back to CPU.",
+              file=sys.stderr)
+        print(
+            "  Install with: pip install 'ember-colorized[gpu]'", file=sys.stderr)
+        args.use_gpu = False
+
+    output_path = args.output or _output_path(input_path, args.colors)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    print(f"Input:    {input_path}")
+    print(f"Palette:  {args.colors} ({PALETTES[args.colors]['name']})")
+    print(f"Mode:     {args.mode}")
+    print(f"Strength: {args.strength}")
+    if args.use_gpu:
+        print(f"GPU:      yes ({gpu_name})")
+    print(f"Output:   {output_path}")
+
+    img = Image.open(input_path)
+    print(f"Loaded:   {img.size[0]}x{img.size[1]} {img.mode}")
+
+    spinner = Spinner("Initializing...")
+    spinner.start()
+
+    start = time.perf_counter()
+
+    def _progress(msg: str):
+        spinner.update(msg)
+
+    if args.mode == "aggressive":
+        result = colorize(
+            img, args.colors,
+            strength=args.strength,
+            n_clusters=args.clusters,
+            use_gpu=args.use_gpu,
+            progress_callback=_progress,
+        )
+    else:
+        result = colorize_fast(
+            img, args.colors, strength=args.strength, progress_callback=_progress)
+
+    spinner.stop()
+
+    result.save(str(output_path))
+
+    elapsed = time.perf_counter() - start
+    sec = int(elapsed)
+    ms = int((elapsed - sec) * 1000)
+    if sec > 0:
+        print(f"Success! in {sec}.{ms:03d}s")
+    else:
+        print(f"Success! in {ms}ms")

src/ember_colorized/colorizer.py

@@ -0,0 +1,252 @@
+# Copyright JesusChapman <jesuschapman@openlat.dev>
+# 2026
+
+"""Core colorization engine — maps image colors to ember palettes."""
+
+from __future__ import annotations
+from sklearn.cluster import KMeans
+
+import numpy as np
+from PIL import Image
+
+from .palettes import get_palette_rgb
+
+# --- Backend detection ---
+
+_backend = "cpu"
+_gpu_name = None
+
+try:
+    from cuml.cluster import KMeans as CuMLKMeans
+    from cuml.common.device_selection import DeviceProperties
+    _backend = "cuml"
+    _gpu_name = f"NVIDIA ({DeviceProperties().name})"
+except ImportError:
+    pass
+
+if _backend == "cpu":
+    try:
+        import torch
+        if torch.cuda.is_available():
+            _backend = "torch-cuda"
+            _gpu_name = f"NVIDIA ({torch.cuda.get_device_name(0)})"
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            _backend = "torch-mps"
+            import platform
+            _gpu_name = f"Apple ({platform.machine()})"
+    except ImportError:
+        pass
+
+
+def gpu_info() -> tuple[bool, str]:
+    """Return (is_gpu, gpu_name_or_empty)."""
+    return _backend != "cpu", _gpu_name or ""
+
+
+# --- K-means backends ---
+
+def _kmeans_cuml(pixels: np.ndarray, n_clusters: int):
+    from cuml.common.frame_utils import input_to_cuml_array
+    cu_pixels, _ = input_to_cuml_array(pixels)
+    km = CuMLKMeans(n_clusters=n_clusters, n_init=5,
+                    max_iter=200, random_state=42)
+    km.fit(cu_pixels)
+    return km.labels_.to_numpy().astype(np.int32), km.cluster_centers_.to_numpy()
+
+
+def _kmeans_torch_gpu(pixels: np.ndarray, n_clusters: int, device: str):
+    """GPU K-means with smart downsampling for large images.
+
+    Strategy: cluster on a downsampled subset, then assign all pixels to
+    nearest center using vectorized distance. This avoids the O(N×K) memory
+    and compute overhead of full K-means on GPU for large N.
+    """
+    import torch
+
+    n = len(pixels)
+
+    # For large images, downsample for clustering (GPU wins on small N)
+    MAX_CLUSTER = 200_000
+    if n > MAX_CLUSTER:
+        idx = np.random.default_rng(42).choice(n, MAX_CLUSTER, replace=False)
+        sample = pixels[idx]
+    else:
+        sample = pixels
+
+    t_sample = torch.tensor(sample, dtype=torch.float32, device=device)
+
+    # K-means++ init on CPU
+    ns = len(sample)
+    centers_np = np.empty((n_clusters, sample.shape[1]), dtype=np.float32)
+    rng = np.random.default_rng(42)
+    centers_np[0] = sample[rng.integers(0, ns)]
+    for i in range(1, n_clusters):
+        c = torch.tensor(centers_np[:i], dtype=torch.float32, device=device)
+        dists = torch.cdist(t_sample.unsqueeze(
+            0), c.unsqueeze(0)).squeeze(0).min(dim=1).values
+        dists_np = dists.cpu().numpy()
+        dists_np = dists_np - dists_np.min()
+        total = dists_np.sum()
+        probs = dists_np / \
+            total if total > 0 else np.ones(ns, dtype=np.float32) / ns
+        centers_np[i] = sample[rng.choice(ns, p=probs)]
+
+    centers = torch.tensor(centers_np, dtype=torch.float32, device=device)
+
+    # Vectorized K-means on sample
+    for _ in range(200):
+        dists = torch.cdist(t_sample.unsqueeze(
+            0), centers.unsqueeze(0)).squeeze(0)
+        labels = dists.argmin(dim=1)
+        new_centers = torch.zeros_like(centers)
+        counts = torch.zeros(n_clusters, device=device)
+        new_centers.scatter_add_(0, labels.unsqueeze(
+            1).expand(-1, sample.shape[1]), t_sample)
+        counts.scatter_add_(0, labels, torch.ones(ns, device=device))
+        counts = counts.clamp(min=1)
+        new_centers /= counts.unsqueeze(1)
+        if torch.allclose(new_centers, centers, atol=1e-4):
+            break
+        centers = new_centers
+
+    # Assign ALL pixels to nearest center (vectorized, fast)
+    t_all = torch.tensor(pixels, dtype=torch.float32, device=device)
+    dists = torch.cdist(t_all.unsqueeze(0), centers.unsqueeze(0)).squeeze(0)
+    labels = dists.argmin(dim=1)
+
+    return labels.cpu().numpy().astype(np.int32), centers.cpu().numpy()
+
+
+def _kmeans_fit(pixels: np.ndarray, n_clusters: int, use_gpu: bool = True):
+    """Run K-means on the best available backend."""
+    if use_gpu:
+        if _backend == "cuml":
+            return _kmeans_cuml(pixels, n_clusters)
+        if _backend in ("torch-cuda", "torch-mps"):
+            return _kmeans_torch_gpu(pixels, n_clusters, "cuda" if _backend == "torch-cuda" else "mps")
+    km = KMeans(n_clusters=n_clusters, n_init=5, max_iter=200, random_state=42)
+    km.fit(pixels)
+    return km.labels_.astype(np.int32), km.cluster_centers_
+
+
+# ---------------------------------------------------------------------------
+# Aggressive mode — K-means clustering → nearest palette
+# ---------------------------------------------------------------------------
+
+def colorize(
+    img: Image.Image,
+    palette_name: str,
+    strength: float = 1.0,
+    n_clusters: int = 12,
+    use_gpu: bool = False,
+    progress_callback=None,
+) -> Image.Image:
+    """K-means clusters image colors, then maps each cluster to nearest palette color."""
+    palette = np.array(get_palette_rgb(palette_name), dtype=np.float64)
+
+    rgba = img.convert("RGBA")
+    rgb = np.array(rgba, dtype=np.float64)
+    alpha = rgb[:, :, 3:4] if rgb.shape[2] == 4 else np.ones(
+        (*rgb.shape[:2], 1))
+    rgb = rgb[:, :, :3]
+
+    h, w, _ = rgb.shape
+    pixels = rgb.reshape(-1, 3)
+
+    opaque_mask = alpha.reshape(-1) > 127
+    opaque_pixels = pixels[opaque_mask]
+
+    if len(opaque_pixels) == 0:
+        return img
+
+    if progress_callback:
+        progress_callback("Clustering colors...")
+
+    n_clusters = min(n_clusters, len(opaque_pixels))
+    labels, centers = _kmeans_fit(opaque_pixels, n_clusters, use_gpu)
+
+    if progress_callback:
+        progress_callback("Mapping to palette...")
+
+    diffs = centers[:, None, :] - palette[None, :, :]
+    dists = np.sum(diffs ** 2, axis=2)
+    nearest = np.argmin(dists, axis=1)
+
+    cluster_to_palette = {i: palette[nearest[i]] for i in range(len(nearest))}
+    mapped = np.array([cluster_to_palette[l]
+                      for l in labels], dtype=np.float64)
+
+    if strength < 1.0:
+        mapped = opaque_pixels * (1 - strength) + mapped * strength
+
+    result = pixels.copy()
+    result[opaque_mask] = mapped
+    result = result.reshape(h, w, 3)
+    result_img = Image.fromarray(result.astype(np.uint8), "RGB")
+
+    if progress_callback:
+        progress_callback("Done")
+
+    if img.mode == "RGBA":
+        result_img = Image.merge(
+            "RGBA", (*result_img.split(), rgba.split()[3]))
+
+    return result_img
+
+
+# ---------------------------------------------------------------------------
+# Fast mode — pure RGB nearest-palette (gruvboxify approach)
+# ---------------------------------------------------------------------------
+
+def colorize_fast(
+    img: Image.Image,
+    palette_name: str,
+    strength: float = 1.0,
+    progress_callback=None,
+) -> Image.Image:
+    """Pixel-by-pixel RGB Euclidean distance to nearest palette color."""
+    palette = np.array(get_palette_rgb(palette_name), dtype=np.float64)
+
+    rgba = img.convert("RGBA")
+    rgb = np.array(rgba, dtype=np.float64)
+    alpha = rgb[:, :, 3:4] if rgb.shape[2] == 4 else np.ones(
+        (*rgb.shape[:2], 1))
+    rgb = rgb[:, :, :3]
+
+    h, w, _ = rgb.shape
+    pixels = rgb.reshape(-1, 3)
+
+    opaque_mask = alpha.reshape(-1)
+    opaque_bool = opaque_mask > 127
+
+    if not opaque_bool.any():
+        return img
+
+    opaque_pixels = pixels[opaque_bool]
+
+    if progress_callback:
+        progress_callback("Mapping colors...")
+
+    diffs = opaque_pixels[:, None, :] - palette[None, :, :]
+    dists = np.sum(diffs ** 2, axis=2)
+    nearest_idx = np.argmin(dists, axis=1)
+
+    mapped = palette[nearest_idx]
+
+    if strength < 1.0:
+        mapped = opaque_pixels * (1 - strength) + mapped * strength
+
+    result = pixels.copy()
+    result[opaque_bool] = mapped
+
+    result = result.reshape(h, w, 3)
+    result_img = Image.fromarray(result.astype(np.uint8), "RGB")
+
+    if progress_callback:
+        progress_callback("Done")
+
+    if img.mode == "RGBA":
+        result_img = Image.merge(
+            "RGBA", (*result_img.split(), rgba.split()[3]))
+
+    return result_img

src/ember_colorized/palettes.py

@@ -0,0 +1,62 @@
+# Copyright JesusChapman <jesuschapman@openlat.dev>
+# 2026
+
+"""Ember color palettes for image colorization."""
+
+from __future__ import annotations
+
+PALETTES: dict[str, dict] = {
+    "ember": {
+        "name": "Ember",
+        "type": "dark",
+        "colors": [
+            "#1c1b19", "#242320", "#252422", "#2e2d2a", "#3e3c38",
+            "#585550", "#706c61", "#908a7e", "#b8b0a0", "#d8d0c0",
+            "#b0a898", "#e08060", "#c09058", "#c8b468", "#8a9868",
+            "#80a090", "#7890a0", "#b07878", "#988090",
+        ],
+    },
+    "ember-soft": {
+        "name": "Ember Soft",
+        "type": "dark",
+        "colors": [
+            "#242320", "#2a2927", "#2c2b28", "#353430", "#444240",
+            "#585550", "#706c61", "#908a7e", "#b8b0a0", "#d8d0c0",
+            "#b0a898", "#e08060", "#c09058", "#c8b468", "#8a9868",
+            "#80a090", "#7890a0", "#b07878", "#988090",
+        ],
+    },
+    "ember-light": {
+        "name": "Ember Light",
+        "type": "light",
+        "colors": [
+            "#e6dac4", "#ddd0b8", "#d8ccb0", "#cec2a8", "#b8ac96",
+            "#989080", "#787060", "#605848", "#484030", "#282418",
+            "#585040", "#b84c30", "#946030", "#7a6820", "#4a6830",
+            "#386858", "#3a6080", "#905050", "#706070",
+        ],
+    },
+    "ember-lighter": {
+        "name": "Ember Lighter",
+        "type": "light",
+        "colors": [
+            "#e8e4de", "#dfd9d4", "#f2efec", "#d0ccc6", "#c8c2b8",
+            "#a09484", "#807868", "#585040", "#3a3428", "#3a3428",
+            "#585040", "#b84c30", "#946030", "#7a6820", "#4a6830",
+            "#386858", "#3a6080", "#905050", "#706070",
+        ],
+    },
+}
+
+
+def hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
+    h = hex_color.lstrip("#")
+    return (int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16))
+
+
+def get_palette_rgb(name: str) -> list[tuple[int, int, int]]:
+    if name not in PALETTES:
+        raise ValueError(
+            f"Unknown palette: {name}. Available: {', '.join(PALETTES.keys())}"
+        )
+    return [hex_to_rgb(c) for c in PALETTES[name]["colors"]]