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

perf(text-mark): footprint-sized arrays in reverse-alpha CPU path

Browse Source 
The reverse-alpha text-mark engine (Doubao/Jimeng/Samsung) allocated
full-frame arrays where only the glyph footprint is ever read:

  - _fixed_alpha_map / _aligned_alpha_map each built a full (h, w) float32
    alpha map non-zero only inside the glyph box, and two were held at once
    during removal (~96 MB of mostly-zeros on a 12 MP frame);
  - extract_mask built a full (h, w) uint8 mask that every caller cropped to
    the located box (~12 MB, rebuilt per text-mark detector on the
    memory-tight identify path).

Both now return footprint-sized arrays: the alpha helpers return the
glyph-sized block plus its placement (ax, ay, gw, gh), and extract_mask
returns the box-sized mask. _apply_reverse_alpha consumes the block
directly; the residual inpaint embeds it into one full-frame uint8 mask only
at cv2.inpaint time (which needs a full-frame mask). remove_watermark_
reverse_alpha tracks the winning region alongside best_amap to place it.

Peak allocation drops from O(image*4)x2 + O(image) to O(footprint)x2 +
one gated O(image*1) uint8 mask -- a win every consumer gets, motivated by
the 512 MB raiw.cc worker that OOMs on large decodes. GPU path untouched.

