Package project as pip-installable augmenator.

Rename pipeline to augmenator, add pyproject.toml and PIL-first public API (augment/augment_batch), and keep Gradio as an optional UI extra.

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Dag Bjornberg
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -11,9 +11,60 @@ startup_duration_timeout: 1h
 pinned: false
 ---
 
-# Text-Driven Image Augmentation
+# augmenator
 
-Upload an image and type an instruction. Transforms are selected by **semantic similarity** between your prompt and a catalog of augmentation keywords (1–5 matches above a cosine-similarity threshold).
+Text-driven image augmentation. Transforms are selected by **semantic similarity** between your prompt and a catalog of augmentation keywords (1–5 matches above a cosine-similarity threshold).
+
+## Install
+
+```bash
+pip install augmenator
+# optional Gradio demo dependency
+pip install "augmenator[ui]"
+```
+
+Editable / local development:
+
+```powershell
+cd "C:\Users\Dag Bjornberg\AI\augmenator"
+python -m venv .venv
+.\.venv\Scripts\Activate.ps1
+pip install -e ".[ui]"
+```
+
+First run downloads the embedding model, RapidOCR, rembg, and style weights as needed.
+
+## Python API
+
+PIL in, PIL out — easy to loop over paired lists:
+
+```python
+from PIL import Image
+from augmenator import augment, augment_batch, warmup
+
+warmup()  # optional; loads the embedding planner once
+
+image = Image.open("photo.jpg")
+out = augment("make it brighter", image)
+
+# Batch: zip instructions with images
+outs = augment_batch(
+    ["blur softly", "cartoon style"],
+    [image, image],
+)
+
+# Or loop yourself
+results = [augment(text, img) for text, img in zip(texts, images)]
+```
+
+For planner metadata (tags, spatial ops, scores), use `augment_detailed(...)`.
+
+Batch folder CLI:
+
+```bash
+augmenator-batch --input ./photos --count 5
+augmenator-batch --input ./photos --count 3 --ignore-ai-tools
+```
 
 ## How planning works
 
@@ -59,19 +110,14 @@ Object insertion (not background swap):
 
 - `put a table behind it`
 
-## Local setup (Windows)
-
-Requires Python 3.10+.
+## Gradio demo
 
 ```powershell
-cd "C:\Users\Dag Bjornberg\AI\cursor_pilot"
-python -m venv .venv
-.\.venv\Scripts\Activate.ps1
-pip install -r requirements.txt
+pip install -e ".[ui]"
 python app.py
 ```
 
-First run downloads the embedding model, RapidOCR, and rembg u2netp weights.
+Live demo: [Hugging Face Space](https://huggingface.co/spaces/dagbjorn/text-driven-image-augmentation)
 
 ## Models
 

app.py

@@ -1,4 +1,4 @@
-import os
+import os
 
 # HF Spaces often set HTTP_PROXY; without NO_PROXY, Gradio's localhost health check fails.
 os.environ.setdefault("NO_PROXY", "localhost,127.0.0.1,::1")
@@ -7,10 +7,10 @@ os.environ.setdefault("no_proxy", "localhost,127.0.0.1,::1")
 import gradio as gr
 from PIL import Image
 
-from pipeline import run_pipeline
-from pipeline.background_replace import warmup as warmup_rembg
-from pipeline.planner import warmup
-from pipeline.text_regions import warmup as warmup_ocr
+from augmenator import run_pipeline
+from augmenator.background_replace import warmup as warmup_rembg
+from augmenator.planner import warmup
+from augmenator.text_regions import warmup as warmup_ocr
 
 _warmed_up = False
 
@@ -145,9 +145,9 @@ with gr.Blocks(title="Text-Driven Image Augmentation") as demo:
     gr.Markdown(
         "# Text-Driven Image Augmentation\n"
         "Upload an image and describe what you want. Transforms are chosen by **semantic similarity** "
-        "to augmentation keywords (1–5 matches above a similarity threshold). "
+        "to augmentation keywords (1–5 matches above a similarity threshold). "
         "Background replacement uses CC-licensed photos from Openverse "
-        "(web only — describe any scene, e.g. `replace background with sunset over Paris`)."
+        "(web only — describe any scene, e.g. `replace background with sunset over Paris`)."
     )
 
     with gr.Row():
@@ -192,3 +192,4 @@ if __name__ == "__main__":
         server_port=int(os.environ.get("PORT", 7860)),
         share=False,
     )
+

