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 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