feat: Batch processing images (#91)

- Added batch_extract_colors() for parallel processing of multiple images
- CLI now supports processing multiple image sources in a single command
- CLI has an enhanced progress display with recently extracted palettes and palette previews
- Added metadata-field to Palette-class. It contains information about the image source, extraction parameters and extraction performance. This is a preparatory step for better palette exports.
- removed --filename and --image-url parameters
- mode parameter to color extraction methods now accepts a ExtractionMethod enum member instead of a string.

Author: qTipTip

Pylette/__init__.py

@@ -1,6 +1,6 @@
 from Pylette.src.color import Color
-from Pylette.src.color_extraction import extract_colors
+from Pylette.src.color_extraction import batch_extract_colors, extract_colors
 from Pylette.src.palette import Palette
 from Pylette.src.types import ImageInput
 
-__all__ = ["extract_colors", "Palette", "Color", "ImageInput"]
+__all__ = ["extract_colors", "batch_extract_colors", "Palette", "Color", "ImageInput"]

Pylette/cmd.py

@@ -1,14 +1,12 @@
 import pathlib
 from enum import Enum
+from typing import Annotated, List
 
 import typer
 
-from Pylette import extract_colors
-
-
-class ExtractionMode(str, Enum):
-    KM = "KM"
-    MC = "MC"
+from Pylette.src.cli_utils import PyletteProgress
+from Pylette.src.color_extraction import batch_extract_colors
+from Pylette.src.types import BatchResult, ExtractionMethod
 
 
 class SortBy(str, Enum):
@@ -25,12 +23,12 @@ class ColorSpace(str, Enum):
 pylette_app = typer.Typer()
 
 
-@pylette_app.command()
+@pylette_app.command(no_args_is_help=True)
 def main(
-    ctx: typer.Context,
-    filename: pathlib.Path | None = None,
-    image_url: str | None = None,
-    mode: ExtractionMode = ExtractionMode.KM,
+    image_sources: Annotated[
+        List[str], typer.Argument(help="A list of paths / directories / URLs pointing to images.")
+    ],  # These can be paths or URLs
+    mode: ExtractionMethod = ExtractionMethod.KM,
     n: int = 5,
     sort_by: SortBy = SortBy.luminance,
     stdout: bool = True,
@@ -43,41 +41,68 @@ def main(
         max=255,
         help="Alpha threshold for transparent image masking (0-255). Pixels with alpha below this value are excluded.",
     ),
+    num_threads: int | None = typer.Option(
+        None, min=1, help="Number of threads used for batch extraction of color palettes"
+    ),
 ):
-    if filename is None and image_url is None:
-        typer.echo(ctx.get_help())
-        raise typer.Exit(code=0)
-
-    if filename is not None and image_url is not None:
-        typer.echo("Please provide either a filename or an image-url, but not both.")
-        raise typer.Exit(code=1)
-
-    if filename is not None:
-        image = filename  # Path
-    else:
-        assert image_url is not None
-        image = image_url  # str (URL)
-
     output_file_path = str(out_filename) if out_filename is not None else None
-    try:
-        palette = extract_colors(
-            image=image,
+
+    # Set up progress bar for CLI
+    with PyletteProgress(palette_size=n) as progress:
+        task_id = progress.add_task("Extracting colors...", total=len(image_sources))
+
+        def progress_callback(task_number: int, result: BatchResult):
+            if result.success and result.palette:
+                progress.mark_task_complete(
+                    task_number=task_number,
+                    task_id=task_id,
+                    completed_task_name=result.palette.metadata["image_source"] if result.palette.metadata else "",
+                    palette_colors=result.palette.colors,
+                )
+            else:
+                progress.update(task_id, advance=1)
+
+        results = batch_extract_colors(
+            images=image_sources,
             palette_size=n,
             sort_mode=sort_by.value,
-            mode=mode.value,
+            mode=mode,
             alpha_mask_threshold=alpha_mask_threshold,
+            max_workers=num_threads,
+            progress_callback=progress_callback,
         )
-    except ValueError as e:
-        typer.echo(str(e))
-        raise typer.Exit(code=1)
-
-    palette.to_csv(filename=output_file_path, frequency=True, stdout=stdout, colorspace=colorspace.value)
-    if display_colors:
-        palette.display()
-
 
-def docs():
-    typer.launch("https://qtiptip.github.io/Pylette/")
+    successful = [r for r in results if r.success]
+    failed = [r for r in results if not r.success]
+    for success in successful:
+        if success.palette is not None:
+            success.palette.to_csv(
+                filename=output_file_path, frequency=True, stdout=stdout, colorspace=colorspace.value
+            )
+            if display_colors:
+                success.palette.display()
+
+    if failed:
+        print_extraction_summary(successful, failed)
+
+    # If we have no successful extractions, return with code 1
+    if not successful:
+        raise typer.Exit(1)
+    # Otherwise, if we have some failures, return with code 2
+    elif failed:
+        raise typer.Exit(2)
+
+
+def print_extraction_summary(successful: list[BatchResult], failed: list[BatchResult]):
+    total = len(successful) + len(failed)
+
+    if successful:
+        typer.secho(f"✓ Processed {len(successful)}/{total} images successfully", fg=typer.colors.GREEN)
+    if failed:
+        typer.secho(f"✗ {len(failed)} images failed:", fg=typer.colors.RED)
+        for result in failed:
+            error_msg = str(result.error)
+            typer.secho(f"{result.source}: {error_msg}", err=True)
 
 
 def main_typer() -> None:

Pylette/src/cli_utils.py

@@ -0,0 +1,103 @@
+from collections import deque
+from typing import Iterator
+
+from rich.console import Console, RenderableType
+from rich.progress import BarColumn, Progress, TaskID, TextColumn, TimeElapsedColumn
+from rich.table import Table
+from rich.text import Text
+
+from Pylette.src.color import Color
+
+
+class RecentlyCompletedDisplay:
+    def __init__(self, num_items: int = 5, max_text_width: int = 20, max_num_preview_colors: int = 6) -> None:
+        self.max_items = num_items
+        self.max_name_width = max_text_width
+        self.items = deque(maxlen=self.max_items)
+        self.num_preview_colors = max_num_preview_colors
+
+        super().__init__()
+
+    def add_completed_task(self, task_number: int, task_name: str, colors: list[Color]) -> None:
+        # Truncate name if too long
+        display_name = task_name
+        if len(display_name) > self.max_name_width:
+            display_name = task_name[:-3] + "..."
+
+        self.items.append(
+            {
+                "number": task_number,
+                "name": display_name,
+                "colors": colors[: self.num_preview_colors],
+            }
+        )
+
+    def render(self) -> RenderableType:
+        table = Table.grid("Number", "Image Source", "Preview", padding=(0, 1))
+        table.add_column()  # Task number
+        table.add_column(width=self.max_name_width)  # Task name
+        table.add_column()  # Palette dots
+
+        # Add each recent task as a row
+        table.add_row("Number", "Image Source", "Preview")
+        for task_info in list(self.items):
+            task_name = task_info["name"]
+            task_number = str(task_info["number"])
+            colors = task_info["colors"]
+
+            # Create colored dots for palette
+            dots_text = Text()
+            for c in colors:
+                r, g, b = c.rgb
+                color = f"rgb({r},{g},{b})"
+                dots_text.append("●", style=color)
+
+            table.add_row(task_number, task_name, dots_text)
+
+        return table
+
+
+class PyletteProgress(Progress):
+    """Custom Progress class for Pylette color extraction with palette preview"""
+
+    def __init__(
+        self,
+        palette_size: int = 5,
+        max_recent_items: int = 5,
+        max_name_width: int = 40,
+        console_width: int = 140,
+        *args,  # pyright: ignore[reportMissingParameterType]
+        **kwargs,  # pyright: ignore[reportMissingParameterType]
+    ):
+        # Create the recently completed display
+        self.recently_completed = RecentlyCompletedDisplay(
+            num_items=max_recent_items, max_text_width=max_name_width, max_num_preview_colors=palette_size
+        )
+
+        # Set up the console with appropriate width
+        console = Console(width=console_width)
+
+        # Initialize Progress with custom columns
+        super().__init__(
+            TextColumn("[bold blue]{task.description}"),
+            BarColumn(bar_width=40),
+            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+            TextColumn("({task.completed}/{task.total})"),
+            TimeElapsedColumn(),
+            console=console,
+            *args,
+            **kwargs,
+        )
+
+    def get_renderables(self) -> Iterator[RenderableType]:
+        for renderable in super().get_renderables():
+            yield renderable
+
+        yield ""
+        yield self.recently_completed.render()
+
+    def mark_task_complete(
+        self, task_id: TaskID, task_number: int, completed_task_name: str, palette_colors: list[Color]
+    ) -> None:
+        self.update(task_id, advance=1)
+        self.recently_completed.add_completed_task(task_number, completed_task_name, palette_colors)