remove-ai-watermarks/src/remove_ai_watermarks/region_eraser.py

"""Universal region eraser: remove anything inside user-given boxes via inpainting.

Position- and content-agnostic. You supply the rectangle(s); the eraser inpaints
whatever is inside, so it removes any visible logo / watermark / object regardless
of color, style, or location. Localization is the user's responsibility (pass the
box); restoration runs on CPU. This is the universal fallback for marks the
registered visible detectors do not cover.

Backends:
  - ``cv2`` (default): ``cv2.inpaint`` (Telea / Navier-Stokes). Instant, no extra
    dependencies, lower quality on large or textured regions.
  - ``migan`` (optional, extra ``migan``): MI-GAN via onnxruntime
    (``andraniksargsyan/migan``, MIT). CPU, ~28 MB model, ~0.19 s/call -- the
    droplet-friendly tier: near-big-LaMa quality on small marks. Model downloaded
    on first use. Like ``lama`` it crops a padded region around the mask before
    inference (at native resolution -- MI-GAN takes arbitrary dims), so peak RAM is
    bounded by the mark size (~0.6-0.9 GB) rather than scaling with the image, which
    is what lets a memory-tight host run it on a large upload.
  - ``lama`` (optional, extra ``lama``): big-LaMa via onnxruntime
    (``Carve/LaMa-ONNX``, Apache-2.0). CPU, resolution-robust, best quality on
    texture but ~200 MB model and ~4.7 GB peak RAM (too heavy for a small host).
    The model is downloaded on first use and cached by huggingface_hub; it is
    never bundled in this repo.

The per-backend RAM and wall-time figures above are reproduced by
``scripts/resource_ceilings.py`` (fresh subprocess per measurement, synthetic inputs).
Re-run it before changing any of them.
"""

# cv2/numpy boundary: cv2 ships no usable type info, so strict pyright cannot know
# its array element types. Relax the unknown-type rules for this file only; the
# public signatures are still annotated with NDArray[Any].
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false, reportConstantRedefinition=false, reportUnnecessaryComparison=false
from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal

import cv2
import numpy as np

if TYPE_CHECKING:
    from numpy.typing import NDArray

logger = logging.getLogger(__name__)

Backend = Literal["cv2", "lama", "migan"]

_LAMA_REPO = "Carve/LaMa-ONNX"
_LAMA_FILE = "lama_fp32.onnx"

_MIGAN_REPO = "andraniksargsyan/migan"
_MIGAN_FILE = "migan.onnx"

# Cached onnxruntime sessions, keyed by backend name (loading is expensive; reuse
# across calls).
_sessions: dict[str, object] = {}


@dataclass(frozen=True)
class _LearnedBackend:
    """One optional model-backed fill: its human label and the module-level names of
    its availability probe and erase function.

    NAMES, not function objects: the fallback tests monkeypatch these by attribute on
    the module, and a table holding bound references would not see the patch.
    """

    label: str
    available: str
    erase: str


# The learned tier, best-quality first. `resolve_backend`'s preference order and the
# CLI's choices are separate literals by design (see FILL_BACKENDS below), but the
# availability probe, install hint and dispatch all read this one table.
_LEARNED_BACKENDS: dict[str, _LearnedBackend] = {
    "lama": _LearnedBackend("LaMa", "lama_available", "erase_lama"),
    "migan": _LearnedBackend("MI-GAN", "migan_available", "erase_migan"),
}

# Every fill backend this module can execute, plus the caller-facing `auto`. Kept as a
# literal tuple (not derived through an import) so `watermark_registry` and the CLI can
# state their choices without importing this cv2-loading module at their import time --
# a test pins the two in sync.
FILL_BACKENDS: tuple[str, ...] = ("auto", "cv2", "migan", "lama")


def boxes_to_mask(
    shape: tuple[int, int],
    boxes: list[tuple[int, int, int, int]],
    dilate: int = 3,
) -> NDArray[Any]:
    """Build a uint8 mask (255 inside boxes) from ``(x, y, w, h)`` rectangles."""
    h, w = shape
    mask = np.zeros((h, w), np.uint8)
    for x, y, bw, bh in boxes:
        x0, y0 = max(0, x), max(0, y)
        x1, y1 = min(w, x + bw), min(h, y + bh)
        if x1 > x0 and y1 > y0:
            mask[y0:y1, x0:x1] = 255
    if dilate > 0 and mask.any():
        k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * dilate + 1, 2 * dilate + 1))
        mask = cv2.dilate(mask, k)
    return mask


