Spaces:
Build error
Build error
File size: 2,894 Bytes
7025ca1 44a3896 7025ca1 44a3896 7025ca1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | 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,
}
|