augmenator/spatial.py

import random
import re

import cv2
import numpy as np
from PIL import Image, ImageDraw

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


def infer_fill(instruction: str) -> str:
    lowered = instruction.lower()
    if any(word in lowered for word in ("blur", "soft", "soften")):
        return "blur"
    return "solid"


def _patch_size(
    image: Image.Image,
    strength: float,
    scale_w: float = 1.0,
    scale_h: float = 1.0,
) -> tuple[int, int]:
    w, h = image.size
    ratio = 0.18 + 0.10 * max(0.5, min(1.5, strength) - 0.5)
    patch_w = max(16, int(w * ratio * scale_w))
    patch_h = max(16, int(h * ratio * scale_h))
    return patch_w, patch_h


def _varied_patch_size(image: Image.Image, strength: float) -> tuple[int, int]:
    scale_w = random.uniform(0.55, 1.45)
    scale_h = random.uniform(0.55, 1.45)
    return _patch_size(image, strength, scale_w=scale_w, scale_h=scale_h)


def _scaled_patch_size(image: Image.Image, strength: float) -> tuple[int, int]:
    scale = random.uniform(0.65, 1.35)
    return _patch_size(image, strength, scale_w=scale, scale_h=scale)


def infer_cover_count(instruction: str, explicit: int | None = None) -> int:
    if explicit is not None:
        return min(3, max(1, explicit))
    return _infer_count(instruction, "cover", default=1, max_count=3)


def _infer_count(instruction: str, kind: str, default: int = 1, max_count: int = 4) -> int:
    lowered = instruction.lower()
    patterns = (
        rf"\b(\d+)\s+{kind}s?\b",
        rf"\b(\d+)\s+{kind}\b",
        r"\b(\d+)\s+(hole|holes|patch|patches)\b",
    )
    for pattern in patterns:
        match = re.search(pattern, lowered)
        if match:
            return min(max_count, max(1, int(match.group(1))))

    if any(word in lowered for word in ("several", "multiple", "many")):
        return random.randint(2, max_count)
    if "few" in lowered:
        return random.randint(2, min(3, max_count))
    return default


def _resolve_cutout_count(op: dict, instruction: str) -> int:
    explicit = op.get("cutout_count")
    if explicit is not None:
        return min(MAX_CUTOUT_COUNT, max(1, int(explicit)))

    lowered = instruction.lower()
    if any(
        phrase in lowered
        for phrase in ("one cutout", "a cutout", "single cutout", "one hole", "a hole")
    ):
        return 1
    if re.search(r"\bone\b", lowered) and "cutout" in lowered:
        return 1

    patterns = (
        r"\b(\d+)\s+cutouts?\b",
        r"\b(\d+)\s+holes?\b",
        r"\b(\d+)\s+(hole|holes|patch|patches)\b",
    )
    for pattern in patterns:
        match = re.search(pattern, lowered)
        if match:
            return min(MAX_CUTOUT_COUNT, max(1, int(match.group(1))))

    if any(word in lowered for word in ("several", "multiple", "many")):
        return random.randint(2, MAX_CUTOUT_COUNT)

    return random.randint(1, MAX_CUTOUT_COUNT)


def _pick_shape(instruction: str, varied: bool = False, force_varied: bool = False) -> str:
    if varied or force_varied:
        return random.choice(SHAPES)

    lowered = instruction.lower()
    if any(word in lowered for word in ("circle", "round", "oval", "ellipse")):
        return "ellipse"
    if any(word in lowered for word in ("rectangle", "square", "rect")):
        return "rect"
    return "rounded_rect"


def _pick_rect(
    image: Image.Image,
    avoid_text: bool,
    patch_w: int,
    patch_h: int,
) -> tuple[tuple[int, int, int, int], int, bool]:
    text_boxes = detect_text_boxes(image) if avoid_text else []
    if avoid_text and text_boxes:
        rect, used_fallback = find_safe_rect(image.size, text_boxes, patch_w, patch_h)
        return rect, len(text_boxes), used_fallback

    w, h = image.size
    max_x = max(0, w - patch_w)
    max_y = max(0, h - patch_h)
    x = random.randint(0, max_x) if max_x > 0 else 0
    y = random.randint(0, max_y) if max_y > 0 else 0
    return (x, y, patch_w, patch_h), len(text_boxes), False


