color-extract.py

#!/usr/bin/env python3
"""
Color Extractor — Extract dominant colors from images.

Usage:
    python color-extract.py <image>
    python color-extract.py <image> --colors 5
    python color-extract.py <image> --format hex
    python color-extract.py <directory> --glob *.jpg

Options:
    --colors N        Number of dominant colors to extract (default: 5)
    --format FMT      Output format: rgb, hex, or both (default: both)
    --show-preview    Show an image preview with color swatches (requires matplotlib)
    --glob PAT        Process multiple files by glob pattern
    --help            Show this help message and exit

Requirements:
    pip install Pillow numpy
    (optional: matplotlib for --show-preview)
"""

import sys
import argparse
from pathlib import Path
from collections import Counter

try:
    from PIL import Image
except ImportError:
    print("Error: Pillow is required. Install with: pip install Pillow")
    sys.exit(1)

try:
    import numpy as np
except ImportError:
    print("Error: numpy is required. Install with: pip install numpy")
    sys.exit(1)


def rgb_to_hex(r: int, g: int, b: int) -> str:
    """Convert RGB tuple to hex string."""
    return f"#{r:02x}{g:02x}{b:02x}"


def extract_colors(img_path: Path, num_colors: int) -> list[dict]:
    """Extract dominant colors from an image using k-means approximation (quantization)."""
    img = Image.open(img_path).convert("RGB")

    # Downsample large images for performance
    max_dim = 200
    if max(img.size) > max_dim:
        ratio = max_dim / max(img.size)
        new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
        img = img.resize(new_size, Image.LANCZOS)

    pixels = np.array(img).reshape(-1, 3)

    # Subtractive color quantization: reduce to 16-bit color space
    # then count frequencies
    quantized = (pixels // 16) * 16  # quantize to 16-step multiples
    color_tuples = [tuple(c) for c in quantized]

    # Count most common colors
    counter = Counter(color_tuples)
    most_common = counter.most_common(num_colors * 3)  # grab more for better sampling

    # Group nearby colors and pick top
    seen = set()
    dominant: list[dict] = []
    for color, count in most_common:
        if color in seen:
            continue
        seen.add(color)
        r, g, b = color
        # Ensure color is at least somewhat distinct
        dominant.append({
            "rgb": (r, g, b),
            "hex": rgb_to_hex(r, g, b),
            "pixel_count": count,
        })
        if len(dominant) >= num_colors:
            break

    return dominant


def process_file(filepath: Path, num_colors: int, fmt: str) -> None:
    """Extract and print colors from a single file."""
    try:
        colors = extract_colors(filepath, num_colors)
    except Exception as e:
        print(f"  ✗ {filepath.name}: {e}")
        return

    print(f"\n── {filepath.name} ──")
    for i, c in enumerate(colors, 1):
        parts = []
        if fmt in ("rgb", "both"):
            parts.append(f"rgb({c['rgb'][0]}, {c['rgb'][1]}, {c['rgb'][2]})")
        if fmt in ("hex", "both"):
            parts.append(c["hex"])
        bar_len = int(c["pixel_count"] / max(cc["pixel_count"] for cc in colors) * 20)
        bar = "█" * bar_len
        print(f"  {i}. {'  '.join(parts)}  {bar}")


def main():
    parser = argparse.ArgumentParser(
        description="Extract dominant colors from images.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            "  python color-extract.py photo.jpg\n"
            "  python color-extract.py photo.jpg --colors 10 --format hex\n"
            "  python color-extract.py ./images --glob *.png --colors 3\n"
        ),
    )
    parser.add_argument("input", help="Image file or directory")
    parser.add_argument("--colors", type=int, default=5, help="Number of colors (default: 5)")
    parser.add_argument("--format", choices=["rgb", "hex", "both"], default="both",
                        help="Output format (default: both)")
    parser.add_argument("--show-preview", action="store_true",
                        help="Show a preview with color swatches (requires matplotlib)")
    parser.add_argument("--glob", default="*", help="Glob pattern when input is a directory")

    args = parser.parse_args()

    input_path = Path(args.input)

    files: list[Path] = []
    if input_path.is_dir():
        files = list(input_path.glob(args.glob))
        files = [f for f in files if f.suffix.lower() in
                 {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"}]
        if not files:
            print(f"No images matching '{args.glob}' in '{input_path}'.")
            sys.exit(0)
    elif input_path.is_file():
        files = [input_path]
    else:
        print(f"Error: '{args.input}' not found.")
        sys.exit(1)

    print(f"Extracting {args.colors} dominant color(s) from {len(files)} file(s)...")

    for f in files:
        process_file(f, args.colors, args.format)

    if args.show_preview:
        try:
            import matplotlib.pyplot as plt
            from matplotlib.patches import Rectangle

            if len(files) > 1:
                print("\n(Preview only available for single image mode)")
            else:
                colors = extract_colors(files[0], args.colors)
                fig, ax = plt.subplots(figsize=(8, 2))
                ax.set_xlim(0, args.colors)
                ax.set_ylim(0, 1)
                ax.axis("off")
                for i, c in enumerate(colors):
                    norm = tuple(x / 255 for x in c["rgb"])
                    ax.add_patch(Rectangle((i, 0), 1, 1, facecolor=norm))
                plt.show()
        except ImportError:
            print("\n(Install matplotlib for preview: pip install matplotlib)")


if __name__ == "__main__":
    main()