"""Diverse synthetic masks used for context-aware inpainting training.""" from __future__ import annotations import math import random from PIL import Image, ImageChops, ImageDraw, ImageFilter MASK_KINDS = ( "brush", "multi_brush", "rectangle", "ellipse", "polygon", "multi_region", "outpaint", ) def _area_fraction(mask: Image.Image) -> float: histogram = mask.convert("L").histogram() white_sum = sum(value * count for value, count in enumerate(histogram)) return white_sum / (255.0 * mask.width * mask.height) def _random_box( size: tuple[int, int], rng: random.Random, target_area: float, ) -> tuple[int, int, int, int]: width, height = size aspect = math.exp(rng.uniform(math.log(0.35), math.log(2.85))) box_width = min(width, max(8, round(math.sqrt(target_area * aspect)))) box_height = min(height, max(8, round(math.sqrt(target_area / aspect)))) left = rng.randint(0, max(0, width - box_width)) top = rng.randint(0, max(0, height - box_height)) return left, top, left + box_width, top + box_height def _draw_brush( draw: ImageDraw.ImageDraw, size: tuple[int, int], rng: random.Random, *, strokes: int, ) -> None: width, height = size for _ in range(strokes): stroke_width = max(6, round(min(width, height) * rng.uniform(0.035, 0.16))) point_count = rng.randint(3, 8) x = rng.uniform(0, width - 1) y = rng.uniform(0, height - 1) points: list[tuple[float, float]] = [(x, y)] angle = rng.uniform(0, 2 * math.pi) for _ in range(point_count - 1): angle += rng.uniform(-1.25, 1.25) distance = rng.uniform(0.06, 0.24) * min(width, height) x = min(width - 1, max(0, x + math.cos(angle) * distance)) y = min(height - 1, max(0, y + math.sin(angle) * distance)) points.append((x, y)) draw.line(points, fill=255, width=stroke_width, joint="curve") radius = stroke_width / 2 for px, py in points: draw.ellipse((px - radius, py - radius, px + radius, py + radius), fill=255) def _draw_outpaint( draw: ImageDraw.ImageDraw, size: tuple[int, int], rng: random.Random, target_fraction: float, ) -> None: width, height = size sides = rng.sample(("left", "right", "top", "bottom"), rng.choice((1, 1, 1, 2))) remaining = max(0.02, target_fraction) for index, side in enumerate(sides): share = remaining if index == len(sides) - 1 else remaining * rng.uniform(0.35, 0.7) if side in {"left", "right"}: thickness = max(4, min(width, round(share * width))) box = (0, 0, thickness, height) if side == "left" else (width - thickness, 0, width, height) else: thickness = max(4, min(height, round(share * height))) box = (0, 0, width, thickness) if side == "top" else (0, height - thickness, width, height) draw.rectangle(box, fill=255) remaining = max(0.0, remaining - share) def _draw_candidate( size: tuple[int, int], rng: random.Random, *, kind: str, target_fraction: float, ) -> Image.Image: width, height = size mask = Image.new("L", size, 0) draw = ImageDraw.Draw(mask) target_area = target_fraction * width * height if kind == "rectangle": draw.rounded_rectangle( _random_box(size, rng, target_area), radius=rng.randint(0, max(1, round(min(width, height) * 0.08))), fill=255, ) elif kind == "ellipse": draw.ellipse(_random_box(size, rng, target_area), fill=255) elif kind == "polygon": center_x = rng.uniform(width * 0.2, width * 0.8) center_y = rng.uniform(height * 0.2, height * 0.8) radius = math.sqrt(target_area / math.pi) points = [] point_count = rng.randint(5, 10) for index in range(point_count): angle = (2 * math.pi * index / point_count) + rng.uniform(-0.25, 0.25) local_radius = radius * rng.uniform(0.65, 1.35) points.append( ( min(width - 1, max(0, center_x + math.cos(angle) * local_radius)), min(height - 1, max(0, center_y + math.sin(angle) * local_radius)), ) ) draw.polygon(points, fill=255) elif kind == "brush": _draw_brush(draw, size, rng, strokes=1) elif kind == "multi_brush": _draw_brush(draw, size, rng, strokes=rng.randint(2, 5)) elif kind == "multi_region": region_count = rng.randint(2, 5) for _ in range(region_count): region_area = target_area * rng.uniform(0.12, 0.5) box = _random_box(size, rng, region_area) if rng.random() < 0.55: draw.ellipse(box, fill=255) else: draw.rounded_rectangle(box, radius=rng.randint(2, 24), fill=255) elif kind == "outpaint": _draw_outpaint(draw, size, rng, target_fraction) else: raise ValueError(f"Unsupported mask kind: {kind}") if kind not in {"rectangle", "outpaint"} and rng.random() < 0.35: # A small close operation removes pinholes without making every edge # unnaturally geometric. mask = mask.filter(ImageFilter.MaxFilter(rng.choice((3, 5, 7)))) return mask.point(lambda value: 255 if value >= 128 else 0, mode="L") def random_mask( size: tuple[int, int], rng: random.Random, *, min_area: float = 0.04, max_area: float = 0.65, kind: str | None = None, ) -> Image.Image: """Return a binary mask where white means regenerate. The sampler deliberately mixes object-like regions, free-form user brush strokes, disconnected edits, and edge/outpainting masks. Candidates are retried so the actual white area—not only a geometric estimate—falls close to the requested range. """ if not 0.0 < min_area < max_area < 1.0: raise ValueError("mask area bounds must satisfy 0 < min < max < 1") if kind is not None and kind not in MASK_KINDS: raise ValueError(f"Unsupported mask kind: {kind}") best_mask: Image.Image | None = None best_distance = float("inf") target = rng.uniform(min_area, max_area) for _ in range(16): selected_kind = kind or rng.choices( MASK_KINDS, weights=(24, 18, 12, 10, 12, 16, 8), k=1, )[0] candidate = _draw_candidate(size, rng, kind=selected_kind, target_fraction=target) area = _area_fraction(candidate) if min_area <= area <= max_area: return candidate distance = min(abs(area - min_area), abs(area - max_area)) if distance < best_distance: best_mask = candidate best_distance = distance target = rng.uniform(min_area, max_area) if best_mask is None: raise RuntimeError("mask generation did not produce a candidate") return best_mask def apply_mask(image: Image.Image, mask: Image.Image) -> Image.Image: """Black out the pixels that the inpainting model must regenerate.""" image = image.convert("RGB") keep = ImageChops.invert(mask.convert("L")) return Image.composite(image, Image.new("RGB", image.size), keep) def mask_area_fraction(mask: Image.Image) -> float: """Expose the exact white-area fraction for validation and reporting.""" return _area_fraction(mask)