def cover_patch(image: Image.Image, rect: tuple[int, int, int, int], fill: str = "blur") -> Image.Image:
    result = image.copy()
    x, y, w, h = rect
    region = result.crop((x, y, x + w, y + h))

    if fill == "blur":
        arr = np.array(region)
        blurred = cv2.GaussianBlur(arr, (0, 0), sigmaX=8, sigmaY=8)
        region = Image.fromarray(blurred)
    else:
        region = Image.new("RGB", (w, h), color=(140, 140, 140))

    result.paste(region, (x, y))
    return result


def _draw_shape_mask(draw: ImageDraw.ImageDraw, xy: tuple, shape: str) -> None:
    x, y, w, h = xy
    box = (x, y, x + w, y + h)
    if shape == "ellipse":
        draw.ellipse(box, fill=255)
    elif shape == "rect":
        draw.rectangle(box, fill=255)
    else:
        radius = min(w, h) // 6
        draw.rounded_rectangle(box, radius=radius, fill=255)


def procedural_cutout(
    image: Image.Image,
    rect: tuple[int, int, int, int],
    shape: str = "rounded_rect",
    style: str = "solid",
    color: tuple[int, int, int] = (255, 200, 80),
) -> Image.Image:
    x, y, w, h = rect
    result = image.convert("RGBA")

    if style == "transparent":
        mask = Image.new("L", result.size, 0)
        draw = ImageDraw.Draw(mask)
        _draw_shape_mask(draw, (x, y, w, h), shape)
        arr = np.array(result)
        arr[..., 3] = np.where(np.array(mask) > 0, 0, arr[..., 3])
        return Image.fromarray(arr)

    overlay = Image.new("RGBA", result.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)
    fill = (*color, 255)
    outline = tuple(max(0, c - 40) for c in color) + (255,)
    box = (x, y, x + w, y + h)

    if shape == "ellipse":
        draw.ellipse(box, fill=fill, outline=outline, width=3)
    elif shape == "rect":
        draw.rectangle(box, fill=fill, outline=outline, width=3)
    else:
        radius = min(w, h) // 6
        draw.rounded_rectangle(box, radius=radius, fill=fill, outline=outline, width=3)

    return Image.alpha_composite(result, overlay)


def apply_cutouts(
    image: Image.Image,
    op: dict,
    strength: float,
    instruction: str,
) -> tuple[Image.Image, list[dict], int, int, bool]:
    style = op.get("cutout_style", "solid")
    color = tuple(op.get("cutout_color", (255, 200, 80)))
    varied_shapes = op.get("varied_shapes", True)
    cutout_count = _resolve_cutout_count(op, instruction)
    avoid_text = op.get("avoid_text_cutouts", op.get("avoid_text", False))

    result = image.copy()
    cutout_details: list[dict] = []
    text_count = 0
    used_fallback = False

    for _ in range(cutout_count):
        patch_w, patch_h = _varied_patch_size(result, strength)
        rect, found, fallback = _pick_rect(
            result, avoid_text=avoid_text, patch_w=patch_w, patch_h=patch_h
        )
        text_count = max(text_count, found)
        used_fallback = used_fallback or fallback
        use_varied = varied_shapes and cutout_count > 1
        shape = _pick_shape(instruction, varied=varied_shapes, force_varied=use_varied)
        result = procedural_cutout(
            result,
            rect,
            shape=shape,
            style=style,
            color=color,
        )
        cutout_details.append({"rect": rect, "shape": shape, "style": style})

    return result, cutout_details, cutout_count, text_count, used_fallback


