CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-08-07 06:28:36 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(visible): capture-less AI生成 pill (#54), inpaint fallback, MI-GAN backend (#56)

Browse Source 
- Add the Jimeng-basic top-left "AI生成" pill as a CAPTURE-LESS mark
  (pill_engine.py): synthetic-silhouette edge-NCC detect + inpaint-only removal.
  Gated in remove_auto_marks: kept only when Jimeng is confirmed (TC260 metadata
  OR the bottom-right "★ 即梦AI" wordmark fired -- the wordmark keeps recall on
  metadata-STRIPPED uploads) AND Doubao did not fire.
- Add an inpaint-fallback removal path + MI-GAN ONNX backend (migan extra, MIT,
  ~28 MB / ~1 GB peak -- droplet-friendly) alongside big-LaMa. New
  --method auto|reverse-alpha|inpaint (shared across visible/all/batch) and
  erase --backend migan; footprint_mask on each engine.
- auto is deterministic: reverse-alpha for capture marks (recovers exact pixels,
  lighter -- measured cleaner than MI-GAN on structured backgrounds) and inpaint
  only for the capture-less pill.
- --mark auto now removes EVERY detected mark in one pass (remove_auto_marks),
  so a Jimeng-basic image's top-left pill AND bottom-right wordmark both clear.
- Bump 0.12.1 -> 0.13.0.

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-07-06 20:38:23 +03:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 0f54c6b54d
commit 0e5a4cbc54
18 changed files with 994 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/remove_ai_watermarks/__init__.py
+1 -1
View File
@@ -11,4 +11,4 @@ _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
_warnings.filterwarnings("ignore", message=r".*ImageProcessorFast.*")
__version__ = "0.12.1"
__version__ = "0.13.0"
src/remove_ai_watermarks/_text_mark_engine.py
+41
View File
@@ -504,3 +504,44 @@ class TextMarkEngine:
rm = cv2.dilate(rm, kernel)
best_out = cv2.inpaint(best_out, rm, c.residual_inpaint_radius, cv2.INPAINT_NS)
return best_out
# ── Inpaint footprint (for the inpaint-fallback removal path) ────────
def footprint_mask(
self, image: NDArray[Any], *, force: bool = False, dilate: int | None = None
) -> NDArray[Any] | None:
"""Full-frame uint8 mask (255 = mark) of the mark footprint, for the
inpaint-fallback removal path (LaMa / cv2), or None if no placement fits.
``force`` is accepted for a uniform engine signature (the caller passes it to
every engine) but ignored here -- the text-mark footprint is always the
geometry-placed captured silhouette, present with or without a detection.
Uses the NCC-ALIGNED captured silhouette, NOT the per-image
:meth:`extract_mask` signature: the signature under-segments the glyphs, so
inpainting it leaves a residual ghost (corpus-validated 2026-07 -- Doubao
left a "三包" remnant). The mask is dilated to absorb alpha-alignment slop
(a scale/position mismatch at low detect confidence otherwise leaves a thin
residual ring); ``dilate`` defaults to a mark-relative margin.
The caller gates on detection -- this returns the geometric footprint
regardless, so a clean corner would be masked too.
"""
image = image_io.to_bgr(image)
h, w = image.shape[:2]
if h < 32 or w < 64:
return None
placed = self._aligned_alpha_map(image) or self._fixed_alpha_map(image)
if placed is None:
return None
block, (ax, ay, gw, gh) = placed
sil = (block > self.config.residual_alpha_floor).astype(np.uint8) * 255
if int((sil > 0).sum()) == 0:
return None
mask = np.zeros((h, w), np.uint8)
ch, cw = min(gh, h - ay), min(gw, w - ax)
mask[ay : ay + ch, ax : ax + cw] = sil[:ch, :cw]
d = dilate if dilate is not None else max(9, int(0.05 * gw))
if d > 0:
mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * d + 1, 2 * d + 1)))
return mask
src/remove_ai_watermarks/assets/jimeng_pill.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 3.1 KiB