{pipeline → augmenator}/__init__.py

@@ -1,8 +1,25 @@
+"""augmenator — text-driven image augmentation."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
 from PIL import Image
 
-from pipeline.ai_tools import strip_ai_tools_from_plan
-from pipeline.augment import apply_augmentations
-from pipeline.planner import plan_from_instruction
+from augmenator.ai_tools import strip_ai_tools_from_plan
+from augmenator.augment import apply_augmentations
+from augmenator.planner import plan_from_instruction
+from augmenator.planner import warmup as warmup
+
+__all__ = [
+    "augment",
+    "augment_batch",
+    "augment_detailed",
+    "run_pipeline",
+    "warmup",
+]
+
+__version__ = "0.1.0"
 
 
 def run_pipeline(
@@ -61,3 +78,60 @@ def run_pipeline(
         "image": augmented,
         "use_ai_tools": use_ai_tools,
     }
+
+
+def augment(
+    instruction: str,
+    image: Image.Image,
+    *,
+    strength: float = 1.0,
+    use_ai_tools: bool = True,
+) -> Image.Image:
+    """Augment a single image from a text instruction. Returns a PIL image.
+
+    If the instruction is unsupported or stripped empty, returns ``image`` unchanged.
+    """
+    return run_pipeline(
+        image,
+        instruction,
+        strength=strength,
+        use_ai_tools=use_ai_tools,
+    )["image"]
+
+
+def augment_batch(
+    instructions: Sequence[str],
+    images: Sequence[Image.Image],
+    *,
+    strength: float = 1.0,
+    use_ai_tools: bool = True,
+) -> list[Image.Image]:
+    """Augment paired lists of instructions and images.
+
+    ``instructions`` and ``images`` must be the same length.
+    """
+    if len(instructions) != len(images):
+        raise ValueError(
+            f"instructions and images must be the same length "
+            f"(got {len(instructions)} and {len(images)})"
+        )
+    return [
+        augment(instruction, image, strength=strength, use_ai_tools=use_ai_tools)
+        for instruction, image in zip(instructions, images)
+    ]
+
+
+def augment_detailed(
+    instruction: str,
+    image: Image.Image,
+    *,
+    strength: float = 1.0,
+    use_ai_tools: bool = True,
+) -> dict:
+    """Like ``augment``, but returns the full result dict (tags, spatial, scores)."""
+    return run_pipeline(
+        image,
+        instruction,
+        strength=strength,
+        use_ai_tools=use_ai_tools,
+    )

{pipeline → augmenator}/ai_tools.py

@@ -1,8 +1,8 @@
-"""AI-powered pipeline features: OCR text avoidance, background replacement, neural style transfer."""
+"""AI-powered pipeline features: OCR text avoidance, background replacement, neural style transfer."""
 
 from copy import deepcopy
 
-from pipeline.style_transfer import STYLE_TAGS
+from augmenator.style_transfer import STYLE_TAGS
 
 AI_TOOL_KEYWORD_IDS = frozenset(
     {
@@ -46,3 +46,4 @@ def strip_ai_tools_from_plan(plan: dict) -> dict:
         if keyword_id not in AI_TOOL_KEYWORD_IDS
     ]
     return updated
+

{pipeline → augmenator}/augment.py

@@ -1,11 +1,11 @@
-import random
+import random
 
 import cv2
 import numpy as np
 from PIL import Image, ImageEnhance, ImageFilter, ImageOps
 
-from pipeline.spatial import apply_spatial_ops
-from pipeline.style_transfer import STYLE_MODELS, STYLE_TAGS, apply_style
+from augmenator.spatial import apply_spatial_ops
+from augmenator.style_transfer import STYLE_MODELS, STYLE_TAGS, apply_style
 
 ALLOWED_TAGS = {
     "brighten",
@@ -147,15 +147,15 @@ def _apply_tag(image: Image.Image, tag: str, strength: float) -> Image.Image:
         angle = random.uniform(0, 360)
         fill = (128, 128, 128) if image.mode == "RGB" else (128, 128, 128, 255)
         rotated = image.rotate(angle, expand=True, fillcolor=fill)
-        return rotated, f"rotate({angle:.1f}°)"
+        return rotated, f"rotate({angle:.1f}°)"
     if tag == "rotate_90_random":
         degrees = random.choice([90, 270])
         direction = "CCW" if degrees == 90 else "CW"
         return _rotate_cardinal(image, degrees), f"rotate_90({direction})"
     if tag == "rotate_left":
-        return _rotate_cardinal(image, 90), "rotate_left(90° CCW)"
+        return _rotate_cardinal(image, 90), "rotate_left(90° CCW)"
     if tag == "rotate_right":
-        return _rotate_cardinal(image, 270), "rotate_right(90° CW)"
+        return _rotate_cardinal(image, 270), "rotate_right(90° CW)"
     if tag == "rotate_180":
         return _rotate_cardinal(image, 180), "rotate_180"
     if tag == "flip":
@@ -266,3 +266,4 @@ def apply_augmentations(
         result = result.convert("RGB")
 
     return result, applied, applied_spatial
+

{pipeline → augmenator}/background_replace.py

@@ -1,10 +1,10 @@
-import re
+import re
 from typing import Literal
 
 from PIL import Image
 from rembg import new_session, remove
 
-from pipeline.background_web import fetch_public_background
+from augmenator.background_web import fetch_public_background
 
 BACKGROUND_TRIGGERS = (
     "replace background",
@@ -105,3 +105,4 @@ def replace_background(
         "requested_query": query,
         **pick_meta,
     }
+

augmenator/cli.py

@@ -0,0 +1,259 @@
+"""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())