def apply_cover_and_cutouts(
    image: Image.Image,
    op: dict,
    strength: float,
    instruction: str,
) -> tuple[Image.Image, dict]:
    fill = op.get("fill") or infer_fill(instruction)
    avoid_text_for_covers = op.get("avoid_text", False)
    cover_count = infer_cover_count(instruction, op.get("cover_count"))

    result = image.copy()
    cover_rects: list[tuple[int, int, int, int]] = []
    text_count = 0
    used_fallback = False

    for _ in range(cover_count):
        patch_w, patch_h = _scaled_patch_size(result, strength)
        rect, found, fallback = _pick_rect(result, avoid_text_for_covers, patch_w, patch_h)
        text_count = max(text_count, found)
        used_fallback = used_fallback or fallback
        result = cover_patch(result, rect, fill=fill)
        cover_rects.append(rect)

    result, cutout_details, cutout_count, cutout_text_count, cutout_fallback = apply_cutouts(
        result, op, strength, instruction
    )
    text_count = max(text_count, cutout_text_count)
    used_fallback = used_fallback or cutout_fallback

    return result, {
        "op": op.get("op", "cover_and_cutout"),
        "cover_rects": cover_rects,
        "cutouts": cutout_details,
        "fill": fill,
        "cutout_style": op.get("cutout_style", "solid"),
        "text_regions_found": text_count,
        "used_fallback": used_fallback,
        "cover_count": cover_count,
        "cutout_count": cutout_count,
    }


def apply_spatial_op(
    image: Image.Image,
    op: dict,
    strength: float,
    instruction: str,
) -> tuple[Image.Image, dict]:
    op_name = op.get("op", "")
    if op_name == "replace_background":
        result, meta = do_replace_background(
            image,
            query=op.get("query"),
            instruction=instruction,
        )
        return result, meta

    if op_name in ("cover_and_cutout", "cover_and_cutout_avoid_text"):
        op_copy = dict(op)
        if op_name == "cover_and_cutout_avoid_text":
            op_copy["avoid_text"] = True
            op_copy["avoid_text_cutouts"] = True
            op_copy["varied_shapes"] = True
        return apply_cover_and_cutouts(image, op_copy, strength, instruction)

    fill = op.get("fill") or infer_fill(instruction)
    avoid_text = op.get("avoid_text", op_name.endswith("avoid_text"))
    patch_w, patch_h = _patch_size(image, strength)

    rect, text_count, used_fallback = _pick_rect(image, avoid_text, patch_w, patch_h)
    meta = {
        "op": op_name,
        "rect": rect,
        "fill": fill,
        "cutout_style": op.get("cutout_style"),
        "text_regions_found": text_count,
        "used_fallback": used_fallback,
    }

    if op_name in ("cover_avoid_text", "cover_random", "cover"):
        result = cover_patch(image, rect, fill=fill)
    elif op_name in ("add_cutout", "cutout"):
        cutout_op = dict(op)
        if op.get("avoid_text"):
            cutout_op["avoid_text_cutouts"] = True
        result, cutout_details, cutout_count, text_count, used_fallback = apply_cutouts(
            image, cutout_op, strength, instruction
        )
        meta = {
            "op": op_name,
            "cutouts": cutout_details,
            "cutout_count": cutout_count,
            "cutout_style": op.get("cutout_style", "solid"),
            "text_regions_found": text_count,
            "used_fallback": used_fallback,
        }
    else:
        return image, meta

    return result, meta


def apply_spatial_ops(
    image: Image.Image,
    spatial_ops: list[dict],
    strength: float,
    instruction: str,
) -> tuple[Image.Image, list[dict]]:
    priority = {
        "replace_background": 0,
        "cover_and_cutout": 1,
        "cover_and_cutout_avoid_text": 1,
        "cover_avoid_text": 1,
        "cover_random": 1,
        "cover": 1,
        "add_cutout": 2,
        "cutout": 2,
    }
    ordered = sorted(spatial_ops, key=lambda op: priority.get(op.get("op", ""), 99))

    result = image
    applied = []
    for op in ordered:
        result, meta = apply_spatial_op(result, op, strength, instruction)
        applied.append(meta)
    return result, applied