355 lines
13 KiB
Python
355 lines
13 KiB
Python
"""
|
||
Procedural mask generation for inpainting training.
|
||
|
||
Provides three mask types that can be combined:
|
||
- cloud : fractional Brownian motion turbulence (layered smooth noise)
|
||
- polygon : random convex or irregular polygon
|
||
- shape : original rectangle / ellipse (kept for backwards compat)
|
||
|
||
All functions return a PIL Image in mode "L" (0 = keep, 255 = mask/inpaint).
|
||
|
||
| Function | Description |
|
||
|---|---|
|
||
| `cloud_mask(w, h, ...)` | Organic blob shapes via fractional Brownian motion (layered noise) |
|
||
| `polygon_mask(w, h, ...)` | Random convex polygons |
|
||
| `shape_mask(w, h, ...)` | Axis-aligned rectangles and ellipses |
|
||
| `combine_masks(*masks)` | Union or intersection of any of the above |
|
||
| `random_mask(w, h, ...)` | Randomly selects and optionally combines mask types |
|
||
|
||
To visually inspect mask output, run:
|
||
|
||
python3 tests/visualize_masks.py --data-dir /path/to/images --out-dir ./mask_viz
|
||
|
||
This produces gallery PNG files for each mask type, showing the raw mask alongside the mask applied to
|
||
source images. If no `--data-dir` is provided (or the directory is empty), synthetic test images
|
||
are used as the background.
|
||
|
||
"""
|
||
|
||
import math
|
||
import random
|
||
from typing import Optional, Tuple
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image, ImageDraw
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Cloud / turbulence masks
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _smooth_noise(w: int, h: int, scale: float, rng: np.random.Generator) -> np.ndarray:
|
||
"""Generate a single octave of smooth noise by upsampling small random values."""
|
||
small_w = max(2, int(w / scale))
|
||
small_h = max(2, int(h / scale))
|
||
small = rng.random((small_h, small_w)).astype(np.float32)
|
||
return cv2.resize(small, (w, h), interpolation=cv2.INTER_CUBIC)
|
||
|
||
|
||
def cloud_mask(
|
||
width: int,
|
||
height: int,
|
||
octaves: int = 4,
|
||
persistence: float = 0.5,
|
||
base_scale: float = 16.0,
|
||
threshold: float = 0.2,
|
||
seed: Optional[int] = None,
|
||
) -> Image.Image:
|
||
"""
|
||
Fractional Brownian Motion cloud mask.
|
||
|
||
Parameters
|
||
----------
|
||
width, height : output size in pixels
|
||
octaves : number of noise layers; more = finer detail, more = grainier
|
||
persistence : amplitude falloff per octave (0..1); lower = smoother / fewer details
|
||
base_scale : spatial scale of the coarsest octave — higher = larger blobs
|
||
threshold : fraction of image area to mask (0..1); keep in ~0.1–0.3 for realistic patches
|
||
seed : optional RNG seed for reproducibility
|
||
"""
|
||
rng = np.random.default_rng(seed)
|
||
|
||
noise = np.zeros((height, width), dtype=np.float32)
|
||
amplitude = 1.0
|
||
total_amplitude = 0.0
|
||
scale = base_scale
|
||
|
||
for _ in range(octaves):
|
||
noise += amplitude * _smooth_noise(width, height, scale, rng)
|
||
total_amplitude += amplitude
|
||
amplitude *= persistence
|
||
scale *= 2.0
|
||
|
||
noise /= total_amplitude # normalise to [0, 1]
|
||
cutoff = np.quantile(noise, 1.0 - threshold)
|
||
binary = (noise >= cutoff).astype(np.uint8) * 255
|
||
|
||
return Image.fromarray(binary, mode="L")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Polygon masks
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _random_convex_polygon(
|
||
cx: float, cy: float, radius: float, n_points: int, irregularity: float, rng: random.Random
|
||
) -> list:
|
||
"""Return vertices of a roughly convex polygon centred at (cx, cy)."""
|
||
angles = sorted(rng.uniform(0, 2 * math.pi) for _ in range(n_points))
|
||
points = []
|
||
for angle in angles:
|
||
r = radius * rng.uniform(1.0 - irregularity, 1.0)
|
||
points.append((cx + r * math.cos(angle), cy + r * math.sin(angle)))
|
||
return points
|
||
|
||
|
||
def polygon_mask(
|
||
width: int,
|
||
height: int,
|
||
n_points: int = 6,
|
||
min_coverage: float = 0.1,
|
||
max_coverage: float = 0.6,
|
||
irregularity: float = 0.4,
|
||
n_polygons: int = 1,
|
||
seed: Optional[int] = None,
|
||
) -> Image.Image:
|
||
"""
|
||
Random polygon mask.
|
||
|
||
Parameters
|
||
----------
|
||
width, height : output size in pixels
|
||
n_points : number of vertices per polygon
|
||
min_coverage : minimum fraction of image area each polygon can cover
|
||
max_coverage : maximum fraction of image area each polygon can cover
|
||
irregularity : how irregular the polygon is (0 = circle, 1 = very jagged)
|
||
n_polygons : how many polygons to draw
|
||
seed : optional RNG seed
|
||
"""
|
||
rng = random.Random(seed)
|
||
mask = Image.new("L", (width, height), 0)
|
||
draw = ImageDraw.Draw(mask)
|
||
|
||
min_r = math.sqrt(min_coverage * width * height / math.pi)
|
||
max_r = math.sqrt(max_coverage * width * height / math.pi)
|
||
|
||
for _ in range(n_polygons):
|
||
radius = rng.uniform(min_r, max_r)
|
||
cx = rng.uniform(radius * 0.2, width - radius * 0.2)
|
||
cy = rng.uniform(radius * 0.2, height - radius * 0.2)
|
||
pts = _random_convex_polygon(cx, cy, radius, n_points, irregularity, rng)
|
||
draw.polygon(pts, fill=255)
|
||
|
||
return mask
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Shape masks (rectangles / ellipses — original behaviour)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def shape_mask(
|
||
width: int,
|
||
height: int,
|
||
min_coverage: float = 0.1,
|
||
max_coverage: float = 0.4,
|
||
shape: str = "random",
|
||
seed: Optional[int] = None,
|
||
) -> Image.Image:
|
||
"""
|
||
Rectangle or ellipse mask (mirrors the original random_mask behaviour).
|
||
|
||
Parameters
|
||
----------
|
||
shape : "rectangle", "ellipse", or "random"
|
||
"""
|
||
rng = random.Random(seed)
|
||
mask = Image.new("L", (width, height), 0)
|
||
draw = ImageDraw.Draw(mask)
|
||
|
||
area = width * height
|
||
size_w = rng.randint(int(width * min_coverage ** 0.5), int(width * max_coverage ** 0.5))
|
||
size_h = rng.randint(int(height * min_coverage ** 0.5), int(height * max_coverage ** 0.5))
|
||
|
||
x1 = rng.randint(0, width - size_w)
|
||
y1 = rng.randint(0, height - size_h)
|
||
x2, y2 = x1 + size_w, y1 + size_h
|
||
|
||
use_ellipse = (shape == "ellipse") or (shape == "random" and rng.random() < 0.5)
|
||
if use_ellipse:
|
||
draw.ellipse([x1, y1, x2, y2], fill=255)
|
||
else:
|
||
draw.rectangle([x1, y1, x2, y2], fill=255)
|
||
|
||
return mask
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Wobbly ellipse mask (fBm radius variation — good default for sampling)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def wobbly_ellipse_mask(
|
||
width: int,
|
||
height: int,
|
||
coverage: float = 0.2,
|
||
wobble_scale: float = 0.5,
|
||
octaves: int = 4,
|
||
persistence: float = 0.5,
|
||
n_points: int = 128,
|
||
seed: Optional[int] = None,
|
||
) -> Image.Image:
|
||
"""
|
||
Ellipse mask whose radius varies along the circumference using periodic
|
||
fractional Brownian motion noise, producing organic blob-like shapes.
|
||
|
||
Unlike cloud masks (which can produce large, fragmented regions), this
|
||
always yields a single connected region, making it well-suited as a
|
||
default mask for inpainting sampling and inference.
|
||
|
||
Parameters
|
||
----------
|
||
width, height : output size in pixels
|
||
coverage : approximate fraction of the image area to mask (0..1)
|
||
wobble_scale : radius variation relative to the base radius.
|
||
Default 0.25 means radius varies ±25% around the base.
|
||
octaves : number of harmonic layers in the radius noise
|
||
persistence : amplitude falloff per octave (0..1)
|
||
n_points : number of polygon vertices used to approximate the shape
|
||
seed : optional RNG seed for reproducibility
|
||
"""
|
||
rng = np.random.RandomState(seed if seed is not None else np.random.randint(0, 2**31))
|
||
|
||
# Center — biased toward image interior so the ellipse stays mostly inside
|
||
cx = rng.uniform(0.25, 0.75) * width
|
||
cy = rng.uniform(0.25, 0.75) * height
|
||
|
||
# Ellipse aspect ratio and rotation
|
||
aspect = rng.uniform(0.6, 1.4) # ratio of semi-major to semi-minor
|
||
rotation = rng.uniform(0, math.pi)
|
||
|
||
# Semi-axes derived from target coverage area (≈ π·a·b)
|
||
area_target = coverage * width * height
|
||
r_base = math.sqrt(area_target / (math.pi * aspect))
|
||
a = r_base * math.sqrt(aspect) # semi-major axis
|
||
b = r_base / math.sqrt(aspect) # semi-minor axis
|
||
|
||
# Periodic fBm using harmonic cosines so the noise wraps seamlessly at 0 = 2π
|
||
thetas = np.linspace(0, 2 * math.pi, n_points, endpoint=False)
|
||
noise = np.zeros(n_points, dtype=np.float32)
|
||
amplitude = 1.0
|
||
total_amplitude = 0.0
|
||
for k in range(1, octaves + 1):
|
||
phase = rng.uniform(0, 2 * math.pi)
|
||
noise += amplitude * np.cos(k * thetas + phase)
|
||
total_amplitude += amplitude
|
||
amplitude *= persistence
|
||
noise /= total_amplitude
|
||
std = noise.std()
|
||
if std > 1e-6:
|
||
noise = noise / std * wobble_scale # normalise so ±wobble_scale bounds hold approximately
|
||
|
||
# Build polygon: for each sample angle, compute the ellipse radius at that
|
||
# angle (accounting for rotation), then scale by the noise modulation.
|
||
points = []
|
||
for i, theta in enumerate(thetas):
|
||
t_rot = theta - rotation
|
||
r = (a * b) / math.sqrt((b * math.cos(t_rot)) ** 2 + (a * math.sin(t_rot)) ** 2)
|
||
r_wobbled = max(r * (1.0 + noise[i]), 1.0)
|
||
points.append((cx + r_wobbled * math.cos(theta), cy + r_wobbled * math.sin(theta)))
|
||
|
||
img = Image.new("L", (width, height), 0)
|
||
ImageDraw.Draw(img).polygon(points, fill=255)
|
||
return img
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Combiner
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def combine_masks(*masks: Image.Image, mode: str = "union") -> Image.Image:
|
||
"""
|
||
Combine multiple "L" mode masks.
|
||
|
||
mode : "union" — pixel is masked if ANY mask covers it (logical OR)
|
||
"intersection" — pixel is masked if ALL masks cover it (logical AND)
|
||
"""
|
||
arrays = [np.array(m, dtype=np.uint8) for m in masks]
|
||
if mode == "union":
|
||
result = arrays[0].copy()
|
||
for a in arrays[1:]:
|
||
result = np.maximum(result, a)
|
||
else:
|
||
result = arrays[0].copy()
|
||
for a in arrays[1:]:
|
||
result = np.minimum(result, a)
|
||
return Image.fromarray(result, mode="L")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# High-level entry point used by BaseDataset.random_mask
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# Probability weights for mask type selection
|
||
_TYPE_WEIGHTS = {
|
||
"cloud": 0.40,
|
||
"polygon": 0.35,
|
||
"shape": 0.25,
|
||
}
|
||
|
||
def random_mask(
|
||
width: int,
|
||
height: int,
|
||
cloud_octaves: int = 4,
|
||
cloud_persistence: float = 0.5,
|
||
cloud_threshold: float = 0.2,
|
||
polygon_n_points: int = 6,
|
||
polygon_irregularity: float = 0.4,
|
||
polygon_n_polygons: int = 1,
|
||
min_coverage: float = 0.1,
|
||
max_coverage: float = 0.6,
|
||
combine_cloud_and_shape: bool = True,
|
||
seed: Optional[int] = None,
|
||
) -> Image.Image:
|
||
"""
|
||
Generate a random inpainting mask using a randomly chosen (or combined) strategy.
|
||
|
||
With combine_cloud_and_shape=True (default) there is a 33% chance the result
|
||
is the union of a cloud mask and a polygon/shape mask, which produces more
|
||
varied and naturalistic training masks.
|
||
"""
|
||
rng = random.Random(seed)
|
||
np_seed = rng.randint(0, 2**31)
|
||
poly_seed = rng.randint(0, 2**31)
|
||
|
||
mask_type = rng.choices(
|
||
list(_TYPE_WEIGHTS.keys()), weights=list(_TYPE_WEIGHTS.values())
|
||
)[0]
|
||
|
||
if mask_type == "cloud":
|
||
threshold = rng.uniform(max(0.05, cloud_threshold - 0.1), min(0.4, cloud_threshold + 0.1))
|
||
m = cloud_mask(width, height, octaves=cloud_octaves, persistence=cloud_persistence,
|
||
threshold=threshold, seed=np_seed)
|
||
if combine_cloud_and_shape and rng.random() < 0.33:
|
||
shape_type = rng.choice(["polygon", "shape"])
|
||
if shape_type == "polygon":
|
||
m2 = polygon_mask(width, height, n_points=polygon_n_points,
|
||
min_coverage=min_coverage, max_coverage=max_coverage,
|
||
irregularity=polygon_irregularity, seed=poly_seed)
|
||
else:
|
||
m2 = shape_mask(width, height, min_coverage=min_coverage,
|
||
max_coverage=max_coverage, seed=poly_seed)
|
||
m = combine_masks(m, m2)
|
||
|
||
elif mask_type == "polygon":
|
||
n_poly = rng.randint(1, max(1, polygon_n_polygons * 2))
|
||
m = polygon_mask(width, height, n_points=polygon_n_points,
|
||
min_coverage=min_coverage, max_coverage=max_coverage,
|
||
irregularity=polygon_irregularity, n_polygons=n_poly,
|
||
seed=poly_seed)
|
||
|
||
else: # shape
|
||
m = shape_mask(width, height, min_coverage=min_coverage,
|
||
max_coverage=max_coverage, seed=poly_seed)
|
||
|
||
return m
|