augmenator/cli.py

"""Batch CLI: `augmenator-batch --input ./photos --count 5`."""

from __future__ import annotations

import argparse
import json
import random
import re
import sys
from datetime import datetime, timezone
from pathlib import Path

from PIL import Image

from augmenator import run_pipeline
from augmenator.ai_tools import AI_TOOL_KEYWORD_IDS
from augmenator.keyword_catalog import AUGMENT_KEYWORDS
from augmenator.planner import warmup as warmup_planner

IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"}

COMPOUND_PROMPTS = (
    "vintage warm look",
    "blur everything softly",
    "replace background with neon city at night and rotate",
    "cover and add cutout",
    "flip horizontally",
)


def collect_input_images(folder: Path) -> list[Path]:
    if not folder.is_dir():
        return []
    images = [
        path
        for path in folder.iterdir()
        if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
    ]
    return sorted(images, key=lambda p: p.name.lower())


def build_prompt_pool(*, use_ai_tools: bool) -> list[str]:
    prompts: list[str] = []
    for keyword in AUGMENT_KEYWORDS:
        if not use_ai_tools and keyword.id in AI_TOOL_KEYWORD_IDS:
            continue
        if not use_ai_tools and keyword.spatial_extra and keyword.spatial_extra.get("avoid_text"):
            continue
        prompts.extend(keyword.phrases)
    prompts.extend(COMPOUND_PROMPTS)
    if not use_ai_tools:
        prompts = [
            p
            for p in prompts
            if "replace background" not in p.lower()
            and "avoid text" not in p.lower()
            and "not text" not in p.lower()
            and "except text" not in p.lower()
        ]
    return sorted(set(prompts))


def slugify(text: str, max_len: int = 48) -> str:
    slug = re.sub(r"[^a-z0-9]+", "_", text.lower()).strip("_")
    return slug[:max_len] or "augmentation"


def augment_image(
    source_path: Path,
    source_image: Image.Image,
    output_dir: Path,
    *,
    count: int,
    prompt_pool: list[str],
    use_ai_tools: bool,
    strength: float,
) -> tuple[list[dict], int]:
    """Create up to `count` augmentations for one source image. Returns manifest rows and created count."""
    items: list[dict] = []
    created = 0
    attempts = 0
    max_attempts = count * 8
    source_stem = slugify(source_path.stem, max_len=32)

    while created < count and attempts < max_attempts:
        attempts += 1
        instruction = random.choice(prompt_pool)
        result = run_pipeline(
            source_image,
            instruction,
            strength=strength,
            use_ai_tools=use_ai_tools,
        )
        if not result["supported"]:
            continue

        created += 1
        filename = f"{source_stem}_{created:03d}_{slugify(instruction)}.png"
        out_path = output_dir / filename
        result["image"].save(out_path)

        items.append(
            {
                "file": filename,
                "source": source_path.name,
                "instruction": instruction,
                "applied_tags": result["applied_tags"],
                "applied_spatial": result["applied_spatial"],
                "use_ai_tools": use_ai_tools,
            }
        )
        tags = ", ".join(result["applied_tags"]) or "(none)"
        print(f"  [{created}/{count}] {filename}  <-  {instruction!r}  [{tags}]")

    return items, created


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate augmented images for each JPG/PNG in an input folder.",
    )
    parser.add_argument(
        "--input",
        "-i",
        required=True,
        type=Path,
        help="Input folder containing .jpg / .jpeg / .png images",
    )
    parser.add_argument(
        "--output",
        "-o",
        type=Path,
        default=Path("generated_augmentations"),
        help="Output directory for augmented images (default: generated_augmentations)",
    )
    parser.add_argument(
        "--count",
        "-n",
        type=int,
        default=5,
        metavar="N",
        help="Augmentations to create per input image (default: 5)",
    )
    parser.add_argument(
        "--ignore-ai-tools",
        dest="use_ai_tools",
        action="store_false",
        default=True,
        help=(
            "Skip AI-powered ops: OCR text avoidance, background replacement (rembg/Openverse), "
            "and neural style transfer"
        ),
    )
    parser.add_argument(
        "--strength",
        type=float,
        default=1.0,
        help="Augmentation strength (default: 1.0)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed for reproducible prompt selection",
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    if args.count < 1:
        print("Error: --count must be at least 1", file=sys.stderr)
        return 1

    input_images = collect_input_images(args.input)
    if not input_images:
        print(
            f"Error: no .jpg / .jpeg / .png images found in {args.input}",
            file=sys.stderr,
        )
        return 1

    if args.seed is not None:
        random.seed(args.seed)

    prompt_pool = build_prompt_pool(use_ai_tools=args.use_ai_tools)
    if not prompt_pool:
        print("Error: no prompts available for the selected mode", file=sys.stderr)
        return 1

    print("Loading embedding planner...")
    warmup_planner()

    run_id = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
    run_dir = args.output / run_id
    run_dir.mkdir(parents=True, exist_ok=True)

    manifest_sources: list[dict] = []
    total_created = 0
    total_requested = len(input_images) * args.count
    partial = False

    print(f"Found {len(input_images)} input image(s). Creating {args.count} augmentation(s) each.\n")

    for source_path in input_images:
        print(f"{source_path.name}:")
        source_image = Image.open(source_path)
        items, created = augment_image(
            source_path,
            source_image,
            run_dir,
            count=args.count,
            prompt_pool=prompt_pool,
            use_ai_tools=args.use_ai_tools,
            strength=args.strength,
        )
        total_created += created
        if created < args.count:
            partial = True
            print(
                f"  Warning: only {created}/{args.count} augmentations for {source_path.name}",
                file=sys.stderr,
            )
        manifest_sources.append(
            {
                "source": source_path.name,
                "count_requested": args.count,
                "count_created": created,
                "items": items,
            }
        )
        print()

    manifest_path = run_dir / "manifest.json"
    manifest_path.write_text(
        json.dumps(
            {
                "input_folder": str(args.input.resolve()),
                "output_folder": str(run_dir.resolve()),
                "images_found": len(input_images),
                "count_per_image": args.count,
                "total_requested": total_requested,
                "total_created": total_created,
                "use_ai_tools": args.use_ai_tools,
                "strength": args.strength,
                "sources": manifest_sources,
            },
            indent=2,
        ),
        encoding="utf-8",
    )

    print(f"Done. Saved {total_created} image(s) to {run_dir.resolve()}")
    print(f"Manifest: {manifest_path}")
    return 2 if partial else 0


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