mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-06-05 10:38:00 +02:00
test-user
bc3228d387
Add deterministic, CPU-only removal of the visible Doubao "豆包AI生成" mark and
a position-agnostic region eraser for any other visible watermark/logo.
- doubao_engine.py: locate (geometry, scales with width) + polarity-aware
white-top-hat glyph mask + cv2 inpaint; coverage-gated detection and a
dense-text safety guard. No GPU, ~30ms.
- region_eraser.py + `erase` command: inpaint arbitrary --region box(es).
Default cv2 backend (no deps); optional big-LaMa via onnxruntime (`lama`
extra, Carve/LaMa-ONNX, model downloaded on first use, never bundled).
- cli `visible --mark auto|gemini|doubao`: auto routes by detector confidence.
- tests for both engines; seed previously-unseeded CLI image fixtures to stop
the Doubao detector flaking on random corners.
- .gitignore: doubao_capture/{seeds,captures} scratch (alpha-map calibration).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
99 lines
3.8 KiB
Python
99 lines
3.8 KiB
Python
"""Tests for the Doubao visible-watermark engine."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from remove_ai_watermarks.doubao_engine import DoubaoEngine, load_image_bgr
|
|
|
|
SAMPLE = Path(__file__).resolve().parents[1] / "data" / "samples" / "doubao-1.png"
|
|
|
|
|
|
# ── Locate ──────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestLocate:
|
|
def test_box_anchored_bottom_right(self):
|
|
eng = DoubaoEngine()
|
|
img = np.zeros((2048, 2048, 3), np.uint8)
|
|
loc = eng.locate(img)
|
|
# right and bottom edges sit close to the image corner (within margins)
|
|
assert 2048 - (loc.x + loc.w) < int(2048 * 0.03)
|
|
assert 2048 - (loc.y + loc.h) < int(2048 * 0.03)
|
|
assert loc.is_fallback # geometry anchor, no bundled template yet
|
|
|
|
def test_box_scales_with_width(self):
|
|
eng = DoubaoEngine()
|
|
small = eng.locate(np.zeros((1024, 1024, 3), np.uint8))
|
|
large = eng.locate(np.zeros((2048, 2048, 3), np.uint8))
|
|
# width-relative geometry: 2x wider image -> ~2x wider box
|
|
assert large.w == pytest.approx(small.w * 2, rel=0.1)
|
|
|
|
|
|
# ── Detect + remove on the real sample ──────────────────────────────
|
|
|
|
|
|
@pytest.mark.skipif(not SAMPLE.exists(), reason="sample image not present")
|
|
class TestRealSample:
|
|
def test_detects_watermark(self):
|
|
eng = DoubaoEngine()
|
|
det = eng.detect(load_image_bgr(SAMPLE))
|
|
assert det.detected
|
|
assert det.confidence > 0.0
|
|
assert det.coverage > 0.04
|
|
|
|
def test_remove_reduces_glyph_coverage(self):
|
|
eng = DoubaoEngine()
|
|
img = load_image_bgr(SAMPLE)
|
|
before = eng.detect(img).coverage
|
|
out = eng.remove_watermark(img)
|
|
after = eng.detect(out).coverage
|
|
# the inpaint should clear most glyph pixels from the corner box
|
|
assert after < before * 0.5
|
|
|
|
def test_pixels_outside_box_untouched(self):
|
|
eng = DoubaoEngine()
|
|
img = load_image_bgr(SAMPLE)
|
|
out = eng.remove_watermark(img)
|
|
# top-left quadrant is far from the bottom-right mark: must be identical
|
|
h, w = img.shape[:2]
|
|
assert np.array_equal(img[: h // 2, : w // 2], out[: h // 2, : w // 2])
|
|
|
|
|
|
# ── Negative + safety guard ─────────────────────────────────────────
|
|
|
|
|
|
class TestNegativeAndGuard:
|
|
def test_clean_image_not_detected(self):
|
|
eng = DoubaoEngine()
|
|
# smooth gradient, no watermark
|
|
ramp = np.tile(np.linspace(0, 255, 1024, dtype=np.uint8), (1024, 1))
|
|
img = cv2.cvtColor(ramp, cv2.COLOR_GRAY2BGR)
|
|
det = eng.detect(img)
|
|
assert not det.detected
|
|
|
|
def test_clean_image_returned_unchanged(self):
|
|
eng = DoubaoEngine()
|
|
ramp = np.tile(np.linspace(0, 255, 1024, dtype=np.uint8), (1024, 1))
|
|
img = cv2.cvtColor(ramp, cv2.COLOR_GRAY2BGR)
|
|
out = eng.remove_watermark(img)
|
|
assert np.array_equal(img, out)
|
|
|
|
def test_document_background_guard(self):
|
|
"""A dense high-frequency corner (document-like) trips the coverage
|
|
guard, so the image is left untouched rather than smeared."""
|
|
eng = DoubaoEngine()
|
|
rng = np.random.default_rng(0)
|
|
img = np.full((1024, 1024, 3), 255, np.uint8)
|
|
# fill the bottom-right box area with random grayish text-like noise
|
|
loc = eng.locate(img)
|
|
x, y, bw, bh = loc.bbox
|
|
noise = rng.integers(150, 246, size=(bh, bw), dtype=np.uint8)
|
|
img[y : y + bh, x : x + bw] = noise[:, :, None]
|
|
out = eng.remove_watermark(img)
|
|
assert np.array_equal(img, out)
|