{pipeline → augmenator}/embedding_planner.py

@@ -1,7 +1,7 @@
-import numpy as np
+import numpy as np
 from sentence_transformers import SentenceTransformer
 
-from pipeline.keyword_catalog import AUGMENT_KEYWORDS, AugmentKeyword
+from augmenator.keyword_catalog import AUGMENT_KEYWORDS, AugmentKeyword
 
 MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
 SIMILARITY_THRESHOLD = 0.38
@@ -128,3 +128,4 @@ def select_keywords(instruction: str) -> list[tuple[AugmentKeyword, float]]:
     ]
     deduped = _apply_mutual_exclusion(above_threshold)
     return deduped[:MAX_SELECTIONS]
+

{pipeline → augmenator}/planner.py

@@ -1,10 +1,10 @@
-import re
+import re
 
-from pipeline.background_replace import parse_background_query
-from pipeline.embedding_planner import select_keywords, warmup as warmup_embeddings
-from pipeline.keyword_catalog import AugmentKeyword
-from pipeline.spatial import infer_cover_count
-from pipeline.spatial_triggers import allows_spatial_keyword
+from augmenator.background_replace import parse_background_query
+from augmenator.embedding_planner import select_keywords, warmup as warmup_embeddings
+from augmenator.keyword_catalog import AugmentKeyword
+from augmenator.spatial import infer_cover_count
+from augmenator.spatial_triggers import allows_spatial_keyword
 
 SPATIAL_OPS = {
     "cover_avoid_text",
@@ -31,7 +31,7 @@ GENERATIVE_PATTERNS = (
 UNSUPPORTED_REASON = (
     "This instruction needs generative editing (adding or replacing objects/scenes). "
     "v1 supports color transforms, classical edits, procedural cutouts, text-aware covering, "
-    "and background replacement — not object insertion."
+    "and background replacement — not object insertion."
 )
 
 
@@ -351,3 +351,4 @@ def plan_from_instruction(instruction: str) -> dict:
             "Try describing an effect more clearly, or use one of the example prompts below."
         ),
     }
+

{pipeline → augmenator}/spatial.py

@@ -1,12 +1,12 @@
-import random
+import random
 import re
 
 import cv2
 import numpy as np
 from PIL import Image, ImageDraw
 
-from pipeline.background_replace import replace_background as do_replace_background
-from pipeline.text_regions import detect_text_boxes, find_safe_rect
+from augmenator.background_replace import replace_background as do_replace_background
+from augmenator.text_regions import detect_text_boxes, find_safe_rect
 
 SHAPES = ("rounded_rect", "ellipse", "rect")
 MAX_CUTOUT_COUNT = 4
