"""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())