augmenator/background_replace.py

import re
from typing import Literal

from PIL import Image
from rembg import new_session, remove

from augmenator.background_web import fetch_public_background

BACKGROUND_TRIGGERS = (
    "replace background",
    "change background",
    "swap background",
    "new background",
    "replace by",
)

QUERY_PREFIXES = (
    r"replace background with\s+",
    r"replace background by\s+",
    r"replace by\s+",
    r"change background to\s+",
    r"change background with\s+",
    r"new background\s+",
    r"swap background with\s+",
    r"swap background for\s+",
)

RANDOM_QUERY_WORDS = {"anything", "something", "random", "whatever", "anything at all"}

TRANSFORM_TAIL = re.compile(
    r"\s+and\s+(rotate|rotation|turn|brighter|brighten|darker|darken|more contrast|"
    r"less contrast|flip|blur|sepia|saturate|invert|warmer|cooler|cover .+)$",
    re.I,
)

_session = None


def warmup() -> None:
    global _session
    if _session is None:
        _session = new_session("u2netp")


def is_background_replace_instruction(instruction: str) -> bool:
    lowered = instruction.lower()
    return any(trigger in lowered for trigger in BACKGROUND_TRIGGERS)


def parse_background_query(instruction: str) -> str | None | Literal[False]:
    """
    Return the Openverse search query, None for random, or False if not a bg-replace instruction.
    """
    if not is_background_replace_instruction(instruction):
        return False

    for pattern in QUERY_PREFIXES:
        match = re.search(pattern + r"(.+)", instruction, re.I)
        if not match:
            continue
        query = _clean_query(match.group(1))
        if query and query.lower() not in RANDOM_QUERY_WORDS:
            return query
        return None

    return None


def _clean_query(text: str) -> str:
    query = text.strip().strip('"').strip("'")
    query = TRANSFORM_TAIL.sub("", query).strip()
    query = re.sub(r"^(a|an|the)\s+", "", query, flags=re.I)
    return query.strip()


def pick_background(
    query: str | None,
    size: tuple[int, int],
) -> tuple[Image.Image, dict]:
    return fetch_public_background(query, size)


def extract_foreground(image: Image.Image) -> Image.Image:
    warmup()
    fg = remove(image.convert("RGB"), session=_session)
    return fg.convert("RGBA")


def replace_background(
    image: Image.Image,
    query: str | None = None,
    instruction: str = "",
) -> tuple[Image.Image, dict]:
    if instruction and query is None and is_background_replace_instruction(instruction):
        parsed = parse_background_query(instruction)
        if parsed is not False:
            query = parsed

    fg = extract_foreground(image)
    bg, pick_meta = pick_background(query, fg.size)
    canvas = bg.convert("RGBA")
    canvas.paste(fg, (0, 0), fg)
    return canvas, {
        "op": "replace_background",
        "requested_query": query,
        **pick_meta,
    }