Introduce diffusion-inpaint restoration engine in HVideoTool: added support for a new restoration method that regenerates masked regions via an external SwarmUI server, requiring YOLO detections for mask creation. Updated configuration management to include diffusion parameters, enhanced the UI for engine selection, and improved documentation in README and CLAUDE.md to guide users on the new functionality.
This commit is contained in:
@@ -0,0 +1,57 @@
|
||||
"""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))
|
||||
Reference in New Issue
Block a user