58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
"""Build an inpaint mask (255 = regenerate) from detections.
|
|
|
|
Used by the diffusion restorer. Unlike DeepMosaics — which locates the mosaic itself —
|
|
a diffusion-inpaint backend needs an explicit mask of the region to redraw. We rasterise
|
|
each detection's polygon (or its bbox when there's no polygon) onto a single-channel
|
|
uint8 mask, optionally growing (dilate) and feathering (blur) the edges so the inpaint
|
|
blends into the surrounding pixels.
|
|
|
|
Pure NumPy/OpenCV — no torch, no Qt.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from ..detection.types import Detection
|
|
|
|
|
|
def detections_to_mask(
|
|
detections: Sequence[Detection],
|
|
shape: tuple[int, ...],
|
|
*,
|
|
dilate: int = 0,
|
|
blur: int = 0,
|
|
) -> np.ndarray:
|
|
"""Rasterise ``detections`` onto a single-channel uint8 mask (255 = regenerate).
|
|
|
|
``shape`` is the image shape (``(h, w)`` or ``(h, w, c)``). ``dilate`` grows the mask
|
|
by that many pixels (ellipse kernel) so the inpaint covers the censored edge; ``blur``
|
|
feathers the edge with a Gaussian so the boundary blends. Both are no-ops at 0.
|
|
"""
|
|
h, w = int(shape[0]), int(shape[1])
|
|
mask = np.zeros((h, w), dtype=np.uint8)
|
|
for d in detections:
|
|
if len(d.polygon) >= 3:
|
|
poly = np.array(
|
|
[[int(round(x)), int(round(y))] for x, y in d.polygon], dtype=np.int32
|
|
)
|
|
cv2.fillPoly(mask, [poly], 255)
|
|
else:
|
|
x, y, bw, bh = (int(round(v)) for v in d.bbox)
|
|
cv2.rectangle(mask, (x, y), (x + bw, y + bh), 255, thickness=-1)
|
|
if dilate > 0:
|
|
k = 2 * int(dilate) + 1
|
|
mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k)))
|
|
if blur > 0:
|
|
k = 2 * int(blur) + 1
|
|
mask = cv2.GaussianBlur(mask, (k, k), 0)
|
|
return mask
|
|
|
|
|
|
def mask_is_empty(mask: np.ndarray) -> bool:
|
|
"""True if nothing is masked (so there's nothing to inpaint)."""
|
|
return not bool(np.any(mask))
|