src/remove_ai_watermarks/cli.py
+108 -51
View File
@@ -293,6 +293,18 @@ _force_option = click.option(
)
_visible_method_option = click.option(
"--method",
"removal_method",
type=click.Choice(["auto", "reverse-alpha", "inpaint"]),
default="auto",
help="Visible-mark removal method. auto: reverse-alpha for capture marks (exact "
"pixels, lighter), inpaint for the capture-less pill. reverse-alpha recovers "
"pixels from a captured alpha map; inpaint erases the footprint (MI-GAN with the "
"'migan' extra, else cv2).",
)
def _resolve_auto_polish(auto: bool, adaptive_polish: bool) -> bool:
"""Warn on the retired ``--auto`` flag, returning ``adaptive_polish`` unchanged.
@@ -327,9 +339,23 @@ def _warn_if_esrgan_unavailable(upscaler: str) -> None:
console.print(" Note: --upscaler esrgan needs the 'esrgan' extra; falling back to Lanczos.")
def _aigc_metadata_present(path: Path) -> bool:
"""True when the file carries a China-AIGC (TC260) metadata label. Used to gate
the weak-detector 'AI生成' pill: metadata confirms Jimeng-class provenance. NB
this is only ONE of two confirmations -- ``remove_auto_marks`` also accepts the
bottom-right wordmark, so a metadata-STRIPPED upload can still be handled."""
with contextlib.suppress(Exception):
from remove_ai_watermarks import metadata
return bool(metadata.aigc_label(path))
return False
def _remove_visible_auto(
image: NDArray[Any],
*,
source_path: Path | None = None,
removal_method: str = "auto",
inpaint: bool = True,
inpaint_method: str = "ns",
inpaint_strength: float = 0.85,
@@ -340,19 +366,26 @@ def _remove_visible_auto(
standalone ``visible`` command uses, so EVERY registered mark is handled (the
Gemini sparkle AND the Doubao/Jimeng/Samsung text marks), not just the sparkle.
Returns ``(result, label-or-None)``; when no ``in_auto`` mark fires the image is
returned unchanged with ``None``. ``inpaint*`` tune the Gemini edge-residual
cleanup only (the text engines ignore them).
returned unchanged with ``None``. ``removal_method`` selects reverse-alpha vs the
inpaint fallback (see ``KnownMark.remove``); ``inpaint*`` tune the Gemini
edge-residual cleanup only (the text engines ignore them).
"""
from remove_ai_watermarks import watermark_registry
best = watermark_registry.best_auto_mark(image)
if best is None:
return image, None
rmethod: watermark_registry.RemovalMethod = removal_method # type: ignore[assignment]
method: Literal["telea", "ns"] = "ns" if inpaint_method == "ns" else "telea"
result, _ = watermark_registry.get_mark(best.key).remove(
image, inpaint_method=method, inpaint=inpaint, inpaint_strength=inpaint_strength, force=False
pill_md = _aigc_metadata_present(source_path) if source_path is not None else False
result, removed = watermark_registry.remove_auto_marks(
image,
pill_metadata=pill_md,
method=rmethod,
inpaint_method=method,
inpaint=inpaint,
inpaint_strength=inpaint_strength,
)
return result, best.label
if not removed:
return image, None
return result, ", ".join(removed)
# Exit code for the standalone ``visible`` command when no visible mark was
@@ -535,8 +568,9 @@ def main(ctx: click.Context, verbose: bool) -> None:
type=click.Choice(["auto", *watermark_registry.mark_keys()]),
default="auto",
help="Which known visible mark to target (auto picks the strongest detected). "
"All marks are removed by exact reverse-alpha against a captured alpha map.",
"Removal method is chosen by --method (default auto).",
)
@_visible_method_option
@click.option("--strip-metadata/--keep-metadata", default=True, help="Strip AI metadata from output.")
@click.pass_context
def cmd_visible(
@@ -548,14 +582,17 @@ def cmd_visible(
inpaint_strength: float,
detect: bool,
mark: str,
removal_method: str,
strip_metadata: bool,
) -> None:
"""Remove a known visible AI watermark from an image.
Finds a known mark in its usual place (Gemini sparkle / Doubao text) via the
watermark registry and removes it by exact reverse-alpha against a captured
alpha map -- recovering the true pixels, not an inpaint guess. ``--mark auto``
picks the strongest detected mark. For arbitrary logos/objects, use ``erase``.
watermark registry and removes it. Default ``--method auto`` recovers the true
pixels by exact reverse-alpha for the capture marks, and inpaints only the
capture-less "AI生成" pill (MI-GAN with the ``migan`` extra, else cv2).
``--mark auto`` picks the strongest detected mark. For arbitrary logos/objects,
use ``erase``.
"""
from remove_ai_watermarks import watermark_registry as registry
@@ -574,41 +611,52 @@ def cmd_visible(
h, w = image.shape[:2]
console.print(f" Input: {source.name} ({w}x{h})")
# Resolve the target mark from the known-watermark registry. ``auto`` scans
# every in-auto mark in its usual place and picks the strongest; an explicit
# ``--mark <key>`` targets that one (the user asserts its presence).
if mark == "auto":
best = registry.best_auto_mark(image)
if best is None:
console.print(" No known visible mark detected (gemini / doubao / jimeng / samsung).")
if detect:
_no_visible_mark_exit(source)
target = "gemini" # forced (no-detect): fall back to the default mark
else:
target = best.key
console.print(f" Mark auto: {best.label} ({best.location}, conf {best.confidence:.2f})")
else:
target = mark
chosen = registry.get_mark(target)
det = chosen.detect(image)
if detect and not det.detected:
console.print(f" {chosen.label} not detected (conf {det.confidence:.2f}). Use --no-detect to force.")
_no_visible_mark_exit(source)
if det.detected:
console.print(f" {chosen.label} detected ({chosen.location}, conf {det.confidence:.2f})")
method: Literal["telea", "ns"] = "ns" if inpaint_method == "ns" else "telea"
t0 = time.monotonic()
with console.status(f"Removing {chosen.label}... ({chosen.recovery})"):
result, _ = chosen.remove(
image,
inpaint_method=method,
inpaint=inpaint,
inpaint_strength=inpaint_strength,
force=not detect,
)
elapsed = time.monotonic() - t0
# ``auto`` removes EVERY detected in_auto mark in one pass (a Jimeng-basic image
# carries the top-left pill AND the bottom-right wordmark); an explicit
# ``--mark <key>`` targets that one (the user asserts its presence).
if mark == "auto" and detect:
t0 = time.monotonic()
with console.status("Detecting & removing visible marks..."):
result, removed = registry.remove_auto_marks(
image,
pill_metadata=_aigc_metadata_present(source),
method=removal_method, # type: ignore[arg-type]
inpaint_method=method,
inpaint=inpaint,
inpaint_strength=inpaint_strength,
)
elapsed = time.monotonic() - t0
if not removed:
console.print(" No known visible mark detected (gemini / doubao / jimeng / jimeng-pill / samsung).")
_no_visible_mark_exit(source)
console.print(f" Removed: {', '.join(removed)}")
else:
target = "gemini" if mark == "auto" else mark # --no-detect auto: gemini fallback
chosen = registry.get_mark(target)
det = chosen.detect(image)
if detect and not det.detected:
console.print(f" {chosen.label} not detected (conf {det.confidence:.2f}). Use --no-detect to force.")
_no_visible_mark_exit(source)
if det.detected:
console.print(f" {chosen.label} detected ({chosen.location}, conf {det.confidence:.2f})")
resolved = registry.resolve_removal_method(removal_method, chosen.has_capture) # type: ignore[arg-type]
if resolved == "inpaint" and not registry.inpaint_model_available():
console.print(
" Note: --method inpaint using cv2 (install the 'migan' extra for a lightweight ONNX model)."
)
t0 = time.monotonic()
with console.status(f"Removing {chosen.label}... ({resolved})"):
result, _ = chosen.remove(
image,
method=removal_method, # type: ignore[arg-type]
inpaint_method=method,
inpaint=inpaint,
inpaint_strength=inpaint_strength,
force=not detect,
)
elapsed = time.monotonic() - t0
# Save (rejoins the original alpha plane unchanged)
output.parent.mkdir(parents=True, exist_ok=True)
@@ -653,9 +701,10 @@ def _parse_region(spec: str) -> tuple[int, int, int, int]:
)
@click.option(
"--backend",
type=click.Choice(["cv2", "lama"]),
type=click.Choice(["cv2", "migan", "lama"]),
default="cv2",
help="Inpaint backend. cv2: instant, no deps. lama: onnxruntime big-LaMa, better quality (extra 'lama').",
help="Inpaint backend. cv2: instant, no deps. migan: light ONNX MI-GAN, ~1 GB RAM, "
"near-LaMa quality (extra 'migan'). lama: big-LaMa, best quality but ~4.7 GB RAM (extra 'lama').",
)
@click.option("--inpaint-method", type=click.Choice(["telea", "ns"]), default="telea", help="cv2 inpaint method.")
@click.option("--dilate", type=int, default=3, help="Grow the box by this many px before inpainting.")
@@ -666,7 +715,7 @@ def cmd_erase(
source: Path,
regions: tuple[str, ...],
output: Path | None,
backend: Literal["cv2", "lama"],
backend: Literal["cv2", "migan", "lama"],
inpaint_method: str,
dilate: int,
strip_metadata: bool,
@@ -999,6 +1048,7 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
@click.option(
"--inpaint-method", type=click.Choice(["ns", "telea", "gaussian"]), default="ns", help="Inpainting method."
)
@_visible_method_option
@_strength_option
@click.option("--steps", type=int, default=50, help="Number of denoising steps for invisible removal.")
@_pipeline_option
@@ -1036,6 +1086,7 @@ def cmd_all(
output: Path | None,
inpaint: bool,
inpaint_method: Literal["ns", "telea", "gaussian"],
removal_method: str,
strength: float | None,
steps: int,
pipeline: str,
@@ -1105,7 +1156,9 @@ def cmd_all(
console.print(f" Input: {source.name} ({w}x{h})")
with console.status("Removing visible watermark..."):
result, removed_label = _remove_visible_auto(image, inpaint=inpaint, inpaint_method=inpaint_method)
result, removed_label = _remove_visible_auto(
image, source_path=source, removal_method=removal_method, inpaint=inpaint, inpaint_method=inpaint_method
)
if removed_label is not None:
console.print(f" Visible watermark removed ({removed_label})")
else:
@@ -1244,6 +1297,7 @@ def _process_batch_image(
seed: int | None,
hf_token: str | None,
humanize: float,
removal_method: str = "auto",
unsharp: float = 0.0,
max_resolution: int = 0,
min_resolution: int = 1024,
@@ -1276,7 +1330,7 @@ def _process_batch_image(
if image is None:
raise ValueError("Failed to read image")
result, _ = _remove_visible_auto(image, inpaint=inpaint)
result, _ = _remove_visible_auto(image, source_path=img_path, removal_method=removal_method, inpaint=inpaint)
_write_bgr_with_alpha(out_path, result, alpha)
saved_alpha = alpha
@@ -1361,6 +1415,7 @@ def _process_batch_image(
@_strength_option
@click.option("--steps", type=int, default=50, help="Number of denoising steps (invisible mode).")
@click.option("--inpaint/--no-inpaint", default=True, help="Apply inpainting (visible mode).")
@_visible_method_option
@click.option(
"--humanize", type=float, default=0.0, help="Analog Humanizer film grain intensity (0 = off, typical: 2.0-6.0)."
)
@@ -1402,6 +1457,7 @@ def cmd_batch(
seed: int | None,
hf_token: str | None,
inpaint: bool,
removal_method: str,
humanize: float,
unsharp: float,
max_resolution: int,
@@ -1468,6 +1524,7 @@ def cmd_batch(
seed=seed,
hf_token=hf_token,
humanize=humanize,
removal_method=removal_method,
unsharp=unsharp,
max_resolution=max_resolution,
min_resolution=min_resolution,
src/remove_ai_watermarks/gemini_engine.py
+34
View File
@@ -603,6 +603,40 @@ class GeminiEngine:
self._reverse_alpha_blend(result, alpha_map, pos)
return self._verify_and_repair(result, alpha_map, pos, size)
def footprint_mask(self, image: NDArray[Any], *, force: bool = False, dilate: int = 13) -> NDArray[Any] | None:
"""Full-frame uint8 mask (255 = sparkle) of the sparkle footprint, for the
inpaint-fallback removal path (LaMa / cv2), or None.
The footprint is the interpolated captured alpha at the detected scale --
the same region reverse-alpha operates on. When ``force`` and nothing is
detected, falls back to the default sparkle slot for the image size (the
``--no-detect`` path). The caller gates on the trust-confidence detection.
"""
image = image_io.to_bgr(image)
h, w = image.shape[:2]
det = self.detect_watermark(image)
if det.detected:
x, y, scale = det.region[0], det.region[1], det.region[2]
elif force:
cfg = get_watermark_config(w, h)
x, y = cfg.get_position(w, h)
scale = cfg.logo_size
else:
return None
alpha = self.get_interpolated_alpha(scale)
fp = self._footprint_indices(alpha, (x, y), image.shape)
if fp is None:
return None
aroi, (y1, y2, x1, x2) = fp
sil = (aroi > 0.10).astype(np.uint8) * 255
if int((sil > 0).sum()) == 0:
return None
mask = np.zeros((h, w), np.uint8)
mask[y1:y2, x1:x2] = sil
if dilate > 0:
mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * dilate + 1, 2 * dilate + 1)))
return mask
def remove_watermark_custom(
self,
image: NDArray[Any],
src/remove_ai_watermarks/pill_engine.py
+126
View File
@@ -0,0 +1,126 @@
"""Jimeng-basic 'AI生成' pill: a CAPTURE-LESS visible mark (issue #54).
The Jimeng free-tier TC260 label is a rounded pill with 'AI生成' in the TOP-LEFT
corner -- distinct from the reverse-alpha ``jimeng`` "★ 即梦AI" mark (bottom-right).
No flat capture / alpha map exists for it, so it is removed by INPAINT, not
reverse-alpha:
* Detect: edge-NCC of a font-rendered SILHOUETTE (``assets/jimeng_pill.png``,
synthetic, data-safe -- see ``scripts/render_pill_silhouette.py``) against the
top-left ROI, at the pill's known width fraction. Corpus-calibrated threshold
(61 real positives + jimeng negatives): ``_DETECT_THRESHOLD`` 0.22.
* Remove: place the pill footprint at the matched location and inpaint it
(MI-GAN / cv2 via the registry). Quality comes from the inpaint backend, so the
silhouette need not be pixel-accurate -- which is why a synthetic render is
sufficient and no corpus-derived asset is committed.
Geometry measured on 51 real examples (8 resolutions, all 3:4): width ~0.161*W,
height ~0.091*W, top-left, margins ~0.02-0.05.
"""
from __future__ import annotations
from pathlib import Path
from typing import TYPE_CHECKING, Any, NamedTuple
import cv2
import numpy as np
from remove_ai_watermarks import image_io
if TYPE_CHECKING:
from numpy.typing import NDArray
# 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, reportOptionalSubscript=false, reportAttributeAccessIssue=false, reportUnnecessaryComparison=false
_ASSET = Path(__file__).parent / "assets" / "jimeng_pill.png"
# Geometry (fractions of image WIDTH unless noted); top-left corner.
_WIDTH_FRAC = 0.161
_ROI_W_FRAC = 0.34 # search window width (of W)
_ROI_H_FRAC = 0.14 # search window height (of H)
_DETECT_THRESHOLD = 0.22 # edge-NCC gate, corpus-calibrated
# Inpaint mask GEOMETRY (fractions of W unless noted): a generous fixed top-left box
# covering the pill (measured ~0.167*W wide, ~0.09*W tall, margin ~0.02-0.05) plus
# margin. The mask uses stable geometry, NOT the NCC match position -- the synthetic
# silhouette localizes only approximately, and the corner is negative space, so
# over-covering is harmless while a match-positioned box leaves outline residue.
_MASK_X0, _MASK_Y0 = 0.012, 0.006 # x0 of W, y0 of H
_MASK_W, _MASK_H = 0.205, 0.115 # width of W, height of W
_silhouette: NDArray[Any] | None = None
class PillDetection(NamedTuple):
detected: bool
confidence: float
region: tuple[int, int, int, int] # x, y, w, h of the matched pill
def _load_silhouette() -> NDArray[Any] | None:
global _silhouette
if _silhouette is None:
if not _ASSET.exists():
return None
_silhouette = image_io.imread(str(_ASSET), cv2.IMREAD_GRAYSCALE)
return _silhouette
def _grad(gray: NDArray[Any]) -> NDArray[Any]:
gx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3)
gy = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3)
return cv2.normalize(cv2.magnitude(gx, gy), None, 0, 255, cv2.NORM_MINMAX)
class PillEngine:
"""Detect + inpaint-mask the top-left 'AI生成' pill (edge-NCC, no reverse-alpha)."""
def _match(self, image: NDArray[Any]) -> tuple[float, tuple[int, int, int, int]] | None:
sil = _load_silhouette()
if sil is None or image is None or image.size == 0:
return None
h, w = image.shape[:2]
if h < 64 or w < 64:
return None
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image
rh, rw = int(h * _ROI_H_FRAC), int(w * _ROI_W_FRAC)
roi = gray[0:rh, 0:rw]
tw = max(24, int(_WIDTH_FRAC * w))
th = max(12, int(tw * sil.shape[0] / sil.shape[1]))
if th >= rh or tw >= rw:
return None
tmpl = cv2.resize(sil, (tw, th))
res = cv2.matchTemplate(_grad(roi.astype(np.float32)), _grad(tmpl.astype(np.float32)), cv2.TM_CCOEFF_NORMED)
_, score, _, loc = cv2.minMaxLoc(res)
return float(score), (int(loc[0]), int(loc[1]), tw, th)
def detect(self, image: NDArray[Any]) -> PillDetection:
m = self._match(image)
if m is None:
return PillDetection(False, 0.0, (0, 0, 0, 0))
score, box = m
return PillDetection(score >= _DETECT_THRESHOLD, score, box)
def footprint_mask(self, image: NDArray[Any], *, force: bool = False) -> NDArray[Any] | None:
"""Full-frame uint8 mask (255 = pill) over the pill's known top-left region.
Uses stable GEOMETRY (a generous fixed box), not the NCC match position: the
synthetic silhouette localizes only approximately, so a match-positioned mask
leaves outline residue, while the top-left corner is negative space, so a
generous geometric box removes the pill cleanly and harmlessly. The caller
gates on :meth:`detect`, so a clean corner is never masked. ``force`` is
accepted for a uniform engine signature but ignored (the geometry box is
fixed regardless)."""
if image is None or image.size == 0:
return None
h, w = image.shape[:2]
x0, y0 = int(_MASK_X0 * w), int(_MASK_Y0 * h)
x1, y1 = min(w, x0 + int(_MASK_W * w)), min(h, y0 + int(_MASK_H * w))
if x1 <= x0 or y1 <= y0:
return None
mask = np.zeros((h, w), np.uint8)
mask[y0:y1, x0:x1] = 255
return mask
src/remove_ai_watermarks/region_eraser.py
+86 -5
View File
@@ -9,10 +9,15 @@ deterministic per-generator engines (Gemini sparkle, Doubao) 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, ~700-950 MB peak RAM,
~0.19 s/call -- the droplet-friendly tier: near-big-LaMa quality on small
marks at ~5x less RAM and ~8x faster. Model downloaded on first use.
- ``lama`` (optional, extra ``lama``): big-LaMa via onnxruntime
(``Carve/LaMa-ONNX``, Apache-2.0). CPU, resolution-robust, much better on
texture. The model (~200 MB) is downloaded on first use and cached by
huggingface_hub; it is never bundled in this repo.
(``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.
"""
# cv2/numpy boundary: cv2 ships no usable type info, so strict pyright cannot know
@@ -32,13 +37,17 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
Backend = Literal["cv2", "lama"]
Backend = Literal["cv2", "lama", "migan"]
_LAMA_REPO = "Carve/LaMa-ONNX"
_LAMA_FILE = "lama_fp32.onnx"
# Cached onnxruntime session (loading is expensive; reuse across calls).
_MIGAN_REPO = "andraniksargsyan/migan"
_MIGAN_FILE = "migan.onnx"
# Cached onnxruntime sessions (loading is expensive; reuse across calls).
_lama_session: object | None = None
_migan_session: object | None = None
def boxes_to_mask(
@@ -167,6 +176,72 @@ def erase_lama(image_bgr: NDArray[Any], mask: NDArray[Any]) -> NDArray[Any]:
return result
def migan_available() -> bool:
"""True when the optional MI-GAN backend can run (onnxruntime installed)."""
from .optional_deps import module_available
return module_available("onnxruntime")
def _get_migan_session() -> object:
"""Load (once) the MI-GAN ONNX session, downloading the model on first use."""
global _migan_session
if _migan_session is not None:
return _migan_session
import onnxruntime as ort
from huggingface_hub import hf_hub_download
model_path = hf_hub_download(repo_id=_MIGAN_REPO, filename=_MIGAN_FILE)
logger.info("Loading MI-GAN ONNX model: %s", model_path)
_migan_session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
return _migan_session
def erase_migan(image_bgr: NDArray[Any], mask: NDArray[Any]) -> NDArray[Any]:
"""Inpaint ``mask`` (255 = erase) with MI-GAN via onnxruntime (CPU).
The MI-GAN ONNX pipeline crops around the mask bbox internally and re-composites,
so the full image is fed at native resolution. Only the masked pixels are pasted
back, so untouched areas stay pixel-exact.
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 (corpus-validated
2026-07; feeding 255=hole regenerates the whole frame into stripes).
Like ``erase_lama``, 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]])
session = _get_migan_session()
inp = session.get_inputs() # type: ignore[attr-defined]
img_name, mask_name = inp[0].name, inp[1].name
h, w = image_bgr.shape[:2]
rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
img_in = np.transpose(rgb, (2, 0, 1))[None].astype(np.uint8) # (1,3,H,W)
# invert to MI-GAN polarity: 255 where KNOWN (keep), 0 where hole (erase)
known = (mask <= 127).astype(np.uint8) * 255
mask_in = known[None, None] # (1,1,H,W)
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) # (H',W',3) RGB
if res.shape[:2] != (h, w):
res = cv2.resize(res, (w, h), interpolation=cv2.INTER_LINEAR)
out_bgr = cv2.cvtColor(res, cv2.COLOR_RGB2BGR)
result = image_bgr.copy()
hole = mask > 127
result[hole] = out_bgr[hole]
return result
def erase(
image_bgr: NDArray[Any],
*,
@@ -191,6 +266,12 @@ def erase(
if not mask.any():
return image_bgr.copy()
if backend == "migan":
if not migan_available():
raise RuntimeError(
"MI-GAN backend requires onnxruntime. Install the extra: pip install 'remove-ai-watermarks[migan]'"
)
return erase_migan(image_bgr, mask)
if backend == "lama":
if not lama_available():
raise RuntimeError(
src/remove_ai_watermarks/watermark_registry.py
+147 -10
View File
@@ -41,6 +41,11 @@ if TYPE_CHECKING:
InpaintMethod = Literal["telea", "ns"]
Region = tuple[int, int, int, int]
# Removal method selection for the visible pass. ``auto`` prefers the inpaint
# fallback (LaMa footprint) when the ``lama`` extra is installed, else reverse-alpha
# for marks with a captured alpha map (cv2 inpaint for capture-less marks).
RemovalMethod = Literal["auto", "reverse-alpha", "inpaint"]
@dataclass(frozen=True)
class MarkDetection:
@@ -62,9 +67,10 @@ class KnownMark:
label: str
location: str # usual place, human-readable ("bottom-right")
in_auto: bool # participate in `--mark auto` scanning
recovery: str # removal strategy (all reverse-alpha today)
recovery: str # default removal strategy label ("reverse-alpha")
_detect: Callable[[NDArray[Any]], MarkDetection]
_remove: Callable[..., tuple[NDArray[Any], Region | None]]
has_capture: bool = True # a captured alpha map exists (reverse-alpha possible)
def detect(self, image: NDArray[Any]) -> MarkDetection:
return self._detect(image)
@@ -73,20 +79,29 @@ class KnownMark:
self,
image: NDArray[Any],
*,
method: RemovalMethod = "auto",
inpaint_method: InpaintMethod = "ns",
inpaint: bool = True,
inpaint_strength: float = 0.85,
force: bool = False,
) -> tuple[NDArray[Any], Region | None]:
"""Remove this mark by reverse-alpha; returns ``(result, region)`` where
``region`` is the removed mark's bbox (for residual-inpaint positioning),
or None if nothing was removed. NB: the CLI does NOT use ``region`` to
clear alpha on save -- that zeroing caused the issue-#30 white box.
"""Remove this mark; returns ``(result, region)`` where ``region`` is the
removed mark's bbox (for residual-inpaint positioning), or None if nothing
was removed. NB: the CLI does NOT use ``region`` to clear alpha on save --
that zeroing caused the issue-#30 white box.
``inpaint`` / ``inpaint_strength`` / ``inpaint_method`` tune the Gemini
reverse-alpha edge-residual cleanup only. ``force`` removes at the mark's
usual location even without a positive detection (the ``--no-detect`` path).
``method`` selects the removal path: ``reverse-alpha`` recovers the true
pixels from the captured alpha map (lighter, exact, better on structured
backgrounds); ``inpaint`` erases the footprint with LaMa (or cv2 when the
``lama`` extra is absent), needing no capture; ``auto`` (default) prefers
inpaint when LaMa is installed and falls back to reverse-alpha otherwise
(cv2 inpaint for capture-less marks). ``inpaint``/``inpaint_strength``/
``inpaint_method`` tune the Gemini reverse-alpha edge-residual cleanup only.
``force`` removes at the mark's usual location even without a positive
detection (the ``--no-detect`` path).
"""
if resolve_removal_method(method, self.has_capture) == "inpaint":
return _inpaint_remove(self, image, force)
return self._remove(image, inpaint_method, inpaint, inpaint_strength, force)
@@ -135,11 +150,74 @@ def _engine(key: str) -> Any:
from remove_ai_watermarks.samsung_engine import SamsungEngine
_engines[key] = SamsungEngine()
elif key == "jimeng_pill":
from remove_ai_watermarks.pill_engine import PillEngine
_engines[key] = PillEngine()
else: # pragma: no cover - guarded by the registry keys
raise KeyError(key)
return _engines[key]
def inpaint_model_available() -> bool:
"""True when any ONNX inpaint-model backend (MI-GAN or big-LaMa) can run."""
from remove_ai_watermarks import region_eraser
return region_eraser.migan_available() or region_eraser.lama_available()
def preferred_inpaint_backend() -> str:
"""Backend used by the inpaint fallback: MI-GAN (light, droplet-friendly, the
default) when its ONNX runtime is available, else cv2. big-LaMa is NOT auto-
selected -- it is a heavier explicit opt-in via ``erase --backend lama`` (both
models run on onnxruntime, so availability alone cannot express the user's
intent; the light model is the safe default)."""
from remove_ai_watermarks import region_eraser
return "migan" if region_eraser.migan_available() else "cv2"
def resolve_removal_method(method: RemovalMethod, has_capture: bool) -> Literal["reverse-alpha", "inpaint"]:
"""Resolve the requested method to a concrete one. A capture-less mark has no
alpha map, so it can only be inpainted -- even explicit ``reverse-alpha`` falls
back to inpaint there.
``auto`` uses **reverse-alpha for capture marks** and **inpaint for capture-less**.
Reverse-alpha recovers the true pixels under the mark (measured cleaner than
MI-GAN inpaint on the capture marks -- doubao/gemini/jimeng -- especially on
structured backgrounds, and it needs no model / RAM), so it stays the default
where a capture exists; inpaint is reserved for marks that have no alpha map
(the Jimeng pill). ``--method inpaint`` still forces inpaint for anyone who
wants it."""
# A capture-less mark can only be inpainted; a capture mark inpaints only when
# explicitly asked (auto + reverse-alpha both recover the true pixels).
if method == "inpaint" or not has_capture:
return "inpaint"
return "reverse-alpha"
def _inpaint_remove(mark: KnownMark, image: NDArray[Any], force: bool) -> tuple[NDArray[Any], Region | None]:
"""Remove ``mark`` by inpainting its footprint: the NCC-aligned captured
silhouette (:meth:`footprint_mask`), erased with MI-GAN when its extra is
installed, else cv2 (see :func:`preferred_inpaint_backend`). No-op (returns a
copy) on a clean corner unless ``force``. Gated on the mark's trust-confidence
detection, so a clean image is never touched."""
from remove_ai_watermarks import region_eraser
det = mark.detect(image)
if not (det.detected or force):
return image.copy(), None
engine = _engine(mark.key)
fm = engine.footprint_mask(image, force=force) # uniform signature; text/pill ignore force
if fm is None or not fm.any():
return image.copy(), None
if preferred_inpaint_backend() == "migan":
result = region_eraser.erase_migan(image, fm)
else:
result = region_eraser.erase_cv2(image, fm, radius=6)
return result, (det.region if det.detected else None)
def _gemini_detect(image: NDArray[Any]) -> MarkDetection:
d = _engine("gemini").detect_watermark(image)
detected = bool(d.detected) and d.confidence >= _GEMINI_AUTO_MIN_CONF
@@ -207,11 +285,29 @@ def _text_mark(key: str, label: str, location: str) -> KnownMark:
)
# ── Capture-less mark: the Jimeng-basic "AI生成" pill (top-left, inpaint-only) ──
# No alpha map exists, so removal is inpaint only (has_capture=False routes every
# method to inpaint). Detection is edge-NCC of a synthetic silhouette; see pill_engine.
def _pill_detect(image: NDArray[Any]) -> MarkDetection:
d = _engine("jimeng_pill").detect(image)
return MarkDetection("jimeng_pill", "Jimeng AI生成 pill", "top-left", d.detected, d.confidence, d.region)
def _pill_noop_remove(
image: NDArray[Any], _im: InpaintMethod, _ip: bool, _st: float, _force: bool
) -> tuple[NDArray[Any], Region | None]:
# Capture-less: reverse-alpha is impossible. resolve_removal_method routes every
# method to inpaint (_inpaint_remove), so this reverse-alpha slot is never reached;
# return an untouched copy defensively.
return image.copy(), None
_REGISTRY: tuple[KnownMark, ...] = (
KnownMark("gemini", "Google Gemini sparkle", "bottom-right", True, "reverse-alpha", _gemini_detect, _gemini_remove),
_text_mark("doubao", "Doubao 豆包AI生成 text", "bottom-right"),
_text_mark("jimeng", "Jimeng 即梦AI wordmark", "bottom-right"),
_text_mark("samsung", "Samsung Galaxy AI text", "bottom-left"),
KnownMark("jimeng_pill", "Jimeng AI生成 pill", "top-left", True, "inpaint", _pill_detect, _pill_noop_remove, False),
)
@@ -238,8 +334,7 @@ def detect_marks(image: NDArray[Any], *, include_explicit: bool = True) -> list[
Returns one MarkDetection per scanned mark (``detected`` flags which fired).
``include_explicit=False`` scans only the ``in_auto`` marks -- the set used
by ``--mark auto``.
"""
by ``--mark auto``."""
return [m.detect(image) for m in _REGISTRY if include_explicit or m.in_auto]
@@ -247,3 +342,45 @@ def best_auto_mark(image: NDArray[Any]) -> MarkDetection | None:
"""The highest-confidence detected ``in_auto`` mark, or None if none fired."""
fired = [d for d in detect_marks(image, include_explicit=False) if d.detected]
return max(fired, key=lambda d: d.confidence) if fired else None
def remove_auto_marks(
image: NDArray[Any],
*,
pill_metadata: bool = False,
method: RemovalMethod = "auto",
inpaint_method: InpaintMethod = "ns",
inpaint: bool = True,
inpaint_strength: float = 0.85,
) -> tuple[NDArray[Any], list[str]]:
"""Remove EVERY detected ``in_auto`` mark in one pass, chaining the result.
Marks coexist in different corners -- a Jimeng-basic image carries BOTH the
top-left "AI生成" pill AND the bottom-right "★ 即梦AI" wordmark -- and their
confidences are on different scales, so ``best_auto_mark`` (single strongest)
would clean only one and leave the other (issue #54). Each mark re-detects at
its own corner on the progressively-cleaned image, so order does not matter.
The capture-less ``jimeng_pill`` has a weak edge-NCC detector (~7% raw
false-fire), so it is kept only when the image is CONFIRMED Jimeng and NOT
Doubao. Confirmation is ``pill_metadata`` (the China-AIGC / TC260 metadata
signal, supplied by the caller) OR the reliable bottom-right wordmark firing --
the wordmark keeps recall on metadata-STRIPPED uploads (screenshots / re-saved
files) that the metadata gate alone would miss. Returns ``(result, [labels
removed])``; an empty list means nothing fired."""
fired = [d for d in detect_marks(image, include_explicit=False) if d.detected]
keys = {d.key for d in fired}
jimeng_confirmed = pill_metadata or "jimeng" in keys
if "jimeng_pill" in keys and (not jimeng_confirmed or "doubao" in keys):
fired = [d for d in fired if d.key != "jimeng_pill"]
result = image
for det in fired:
result, _ = get_mark(det.key).remove(
result,
method=method,
inpaint_method=inpaint_method,
inpaint=inpaint,
inpaint_strength=inpaint_strength,
force=False,
)
return result, [d.label for d in fired]