def _padded_crop_box(
    mask: NDArray[Any], h: int, w: int, *, pad_frac: float, pad_min: int
) -> tuple[int, int, int, int] | None:
    """Bounding box of the set mask pixels, padded and clamped to the image.

    Returns ``(x0, y0, x1, y1)`` or ``None`` when the mask is empty. Both learned
    backends crop to this box so the ONNX working set is bounded by the mark size
    rather than the whole image; ``pad_frac``/``pad_min`` tune how much surrounding
    context the inpainter sees (LaMa then resizes the crop to its fixed square,
    MI-GAN feeds it at native resolution).
    """
    ys, xs = np.where(mask > 0)
    if len(xs) == 0:
        return None
    pad = max(pad_min, int(pad_frac * max(xs.max() - xs.min() + 1, ys.max() - ys.min() + 1)))
    x0, y0 = max(0, int(xs.min()) - pad), max(0, int(ys.min()) - pad)
    x1, y1 = min(w, int(xs.max()) + 1 + pad), min(h, int(ys.max()) + 1 + pad)
    return x0, y0, x1, y1


def erase_cv2(
    image_bgr: NDArray[Any],
    mask: NDArray[Any],
    *,
    method: Literal["telea", "ns"] = "telea",
    radius: int = 6,
) -> NDArray[Any]:
    """Inpaint ``mask`` with classical cv2 inpainting (CPU, no extra deps).

    Accepts 1-/3-channel BGR (passed straight to ``cv2.inpaint``) and 4-channel
    BGRA: ``cv2.inpaint`` rejects 4 channels, so the alpha plane is split off,
    the BGR is inpainted, and alpha is re-attached unchanged.
    """
    flag = cv2.INPAINT_TELEA if method == "telea" else cv2.INPAINT_NS
    if image_bgr.ndim == 3 and image_bgr.shape[2] == 4:
        bgr = cv2.inpaint(image_bgr[:, :, :3], mask, radius, flag)
        return np.dstack([bgr, image_bgr[:, :, 3]])
    return cv2.inpaint(image_bgr, mask, radius, flag)


def lama_available() -> bool:
    """True when the optional LaMa-ONNX backend can run (onnxruntime installed)."""
    from .optional_deps import module_available

    return module_available("onnxruntime")


def _get_session(name: str, repo_id: str, filename: str, label: str) -> object:
    """Load (once) an ONNX session, downloading the model on first use."""
    cached = _sessions.get(name)
    if cached is not None:
        return cached

    import onnxruntime as ort
    from huggingface_hub import hf_hub_download

    model_path = hf_hub_download(repo_id=repo_id, filename=filename)
    logger.info("Loading %s model: %s", label, model_path)
    session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
    _sessions[name] = session
    return session


def _get_lama_session() -> object:
    """The big-LaMa ONNX session (kept as a named seam the tests monkeypatch)."""
    return _get_session("lama", _LAMA_REPO, _LAMA_FILE, "LaMa-ONNX")