Byte-identical to the old full-frame path (verified: 17 output hashes
across the three engines, inpaint/no-inpaint, detect, and the real
doubao-1.png fixture, unchanged before/after). tests/test_text_mark_memory.py
guards it by reconstructing the old full-frame path inline and asserting
equality, so the proof survives a cv2/asset bump, and pins the O(footprint)
shape so a regression to full-frame fails loudly.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-06-19 10:01:07 -07:00
parent 9614615001
commit 99e57c872f
3 changed files with 184 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/module-internals.md
+3 -1
View File
@@ -103,7 +103,9 @@ The 11 survivors are near-white ill-conditioning (reverse-alpha divides by `1-a`
`_text_mark_engine.py` — **shared base for the three reverse-alpha text-mark engines (Doubao/Jimeng/Samsung), extracted 2026-06-09** (they were ~90% byte-identical clones). `TextMarkEngine(config: TextMarkConfig)` owns the whole `locate → extract_mask → detect → _fixed/_aligned_alpha_map → _apply_reverse_alpha → remove_watermark_reverse_alpha` pipeline (+ the asset-keyed `load_alpha_template`/`glyph_silhouette`/`template_match_score` caches). Each engine module is now a thin subclass: it supplies only its `TextMarkConfig` (the tuned constants, the bundled asset, and the bounded structural deltas — `corner` br/bl, `margin_floor` 4/2, `morph_open_size` 5/3, `min_gw` 8/16) plus the test-facing module shims (`_alpha_template`/`_glyph_silhouette`/`_template_match_score` + the constants). Behavior is byte-exact vs the old per-engine code (the three engine test suites pass unchanged). Gemini stays a SEPARATE engine (its multi-size fixed-slot sparkle model is genuinely different). Add a new text mark = a new `TextMarkConfig` + a thin subclass + one registry `_text_mark(...)` row. The engine bullets below describe each mark's calibration history; the LOGIC lives here.
**`_apply_reverse_alpha` runs on the glyph crop only:** `amap` is zero outside the glyph `region` (x, y, w, h), so the blend is a no-op there (`(wm - 0)/(1 - 0) == wm`, and a uint8→float32→uint8 round-trip is exact). It copies the frame through and computes the reverse-alpha math on the `region` crop only — byte-identical to the old full-frame pass (verified: Doubao 130 + Jimeng 22 placements, 0 mismatches) but O(glyph) not O(image). The full-frame pass cost ~275 ms on a 12 MP frame for a glyph that is <0.1% of it, once per candidate placement (fixed + aligned ≈ 2×/removal); the crop drops that to ~2 ms. Mirror of the Gemini `_core_and_bg` crop. `remove_watermark_reverse_alpha` passes the `region` each `_fixed/_aligned_alpha_map` returns.
**`_apply_reverse_alpha` runs on the glyph crop only:** the blend is a no-op outside the glyph `region` (x, y, w, h) (`(wm - 0)/(1 - 0) == wm`, and a uint8→float32→uint8 round-trip is exact). It copies the frame through and computes the reverse-alpha math on the `region` crop only — byte-identical to the old full-frame pass (verified: Doubao 130 + Jimeng 22 placements, 0 mismatches) but O(glyph) not O(image). The full-frame pass cost ~275 ms on a 12 MP frame for a glyph that is <0.1% of it, once per candidate placement (fixed + aligned ≈ 2×/removal); the crop drops that to ~2 ms. Mirror of the Gemini `_core_and_bg` crop.
**`_fixed/_aligned_alpha_map` and `extract_mask` return footprint-sized arrays, not full frames (memory):** the alpha-map helpers return the glyph-sized alpha **block** (`(gh, gw)` float32) plus its placement `(ax, ay, gw, gh)`, and `extract_mask` returns the box-sized glyph mask (`(loc.h, loc.w)` uint8) — both used to allocate a full `(h, w)` array that is read only inside the small glyph/box. A full-frame float32 alpha map is ~48 MB on a 12 MP frame and two were held at once during removal (fixed + aligned ≈ 96 MB of mostly-zeros); the box mask was a ~12 MB uint8 allocation rebuilt per text-mark `detect` on the memory-tight `identify` path. `_apply_reverse_alpha` consumes the block directly; the residual inpaint embeds it into one full-frame uint8 mask only at `cv2.inpaint` time (which needs a full-frame mask). Byte-identical to the old full-frame path — the block equals the old map's `[ay:ay+gh, ax:ax+gw]` slice and the box equals the old mask cropped to `loc.bbox` (regression-guarded by `tests/test_text_mark_memory.py`, which reconstructs the old full-frame path inline and asserts equality, so the proof survives a cv2/asset bump). `remove_watermark_reverse_alpha` tracks the winning `region` alongside `best_amap` to place that mask.
## `doubao_engine.py`
src/remove_ai_watermarks/_text_mark_engine.py
+47 -28
View File
@@ -184,20 +184,27 @@ class TextMarkEngine:
# ── Mask ────────────────────────────────────────────────────────────
def extract_mask(self, image: NDArray[Any], loc: TextMarkLocation) -> NDArray[Any]:
"""Build a full-image uint8 mask (255 = watermark glyph) for the box.
"""Build a box-sized uint8 mask (255 = watermark glyph) for ``loc``.
Returns just the glyph mask of the located box (shape ``(loc.h, loc.w)``),
not a full-frame array: every caller immediately crops to ``loc.bbox``, so
allocating a full ``(h, w)`` mask and embedding the box was O(image) work
and memory for an O(box) result -- a wasted full-frame uint8 allocation on
each detect (~12 MB on a 12 MP frame, recomputed per text-mark detector on
the memory-tight identify path). The box mask is byte-identical to the old
full-frame mask cropped to ``loc.bbox``.
Polarity-aware: the mark is a light, low-saturation gray rendered brighter
than the local background (white top-hat), so a white-paper document is left
untouched (nothing brighter than its surroundings is masked there).
"""
c = self.config
h, w = image.shape[:2]
x, y, bw, bh = loc.bbox
# A degenerate ROI (a sliver from an extremely wide/short image) cannot hold
# the mark and would feed cv2's GaussianBlur/morphology a ~1-px-tall array,
# which can fault native code on some platforms. Skip the cv2 pipeline.
if bh < 16 or bw < 16:
return np.zeros((h, w), np.uint8)
return np.zeros((bh, bw), np.uint8)
# Normalize the ROI to 3-channel BGR (grayscale / BGRA would break axis=2).
roi = image_io.to_bgr(image[y : y + bh, x : x + bw]).astype(np.float32)
@@ -216,11 +223,7 @@ class TextMarkEngine:
glyph = cand.astype(np.uint8) * 255
glyph = cv2.morphologyEx(glyph, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))
k = c.morph_open_size
glyph = cv2.morphologyEx(glyph, cv2.MORPH_OPEN, np.ones((k, k), np.uint8))
mask = np.zeros((h, w), np.uint8)
mask[y : y + bh, x : x + bw] = glyph
return mask
return cv2.morphologyEx(glyph, cv2.MORPH_OPEN, np.ones((k, k), np.uint8))
# ── Detect ──────────────────────────────────────────────────────────
@@ -232,9 +235,8 @@ class TextMarkEngine:
if image is None or image.size == 0:
return det
loc = self.locate(image)
mask = self.extract_mask(image, loc)
x, y, bw, bh = loc.bbox
box = mask[y : y + bh, x : x + bw]
box = self.extract_mask(image, loc) # box-sized mask (== old full-frame cropped to bbox)
_x, _y, bw, bh = loc.bbox
coverage = float((box > 0).sum()) / float(max(1, bw * bh))
det.region = loc.bbox
det.coverage = coverage
@@ -254,7 +256,15 @@ class TextMarkEngine:
def _fixed_alpha_map(self, image: NDArray[Any]) -> tuple[NDArray[Any], tuple[int, int, int, int]] | None:
"""Place the template by fixed width-relative geometry (pixel-exact at the
captured width)."""
captured width).
Returns the glyph-sized alpha BLOCK (shape ``(gh, gw)``) plus its placement
``(ax, ay, gw, gh)``, not a full-frame ``(h, w)`` map. The map is non-zero
only inside the glyph box and every consumer reads exactly that box, so a
full-frame float32 map was O(image*4 bytes) of mostly zeros -- ~48 MB on a
12 MP frame, and two were held at once (fixed + aligned). The block is
byte-identical to the old full-frame map's ``[ay:ay+gh, ax:ax+gw]`` slice.
"""
c = self.config
at = self._alpha_template()
if at is None:
@@ -268,22 +278,23 @@ class TextMarkEngine:
else: # bottom-left
ax = min(max(0, int(c.alpha_margin_x_frac * w)), max(0, w - gw))
ay = max(0, h - int(c.alpha_margin_bottom_frac * w) - gh)
amap = np.zeros((h, w), np.float32)
amap[ay : ay + gh, ax : ax + gw] = cv2.resize(at, (gw, gh), interpolation=cv2.INTER_LINEAR)
return amap, (ax, ay, gw, gh)
block = cv2.resize(at, (gw, gh), interpolation=cv2.INTER_LINEAR)
return block, (ax, ay, gw, gh)
def _aligned_alpha_map(self, image: NDArray[Any]) -> tuple[NDArray[Any], tuple[int, int, int, int]] | None:
"""Register the captured template to the actual mark via a TM_CCOEFF_NORMED
scale + position search. Returns ``(alpha_map, glyph_bbox)`` or None."""
scale + position search. Returns the glyph-sized alpha BLOCK and its
placement ``(ax, ay, gw, gh)`` (see :meth:`_fixed_alpha_map` for why the
block, not a full-frame map), or None."""
c = self.config
at = self._alpha_template()
sil = self._glyph_silhouette()
if at is None or sil is None:
return None
h, w = image.shape[:2]
w = image.shape[1]
loc = self.locate(image)
bx, by, bw, bh = loc.bbox
box_mask = self.extract_mask(image, loc)[by : by + bh, bx : bx + bw]
box_mask = self.extract_mask(image, loc) # box-sized (== old full-frame cropped to bbox)
expected = c.alpha_width_frac * w
best: tuple[float, int, int, int, int] | None = None
for scale in np.linspace(*c.alpha_align_search):
@@ -298,18 +309,17 @@ class TextMarkEngine:
return None
_, gw, gh, ox, oy = best
ax, ay = bx + ox, by + oy
amap = np.zeros((h, w), np.float32)
amap[ay : ay + gh, ax : ax + gw] = cv2.resize(at, (gw, gh), interpolation=cv2.INTER_LINEAR)
return amap, (ax, ay, gw, gh)
block = cv2.resize(at, (gw, gh), interpolation=cv2.INTER_LINEAR)
return block, (ax, ay, gw, gh)
def _apply_reverse_alpha(
self, image: NDArray[Any], amap: NDArray[Any], region: tuple[int, int, int, int]
) -> NDArray[Any]:
"""Invert the alpha blend with ``amap``: ``original = (wm - a*logo)/(1-a)``.
``amap`` is zero everywhere except the glyph ``region`` (x, y, w, h), so the
blend is a no-op (``(wm - 0)/(1 - 0) == wm``) outside it. Compute the math on
the glyph crop only and copy the rest through unchanged -- byte-identical to a
``amap`` is the glyph-sized alpha BLOCK for ``region`` (x, y, w, h); outside
it the blend is a no-op (``(wm - 0)/(1 - 0) == wm``). Compute the math on the
glyph crop only and copy the rest through unchanged -- byte-identical to a
full-frame pass (a uint8 round-trip through float32 is exact), but O(glyph)
instead of O(image): a full-frame pass costs ~275 ms on a 12 MP frame for a
glyph that is <0.1% of it, and it runs once per candidate placement.
@@ -319,7 +329,7 @@ class TextMarkEngine:
x2, y2 = x1 + gw, y1 + gh
if y1 >= y2 or x1 >= x2:
return out
a3 = np.clip(amap[y1:y2, x1:x2], 0.0, 1.0)[:, :, None]
a3 = np.clip(amap, 0.0, 1.0)[:, :, None]
logo = np.array(self.config.alpha_logo_bgr, np.float32)
roi = out[y1:y2, x1:x2].astype(np.float32)
out[y1:y2, x1:x2] = np.clip((roi - a3 * logo) / np.clip(1.0 - a3, 0.25, 1.0), 0, 255).astype(np.uint8)
@@ -351,16 +361,25 @@ class TextMarkEngine:
return image.copy()
best_out: NDArray[Any] | None = None
best_amap: NDArray[Any] | None = None
best_region: tuple[int, int, int, int] | None = None
best_residual = float("inf")
for amap, region in maps:
out = self._apply_reverse_alpha(image, amap, region)
residual = self.detect(out).confidence
if residual < best_residual:
best_residual, best_out, best_amap = residual, out, amap
if best_out is None or best_amap is None: # pragma: no cover - maps is non-empty
best_residual, best_out, best_amap, best_region = residual, out, amap, region
if best_out is None or best_amap is None or best_region is None: # pragma: no cover - maps is non-empty
return image.copy()
if residual_inpaint:
# Embed the glyph-sized alpha block into a full-frame uint8 mask only for
# the inpaint (cv2.inpaint needs a mask matching best_out). One uint8
# full-frame array, built once, vs the old two full-frame float32 maps;
# byte-identical to thresholding the old full-frame float32 map (zero
# outside the block, so the dilate/inpaint see the same mask).
ax, ay, gw, gh = best_region
rm = np.zeros(best_out.shape[:2], np.uint8)
rm[ay : ay + gh, ax : ax + gw] = (best_amap > c.residual_alpha_floor).astype(np.uint8) * 255
kernel = np.ones((c.residual_dilate, c.residual_dilate), np.uint8)
rm = cv2.dilate((best_amap > c.residual_alpha_floor).astype(np.uint8) * 255, kernel)
rm = cv2.dilate(rm, kernel)
best_out = cv2.inpaint(best_out, rm, c.residual_inpaint_radius, cv2.INPAINT_NS)
return best_out
tests/test_text_mark_memory.py
+134
View File
@@ -0,0 +1,134 @@
"""Byte-identity guards for the text-mark engine memory optimization.
The reverse-alpha text-mark engine used to allocate full-frame arrays where only
the glyph footprint is ever read:
* ``extract_mask`` built a full ``(h, w)`` uint8 mask and every caller cropped
it to the located box;
* ``_fixed_alpha_map`` / ``_aligned_alpha_map`` each built a full ``(h, w)``
float32 alpha map that is non-zero only inside the glyph box, and two were
held at once during removal.
Both now return footprint-sized arrays. These tests prove the new footprint-sized
path is BYTE-IDENTICAL to the old full-frame path by reconstructing the old
behavior inline from the new building blocks (so the proof survives a cv2/asset
version bump, unlike a pinned output hash), and lock in the O(footprint) memory
characteristic so a regression back to a full-frame allocation fails loudly.
"""
from __future__ import annotations
import cv2
import numpy as np
import pytest
import remove_ai_watermarks.doubao_engine as D
import remove_ai_watermarks.jimeng_engine as J
import remove_ai_watermarks.samsung_engine as S
from remove_ai_watermarks.doubao_engine import DoubaoEngine
from remove_ai_watermarks.jimeng_engine import JimengEngine
from remove_ai_watermarks.samsung_engine import SamsungEngine
# (engine factory, engine module) for each reverse-alpha text mark.
ENGINES = [
pytest.param(DoubaoEngine, D, id="doubao"),
pytest.param(JimengEngine, J, id="jimeng"),
pytest.param(SamsungEngine, S, id="samsung"),
]
def _watermarked(engine, module) -> np.ndarray:
"""Composite the engine's real alpha glyph onto a flat mid-gray field at the
captured native width (so both placement candidates fire)."""
cfg = engine.config
nw = module._ALPHA_NATIVE_WIDTH
at = module._alpha_template()
gw, gh = int(cfg.alpha_width_frac * nw), int(cfg.alpha_height_frac * nw)
ax = (nw - int(cfg.alpha_margin_x_frac * nw) - gw) if cfg.corner == "br" else int(cfg.alpha_margin_x_frac * nw)
ay = nw - int(cfg.alpha_margin_bottom_frac * nw) - gh
amap = np.zeros((nw, nw), np.float32)
amap[ay : ay + gh, ax : ax + gw] = cv2.resize(at, (gw, gh))
a3 = amap[:, :, None]
img = np.full((nw, nw, 3), 100.0, np.float32)
return (a3 * np.array(cfg.alpha_logo_bgr, np.float32) + (1 - a3) * img).clip(0, 255).astype(np.uint8)
@pytest.mark.parametrize(("factory", "module"), ENGINES)
class TestExtractMaskFootprint:
def test_returns_box_sized_mask(self, factory, module):
eng = factory()
img = _watermarked(eng, module)
loc = eng.locate(img)
box = eng.extract_mask(img, loc)
assert box.dtype == np.uint8
# Shape == loc.bbox, i.e. the old full-frame mask's [y:y+bh, x:x+bw] crop.
assert box.shape == (loc.h, loc.w)
# Footprint, not full frame: the box is a tiny fraction of the image.
assert box.size * 4 < img.shape[0] * img.shape[1]
@pytest.mark.parametrize(("factory", "module"), ENGINES)
class TestAlphaMapFootprint:
def test_maps_are_footprint_sized_blocks(self, factory, module):
eng = factory()
img = _watermarked(eng, module)
for placed in (eng._fixed_alpha_map(img), eng._aligned_alpha_map(img)):
assert placed is not None
block, (ax, ay, gw, gh) = placed
assert block.dtype == np.float32
assert block.shape == (gh, gw)
# The placement stays fully inside the image (no clipping needed).
assert ax >= 0
assert ax + gw <= img.shape[1]
assert ay >= 0
assert ay + gh <= img.shape[0]
# O(footprint): far smaller than the frame.
assert block.size * 4 < img.shape[0] * img.shape[1]
def test_apply_reverse_alpha_equals_old_fullframe(self, factory, module):
"""``_apply_reverse_alpha`` with the glyph block is byte-identical to the
old full-frame path: rebuild the full ``(h, w)`` map, run the old-style
full-frame reverse-alpha, and compare to the new block-based output."""
eng = factory()
img = _watermarked(eng, module)
h, w = img.shape[:2]
for placed in (eng._fixed_alpha_map(img), eng._aligned_alpha_map(img)):
assert placed is not None
block, region = placed
ax, ay, gw, gh = region
new_out = eng._apply_reverse_alpha(img, block, region)
# Old behavior: a full-frame map, indexed by region inside _apply_reverse_alpha.
full = np.zeros((h, w), np.float32)
full[ay : ay + gh, ax : ax + gw] = block
old_out = img.copy()
a3 = np.clip(full[ay : ay + gh, ax : ax + gw], 0.0, 1.0)[:, :, None]
logo = np.array(eng.config.alpha_logo_bgr, np.float32)
roi = old_out[ay : ay + gh, ax : ax + gw].astype(np.float32)
old_out[ay : ay + gh, ax : ax + gw] = np.clip(
(roi - a3 * logo) / np.clip(1.0 - a3, 0.25, 1.0), 0, 255
).astype(np.uint8)
assert np.array_equal(new_out, old_out)
def test_residual_mask_equals_old_fullframe(self, factory, module):
"""The residual inpaint mask built from the block embedded in a full-frame
canvas equals thresholding the old full-frame float32 map (zero outside the
block), so the dilate + inpaint see the same mask."""
eng = factory()
img = _watermarked(eng, module)
h, w = img.shape[:2]
cfg = eng.config
block, (ax, ay, gw, gh) = eng._fixed_alpha_map(img)
# New: embed the block into a uint8 canvas, then threshold.
new_mask = np.zeros((h, w), np.uint8)
new_mask[ay : ay + gh, ax : ax + gw] = (block > cfg.residual_alpha_floor).astype(np.uint8) * 255
# Old: a full-frame float32 map, thresholded everywhere.
old_full = np.zeros((h, w), np.float32)
old_full[ay : ay + gh, ax : ax + gw] = block
old_mask = (old_full > cfg.residual_alpha_floor).astype(np.uint8) * 255
assert np.array_equal(new_mask, old_mask)