From 99e57c872f4a5ea6cd0feb89d7c4daa326a23bad Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 19 Jun 2026 10:01:07 -0700 Subject: [PATCH] perf(text-mark): footprint-sized arrays in reverse-alpha CPU path 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 --- docs/module-internals.md | 4 +- src/remove_ai_watermarks/_text_mark_engine.py | 75 ++++++---- tests/test_text_mark_memory.py | 134 ++++++++++++++++++ 3 files changed, 184 insertions(+), 29 deletions(-) create mode 100644 tests/test_text_mark_memory.py diff --git a/docs/module-internals.md b/docs/module-internals.md index 3818e05..eb185b2 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -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` diff --git a/src/remove_ai_watermarks/_text_mark_engine.py b/src/remove_ai_watermarks/_text_mark_engine.py index 708a80a..5e4fbc8 100644 --- a/src/remove_ai_watermarks/_text_mark_engine.py +++ b/src/remove_ai_watermarks/_text_mark_engine.py @@ -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 diff --git a/tests/test_text_mark_memory.py b/tests/test_text_mark_memory.py new file mode 100644 index 0000000..44d1931 --- /dev/null +++ b/tests/test_text_mark_memory.py @@ -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)