def erase_lama(image_bgr: NDArray[Any], mask: NDArray[Any]) -> NDArray[Any]:
    """Inpaint ``mask`` with big-LaMa via onnxruntime (CPU).

    LaMa runs at a fixed square input size. To preserve full-image resolution we
    crop a padded region around the mask, inpaint that crop at the model size,
    and paste only the masked pixels back -- so untouched areas stay pixel-exact.

    Like ``erase_cv2``, accepts 1-channel (grayscale) and 4-channel (BGRA) input:
    LaMa runs on 3-channel BGR, so grayscale is promoted to BGR (result demoted
    back) and a BGRA alpha plane is split off and re-attached unchanged. Without
    this the ``cv2.cvtColor(..., BGR2RGB)`` below would crash on grayscale and
    silently drop alpha on BGRA.
    """
    if image_bgr.ndim == 2:
        bgr = erase_lama(cv2.cvtColor(image_bgr, cv2.COLOR_GRAY2BGR), mask)
        return cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    if image_bgr.ndim == 3 and image_bgr.shape[2] == 4:
        bgr = erase_lama(np.ascontiguousarray(image_bgr[:, :, :3]), mask)
        return np.dstack([bgr, image_bgr[:, :, 3]])
    session = _get_lama_session()
    inp = session.get_inputs()  # type: ignore[attr-defined]
    img_name = inp[0].name
    mask_name = inp[1].name
    # Model declares a fixed square spatial size (e.g. 512); fall back to 512.
    dims = inp[0].shape
    size = next((d for d in reversed(dims) if isinstance(d, int) and d > 1), 512)

    h, w = image_bgr.shape[:2]
    box = _padded_crop_box(mask, h, w, pad_frac=0.4, pad_min=16)
    if box is None:
        return image_bgr.copy()
    cx0, cy0, cx1, cy1 = box
    crop = image_bgr[cy0:cy1, cx0:cx1]
    crop_mask = mask[cy0:cy1, cx0:cx1]
    ch, cw = crop.shape[:2]

    # Resize crop + mask to the model size, normalize to [0,1] RGB CHW.
    crop_rs = cv2.resize(crop, (size, size), interpolation=cv2.INTER_AREA)
    mask_rs = cv2.resize(crop_mask, (size, size), interpolation=cv2.INTER_NEAREST)
    img_in = cv2.cvtColor(crop_rs, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    img_in = np.transpose(img_in, (2, 0, 1))[None]  # (1,3,size,size)
    mask_in = (mask_rs > 127).astype(np.float32)[None, None]  # (1,1,size,size), 1=hole

    out = session.run(None, {img_name: img_in, mask_name: mask_in})[0]  # type: ignore[attr-defined]
    out = np.asarray(out)[0]  # (3,size,size)
    out = np.transpose(out, (1, 2, 0))
    if float(out.max()) <= 1.5:  # model emits [0,1]; otherwise already [0,255]
        out = out * 255.0
    out = np.clip(out, 0, 255).astype(np.uint8)
    out_bgr = cv2.cvtColor(out, cv2.COLOR_RGB2BGR)

    # Resize back to crop size and paste only the masked pixels.
    out_crop = cv2.resize(out_bgr, (cw, ch), interpolation=cv2.INTER_LINEAR)
    result = image_bgr.copy()
    region = result[cy0:cy1, cx0:cx1]
    paste = crop_mask > 127
    region[paste] = out_crop[paste]
    result[cy0:cy1, cx0:cx1] = region
    return result


def migan_available() -> bool:
    """True when the optional MI-GAN backend can run (onnxruntime installed).

    Deliberately a separate function from :func:`lama_available` even though both
    currently reduce to the same onnxruntime probe: they are independent capability
    questions and `auto` resolves them separately (see
    ``watermark_registry.preferred_inpaint_backend``).
    """
    from .optional_deps import module_available

    return module_available("onnxruntime")


def _get_migan_session() -> object:
    """The MI-GAN ONNX session (kept as a named seam the tests monkeypatch)."""
    return _get_session("migan", _MIGAN_REPO, _MIGAN_FILE, "MI-GAN ONNX")


def erase_migan(image_bgr: NDArray[Any], mask: NDArray[Any]) -> NDArray[Any]:
    """Inpaint ``mask`` (255 = erase) with MI-GAN via onnxruntime (CPU).

    Like ``erase_lama``, we crop a padded region around the mask, feed only that crop
    to the ONNX model, and paste only the masked pixels back -- so untouched areas stay
    pixel-exact and the ONNX working set is bounded by the MARK size, not the whole
    image. MI-GAN accepts arbitrary spatial dims, so (unlike LaMa's fixed 512 square)
    the crop is fed at NATIVE resolution -- no resize, so the mark is seen at full
    scale. Feeding the whole frame instead made peak RAM scale with the image
    (~0.6 GB at 4 MP up to ~2.4 GB at 25 MP, measured 2026-07); cropping holds the
    inpaint working set roughly constant, so a memory-tight host (e.g. a 1-2 GB web
    worker) can run MI-GAN on a 25 MP upload. Cropping does not degrade the fill:
    a small mark only needs local context, and on real marks the cropped fill is on
    par with -- sometimes cleaner than -- the full-frame fill (a tighter view gives
    the GAN less room to hallucinate large background structure).

    Mask polarity: the shipped ``andraniksargsyan/migan`` ONNX expects 0 = hole
    (inpaint) / 255 = known (keep) -- the INVERSE of this package's 255-erase
    convention -- so the mask is inverted before feeding the model. Feeding 255=hole
    regenerates the whole frame into stripes.

    Accepts 1-channel (grayscale) and 4-channel (BGRA) input.
    """
    if image_bgr.ndim == 2:
        bgr = erase_migan(cv2.cvtColor(image_bgr, cv2.COLOR_GRAY2BGR), mask)
        return cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    if image_bgr.ndim == 3 and image_bgr.shape[2] == 4:
        bgr = erase_migan(np.ascontiguousarray(image_bgr[:, :, :3]), mask)
        return np.dstack([bgr, image_bgr[:, :, 3]])

    h, w = image_bgr.shape[:2]
    # 2x the mark size (min 256 px) gives ample local context while a small corner
    # mask keeps the crop small regardless of the full image resolution; the crop is
    # clamped to the image, so a mark already spanning the frame degrades to the
    # whole image (the old behavior) rather than erroring.
    box = _padded_crop_box(mask, h, w, pad_frac=2.0, pad_min=256)
    if box is None:
        return image_bgr.copy()
    cx0, cy0, cx1, cy1 = box
    crop = np.ascontiguousarray(image_bgr[cy0:cy1, cx0:cx1])
    crop_mask = mask[cy0:cy1, cx0:cx1]
    ch, cw = crop.shape[:2]

    session = _get_migan_session()
    inp = session.get_inputs()  # type: ignore[attr-defined]
    img_name, mask_name = inp[0].name, inp[1].name

    rgb = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
    img_in = np.transpose(rgb, (2, 0, 1))[None].astype(np.uint8)  # (1,3,ch,cw)
    # invert to MI-GAN polarity: 255 where KNOWN (keep), 0 where hole (erase)
    known = (crop_mask <= 127).astype(np.uint8) * 255
    mask_in = known[None, None]  # (1,1,ch,cw)

    out = session.run(None, {img_name: img_in, mask_name: mask_in})[0]  # type: ignore[attr-defined]
    res = np.transpose(np.asarray(out)[0], (1, 2, 0)).astype(np.uint8)  # (ch',cw',3) RGB
    if res.shape[:2] != (ch, cw):
        res = cv2.resize(res, (cw, ch), interpolation=cv2.INTER_LINEAR)
    out_bgr = cv2.cvtColor(res, cv2.COLOR_RGB2BGR)

    result = image_bgr.copy()
    region = result[cy0:cy1, cx0:cx1]
    hole = crop_mask > 127
    region[hole] = out_bgr[hole]
    result[cy0:cy1, cx0:cx1] = region
    return result


def erase(
    image_bgr: NDArray[Any],
    *,
    boxes: list[tuple[int, int, int, int]] | None = None,
    mask: NDArray[Any] | None = None,
    backend: Backend = "cv2",
    dilate: int = 3,
    cv2_method: Literal["telea", "ns"] = "telea",
    cv2_radius: int = 6,
) -> NDArray[Any]:
    """Erase the given boxes (or mask) via the chosen inpainting backend.

    Provide either ``boxes`` (list of ``(x, y, w, h)``) or a precomputed ``mask``
    (uint8, 255 = erase). Returns an unmodified copy when nothing is selected.
    """
    if image_bgr is None or image_bgr.size == 0:
        return image_bgr
    if mask is None:
        if not boxes:
            return image_bgr.copy()
        mask = boxes_to_mask(image_bgr.shape[:2], boxes, dilate=dilate)
    if not mask.any():
        return image_bgr.copy()

    learned = _LEARNED_BACKENDS.get(backend)
    if learned is not None:
        # Probe and erase by MODULE NAME, not by the captured function object: the
        # availability probes and the erase functions are monkeypatched by name in the
        # fallback tests, and a table of bound references would not see those patches.
        if not globals()[learned.available]():
            raise RuntimeError(
                f"{learned.label} backend requires onnxruntime. "
                f"Install the extra: pip install 'remove-ai-watermarks[{backend}]'"
            )
        return globals()[learned.erase](image_bgr, mask)
    # cv2 and anything unrecognized (including "auto", which a library caller may pass
    # straight through) degrade to the classical fill rather than raising.
    return erase_cv2(image_bgr, mask, method=cv2_method, radius=cv2_radius)