diff --git a/docs/module-internals.md b/docs/module-internals.md index aae8984..83f771a 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -850,6 +850,17 @@ box grows), not the whole `footprint_mask`. Baidu extends right to the corner ta and LiblibAI extends left to the triangle logo; both inherit every guard around that arithmetic. +Doubao is the deliberate exception because its captured alpha supplies more +information than a rectangle. Its continuous top-hat response locates the mark, +then `DoubaoEngine.footprint_mask` resizes the alpha to that same winning +`match_box` and masks only the glyphs. The canonical 2048-pixel fixture has bright +branch texture behind the bottom-edge wordmark; bounding the thresholded response, +padding it, and dilating it produced a solid 404-by-134 mask clamped to the bottom +and right edges. OpenCV then had no context beyond either edge and filled the hole +with large triangular wedges. The aligned sparse mask keeps both frame edges +untouched and still clears the detector. `force` cannot align a missing detection, +so it retains the shared geometry-box fallback. + Yuanbao uses the polarity-independent `contrast` front end because its standard two-line mark can be light on dark scenes or dark on light scenes. Its detector and footprint both use the same best-match box. The separate one-line overlay diff --git a/src/remove_ai_watermarks/doubao_engine.py b/src/remove_ai_watermarks/doubao_engine.py index 6ee0b29..a599e9b 100644 --- a/src/remove_ai_watermarks/doubao_engine.py +++ b/src/remove_ai_watermarks/doubao_engine.py @@ -5,8 +5,10 @@ Doubao (ByteDance) stamps every generated image with a visible "豆包AI生成" label mandated by China's TC260 standard, a near-white semi-transparent overlay. Detection matches the bundled glyph silhouette against the corner candidate; removal -is the shared **localize -> fill** (the glyph-bbox :meth:`footprint_mask` feeds -``region_eraser``), NOT reverse-alpha. This module shares +is the shared **localize -> fill**. Doubao reuses that aligned silhouette as a sparse +glyph footprint instead of bounding the thresholded corner response: bright scene +texture behind this bottom-edge mark can otherwise expand a solid mask to the frame +edges, where classical inpainting has no outside context. This module shares :class:`remove_ai_watermarks._text_mark_engine.TextMarkEngine` and supplies only Doubao's tuned :class:`TextMarkConfig` (bottom-right corner, ``assets/doubao_alpha.png`` -- the detection silhouette, rebuilt by @@ -22,8 +24,11 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any -from remove_ai_watermarks import _text_mark_engine -from remove_ai_watermarks._text_mark_engine import TextMarkConfig, TextMarkEngine +import cv2 +import numpy as np + +from remove_ai_watermarks import _text_mark_engine, image_io +from remove_ai_watermarks._text_mark_engine import TextMarkConfig, TextMarkDetection, TextMarkEngine if TYPE_CHECKING: from numpy.typing import NDArray @@ -62,12 +67,20 @@ DETECT_MIN_COVERAGE = 0.04 DETECT_NCC_THRESHOLD = 0.50 # Detection-silhouette geometry, emitted by scripts/visible_alpha_solve.py at the -# captured width. Sizes the glyph silhouette for the TM_CCOEFF_NORMED detection match -# (removal is the template-free glyph-bbox footprint mask, not this template). +# captured width. Sizes the glyph silhouette for the TM_CCOEFF_NORMED detection match; +# Doubao also reuses the aligned alpha as its sparse removal footprint. _ALPHA_NATIVE_WIDTH = 2048 _ALPHA_WIDTH_FRAC = 0.1636 # asset width / image width -- sizes the detection silhouette _ALPHA_HEIGHT_FRAC = 0.0405 +# The captured alpha is also the removal footprint after the detector aligns it. +# Keep faint anti-aliased edges, then grow by one pixel so cv2 does not leave a halo. +# A solid enclosing rectangle is unsafe here: the mark sits close to two frame edges, +# and a textured branch in the canonical sample made the thresholded response reach +# both edges before OpenCV inpainted the whole block into triangular wedges. +_FOOTPRINT_ALPHA_FLOOR = 0.05 +_FOOTPRINT_DILATE = 1 + _CONFIG = TextMarkConfig( name="Doubao", asset_name="doubao_alpha.png", @@ -115,3 +128,52 @@ class DoubaoEngine(TextMarkEngine): def __init__(self) -> None: super().__init__(_CONFIG) + + def footprint_mask( + self, + image: NDArray[Any] | None, + *, + force: bool = False, + dilate: int | None = None, + detection: TextMarkDetection | None = None, + ) -> NDArray[Any] | None: + """Return the detector-aligned Doubao glyph footprint. + + The continuous top-hat response is suitable for locating the wordmark but is + not a safe removal mask: unrelated bright texture can join the response and + enlarge its bounding rectangle. The detector already carries the winning + template box, so resize the captured alpha to that exact box and mask only its + glyphs. Explicit ``force`` has no trustworthy alignment and retains the shared + geometry-box behavior. + """ + if force: + return super().footprint_mask(image, force=True, dilate=dilate, detection=detection) + if image is None or image.size == 0: + return None + + image = image_io.to_bgr(image) + det = detection if detection is not None else self.detect(image) + if not det.detected or det.match_box is None: + return super().footprint_mask(image, force=False, dilate=dilate, detection=det) + alpha = _alpha_template() + if alpha is None: + return super().footprint_mask(image, force=False, dilate=dilate, detection=det) + + loc = self.locate(image) + x0, y0, x1, y1 = det.match_box # inclusive ROI-local coordinates + x0, y0 = max(0, x0), max(0, y0) + x1, y1 = min(loc.w - 1, x1), min(loc.h - 1, y1) + if x1 < x0 or y1 < y0: + return None + glyph_w, glyph_h = x1 - x0 + 1, y1 - y0 + 1 + aligned = cv2.resize(alpha, (glyph_w, glyph_h), interpolation=cv2.INTER_LINEAR) + + local_mask = np.zeros((loc.h, loc.w), np.uint8) + local_mask[y0 : y0 + glyph_h, x0 : x0 + glyph_w] = (aligned > _FOOTPRINT_ALPHA_FLOOR).astype(np.uint8) * 255 + radius = _FOOTPRINT_DILATE if dilate is None else max(0, dilate) + if radius > 0: + kernel = np.ones((2 * radius + 1, 2 * radius + 1), np.uint8) + local_mask = cv2.dilate(local_mask, kernel) + mask = np.zeros(image.shape[:2], np.uint8) + mask[loc.y : loc.y + loc.h, loc.x : loc.x + loc.w] = local_mask + return mask diff --git a/src/remove_ai_watermarks/watermark_registry.py b/src/remove_ai_watermarks/watermark_registry.py index 669b509..b53ce35 100644 --- a/src/remove_ai_watermarks/watermark_registry.py +++ b/src/remove_ai_watermarks/watermark_registry.py @@ -5,8 +5,8 @@ A single catalog that ties each known visible mark to (a) where it usually sits, registry detects every known mark in its usual place and removes the ones present. -**Localize -> fill.** A known mark is removed by LOCALIZING it (a template-free, -version-robust detector that returns a binary footprint MASK) and then handing +**Localize -> fill.** A known mark is removed by LOCALIZING it (a version-robust +detector that returns a binary footprint MASK) and then handing that mask to ONE shared, swappable fill backend (``region_eraser``: cv2 Telea/NS, MI-GAN, or big-LaMa). No mark carries a reverse-alpha step any more: the old ``original = (wm - a*logo)/(1-a)`` recovery depended on a fixed captured alpha map @@ -519,8 +519,9 @@ def _gemini_mask( # The text-mark engines share the TextMarkEngine interface, so one parameterized # adapter pair drives all of them -- a new # text mark is one `_text_mark(...)` row below, not another copy-paste of these -# bodies. Detection matches the glyph silhouette; the mask is the template-free -# glyph-bbox footprint (see TextMarkEngine.footprint_mask). +# bodies. Detection matches the glyph silhouette; the default mask is the +# template-free glyph-bbox footprint (see TextMarkEngine.footprint_mask), while an +# engine may override it when the mark has a measured safer footprint. def _text_mark_detect(key: str, label: str, location: str) -> Callable[..., MarkDetection]: def detect(image: NDArray[Any], *, provenance: bool = False) -> MarkDetection: d = _engine(key).detect(image, provenance=provenance) diff --git a/tests/test_doubao_engine.py b/tests/test_doubao_engine.py index 95506b4..81805a0 100644 --- a/tests/test_doubao_engine.py +++ b/tests/test_doubao_engine.py @@ -114,7 +114,9 @@ class TestRealSample: assert before.detected out, region = registry.get_mark("doubao").remove(img, backend="cv2") assert region is not None - assert DoubaoEngine().detect(out).confidence < before.confidence + after = DoubaoEngine().detect(out) + assert not after.detected + assert after.confidence < before.confidence def test_far_region_untouched(self): img = load_image_bgr(SAMPLE) @@ -122,6 +124,24 @@ class TestRealSample: h, w = img.shape[:2] assert np.array_equal(img[: h // 2, : w // 2], out[: h // 2, : w // 2]) + def test_footprint_stays_on_glyphs_and_off_frame_edges(self): + """Do not turn textured content behind the mark into a solid edge mask. + + The sample's branch contributes bright top-hat pixels below the wordmark. + Bounding that response and dilating it used to clamp a solid 404x134 mask + to both frame edges, so OpenCV propagated the remaining context into large + triangular artifacts. The detector already aligned the captured glyph; the + removal footprint must follow that glyph instead of the background response. + """ + img = load_image_bgr(SAMPLE) + mask = DoubaoEngine().footprint_mask(img) + assert mask is not None + assert not mask[-1].any() + assert not mask[:, -1].any() + ys, xs = np.where(mask > 0) + bounding_area = (int(xs.max()) - int(xs.min()) + 1) * (int(ys.max()) - int(ys.min()) + 1) + assert np.count_nonzero(mask) < 0.65 * bounding_area + # ── Alpha asset + footprint mask + localize -> fill removal ───────── @@ -159,7 +179,7 @@ class TestFootprintMaskAndRemoval: ("w", "h"), [ (_ALPHA_NATIVE_WIDTH, _ALPHA_NATIVE_WIDTH), # captured width - (1773, 2364), # 3:4 portrait -> template-free footprint generalizes + (1773, 2364), # 3:4 portrait -> aligned footprint generalizes ], ) def test_fill_removes_and_leaves_far_region(self, w, h):