mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-06 22:18:36 +02:00
95 lines
3.6 KiB
Python
95 lines
3.6 KiB
Python
"""Tests for the Tencent Yuanbao (元宝 / AI生成) visible-watermark engine."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from remove_ai_watermarks import watermark_registry as registry
|
|
from remove_ai_watermarks.yuanbao_engine import (
|
|
_ALPHA_HEIGHT_FRAC,
|
|
_ALPHA_WIDTH_FRAC,
|
|
YuanbaoEngine,
|
|
_alpha_template,
|
|
)
|
|
|
|
_MARK_WIDTH_FRAC = 0.08
|
|
_RIGHT_MARGIN_FRAC = 0.028
|
|
_BOTTOM_MARGIN_FRAC = 0.031
|
|
|
|
|
|
def _compose(w: int, h: int, *, bg: float, foreground: float):
|
|
"""Composite the synthetic Yuanbao silhouette at its measured geometry."""
|
|
image = np.full((h, w, 3), bg, np.float32)
|
|
alpha = _alpha_template()
|
|
assert alpha is not None
|
|
short = min(w, h)
|
|
gw = int(_MARK_WIDTH_FRAC * short)
|
|
gh = max(6, int((_ALPHA_HEIGHT_FRAC / _ALPHA_WIDTH_FRAC) * gw))
|
|
ax = w - int(_RIGHT_MARGIN_FRAC * short) - gw
|
|
ay = h - int(_BOTTOM_MARGIN_FRAC * short) - gh
|
|
mark_alpha = np.zeros((h, w), np.float32)
|
|
mark_alpha[ay : ay + gh, ax : ax + gw] = cv2.resize(alpha, (gw, gh))
|
|
a3 = mark_alpha[:, :, None]
|
|
composed = (a3 * foreground + (1 - a3) * image).clip(0, 255).astype(np.uint8)
|
|
return composed, (ax, ay, gw, gh)
|
|
|
|
|
|
class TestConfig:
|
|
def test_uses_two_polarity_contrast_frontend(self):
|
|
assert YuanbaoEngine().config.detect_frontend == "contrast"
|
|
|
|
def test_strict_only(self):
|
|
assert YuanbaoEngine().config.provenance_ncc_factor == 1.0
|
|
|
|
def test_registry_row(self):
|
|
mark = registry.get_mark("yuanbao")
|
|
assert mark.location == "bottom-right"
|
|
assert "元宝" in mark.label
|
|
assert mark.in_auto
|
|
|
|
|
|
class TestDetectAndMask:
|
|
def test_detects_light_mark_on_dark_background(self):
|
|
watermark, _ = _compose(1024, 1024, bg=60, foreground=230)
|
|
detection = YuanbaoEngine().detect(watermark)
|
|
assert detection.detected
|
|
assert detection.confidence >= 0.80
|
|
|
|
def test_detects_dark_mark_on_light_background(self):
|
|
"""Yuanbao switches mark polarity with the background.
|
|
|
|
A white top-hat alone misses the dark-gray stamp used on pale scenes.
|
|
"""
|
|
watermark, _ = _compose(1024, 1024, bg=230, foreground=110)
|
|
detection = YuanbaoEngine().detect(watermark)
|
|
assert detection.detected
|
|
assert detection.confidence >= 0.80
|
|
|
|
def test_clean_gradient_stays_quiet(self):
|
|
ramp = np.tile(np.linspace(40, 220, 1024, dtype=np.uint8), (1024, 1))
|
|
image = cv2.cvtColor(ramp, cv2.COLOR_GRAY2BGR)
|
|
assert not YuanbaoEngine().detect(image).detected
|
|
|
|
def test_match_must_hug_bottom_right_anchor(self):
|
|
watermark, (ax, ay, gw, gh) = _compose(1024, 1024, bg=60, foreground=230)
|
|
assert YuanbaoEngine().detect(watermark).detected
|
|
shifted = np.full_like(watermark, 60)
|
|
shifted[600 : 600 + gh, 600 : 600 + gw] = watermark[ay : ay + gh, ax : ax + gw]
|
|
assert not YuanbaoEngine().detect(shifted).detected
|
|
|
|
def test_mask_uses_detector_box_for_dark_mark(self):
|
|
watermark, (ax, ay, gw, gh) = _compose(1024, 1024, bg=230, foreground=110)
|
|
mask = YuanbaoEngine().footprint_mask(watermark)
|
|
assert mask is not None
|
|
ys, xs = np.where(mask > 0)
|
|
assert xs.min() <= ax + int(0.05 * gw)
|
|
assert xs.max() >= ax + gw - int(0.05 * gw)
|
|
assert ys.min() <= ay + gh // 2 <= ys.max()
|
|
|
|
def test_remove_clears_detector(self):
|
|
watermark, _ = _compose(1024, 1024, bg=60, foreground=230)
|
|
output, region = registry.get_mark("yuanbao").remove(watermark, backend="cv2")
|
|
assert region is not None
|
|
assert not YuanbaoEngine().detect(output).detected
|