@@ -356,3 +356,4 @@ def apply_spatial_ops(
         result, meta = apply_spatial_op(result, op, strength, instruction)
         applied.append(meta)
     return result, applied
+

{pipeline → augmenator}/style_transfer.py

@@ -1,4 +1,4 @@
-"""Neural style transfer: ONNX Model Zoo + PyTorch HF weights (CPU-friendly)."""
+"""Neural style transfer: ONNX Model Zoo + PyTorch HF weights (CPU-friendly)."""
 
 from __future__ import annotations
 
@@ -14,8 +14,8 @@ from huggingface_hub import hf_hub_download
 from PIL import Image
 from torchvision import transforms
 
-from pipeline.style_net import StyleNet
-from pipeline.transform_net import TransformNet
+from augmenator.style_net import StyleNet
+from augmenator.transform_net import TransformNet
 
 StyleBackend = Literal["onnx", "transformnet", "stylenet", "classical"]
 
@@ -239,3 +239,4 @@ def apply_style(image: Image.Image, tag: str, strength: float = 1.0) -> Image.Im
         styled = styled.resize((width, height), Image.Resampling.LANCZOS)
 
     return _blend_strength(original, styled, strength)
+

pyproject.toml

@@ -0,0 +1,53 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "augmenator"
+version = "0.1.0"
+description = "Text-driven image augmentation via embedding-matched keywords"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.10"
+authors = [{ name = "Dag Bjornberg" }]
+keywords = ["image", "augmentation", "style-transfer", "computer-vision"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Image Processing",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "torch>=2.1.0",
+    "torchvision>=0.16.0",
+    "transformers>=4.40.0",
+    "sentence-transformers>=2.7.0",
+    "Pillow>=10.0.0",
+    "numpy>=1.24.0",
+    "rapidocr-onnxruntime>=1.3.0",
+    "opencv-python-headless>=4.8.0",
+    "rembg[cpu]>=2.0.50",
+    "onnxruntime>=1.16.0",
+    "requests>=2.28.0",
+    "huggingface_hub>=0.20.0",
+]
+
+[project.optional-dependencies]
+ui = ["gradio>=5.23.1"]
+dev = ["build", "twine"]
+
+[project.urls]
+Homepage = "https://huggingface.co/spaces/dagbjorn/text-driven-image-augmentation"
+Repository = "https://huggingface.co/spaces/dagbjorn/text-driven-image-augmentation"
+
+[project.scripts]
+augmenator-batch = "augmenator.cli:main"
+
+[tool.setuptools.packages.find]
+include = ["augmenator*"]
+exclude = ["scripts*", "backgrounds*", "generated_augmentations*"]

requirements.txt

@@ -1,12 +1,2 @@
-torch>=2.1.0
-torchvision>=0.16.0
-transformers>=4.40.0
-sentence-transformers>=2.7.0
-gradio>=5.23.1
-Pillow>=10.0.0
-numpy>=1.24.0
-rapidocr-onnxruntime>=1.3.0
-opencv-python-headless>=4.8.0
-rembg[cpu]>=2.0.50
-onnxruntime>=1.16.0
-requests>=2.28.0
+# Core package (also installed via pyproject.toml). Gradio is Space / UI-only.
+-e .[ui]

scripts/generate_augmentations.py

@@ -1,267 +1,14 @@
-"""
+"""
 Batch-generate augmented images from every image in an input folder.
 
 Example:
   python scripts/generate_augmentations.py --input ./photos --count 5
-  python scripts/generate_augmentations.py --input ./photos --count 3 --ignore-ai-tools
+  augmenator-batch --input ./photos --count 3 --ignore-ai-tools
 """
 
 from __future__ import annotations
 
-import argparse
-import json
-import random
-import re
-import sys
-from datetime import datetime, timezone
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-
-from PIL import Image
-
-from pipeline import run_pipeline
-from pipeline.ai_tools import AI_TOOL_KEYWORD_IDS
-from pipeline.keyword_catalog import AUGMENT_KEYWORDS
-from pipeline.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() -> 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 passed to the pipeline (default: 1.0)",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=None,
-        help="Random seed for reproducible prompt selection",
-    )
-    return parser.parse_args()
-
-
-def main() -> int:
-    args = parse_args()
-    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
-
+from augmenator.cli import main
 
 if __name__ == "__main__":
     raise SystemExit(main())

scripts/validate_planner.py

@@ -1,13 +1,10 @@
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+import sys
 
 from PIL import Image
 
-from pipeline.embedding_planner import SIMILARITY_THRESHOLD
-from pipeline.planner import plan_from_instruction, warmup
-from pipeline.spatial import apply_spatial_ops
+from augmenator.embedding_planner import SIMILARITY_THRESHOLD
+from augmenator.planner import plan_from_instruction, warmup
+from augmenator.spatial import apply_spatial_ops
 
 warmup()
 
@@ -127,3 +124,4 @@ if failed:
     raise SystemExit(f"{failed} test(s) failed")
 
 print("All validation tests passed")
+