augmenator/__init__.py

"""augmenator — text-driven image augmentation."""

from __future__ import annotations

from collections.abc import Sequence

from PIL import Image

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(
    image: Image.Image,
    instruction: str,
    strength: float = 1.0,
    *,
    use_ai_tools: bool = True,
) -> dict:
    plan = plan_from_instruction(instruction)
    if not plan["supported"]:
        return {
            "supported": False,
            "reason": plan["reason"],
            "applied_tags": [],
            "applied_spatial": [],
            "text_regions_found": 0,
            "image": image,
            "use_ai_tools": use_ai_tools,
        }

    if not use_ai_tools:
        plan = strip_ai_tools_from_plan(plan)
        if not plan.get("tags") and not plan.get("spatial"):
            return {
                "supported": False,
                "reason": "Instruction only matched AI-tool ops (disabled with --ignore-ai-tools).",
                "applied_tags": [],
                "applied_spatial": [],
                "text_regions_found": 0,
                "image": image,
                "use_ai_tools": use_ai_tools,
            }

    augmented, applied_tags, applied_spatial = apply_augmentations(
        image,
        plan["tags"],
        strength=strength,
        spatial_ops=plan.get("spatial", []),
        instruction=instruction,
    )

    text_regions = 0
    if applied_spatial:
        text_regions = max(s.get("text_regions_found", 0) for s in applied_spatial)

    return {
        "supported": True,
        "instruction": instruction,
        "planned_tags": plan["tags"],
        "planned_spatial": plan.get("spatial", []),
        "planned_scores": plan.get("scores", []),
        "applied_tags": applied_tags,
        "applied_spatial": applied_spatial,
        "text_regions_found": text_regions,
        "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,
    )