Add Tencent Yuanbao visible watermark removal

tests/test_watermark_registry.py

@@ -20,6 +20,7 @@ class TestCatalog:
             "jimeng",
             "qwen",
             "kling",
+            "yuanbao",
             "samsung",
             "runninghub",
             "baidu",
@@ -42,6 +43,7 @@ class TestCatalog:
         assert by_key["gemini"].location == "bottom-right"
         assert by_key["doubao"].location == "bottom-right"
         assert by_key["jimeng"].location == "bottom-right"
+        assert by_key["yuanbao"].location == "bottom-right"
         assert by_key["samsung"].location == "bottom-left"
         assert by_key["jimeng_pill"].location == "top-left"
 
@@ -60,6 +62,7 @@ class TestScan:
             "jimeng",
             "qwen",
             "kling",
+            "yuanbao",
             "samsung",
             "runninghub",
             "baidu",
@@ -89,7 +92,7 @@ class TestScan:
         forced remove on a zero-size ndarray crashed (cv2.error on an empty Mat). detect
         already guarded this; footprint_mask must too. Covers the text + gemini engines."""
         empty = np.zeros(shape, np.uint8)
-        for key in ("doubao", "jimeng", "qwen", "samsung", "gemini"):
+        for key in ("doubao", "jimeng", "qwen", "yuanbao", "samsung", "gemini"):
             _result, mask = reg.get_mark(key).remove(empty, force=True)
             assert mask is None
 
@@ -354,6 +357,16 @@ class TestArbiter:
         assert "qwen" in keys
         assert "jimeng_pill" not in keys
 
+    def test_pill_dropped_on_yuanbao(self):
+        # The standard Yuanbao mark identifies a different TC260 product, so a
+        # coincident top-left pill match must not be treated as Jimeng-basic.
+        cands = [
+            self._c("yuanbao", strict=True, relaxed=True),
+            self._c("jimeng_pill", strict=True, relaxed=True, flat=True),
+        ]
+        keys = self._keys(cands, reg.Context(provenance=frozenset({"jimeng"})))
+        assert keys == {"yuanbao"}
+
     def test_pill_metadata_arm_gated_on_flatness(self):
         ctx = reg.Context(provenance=frozenset({"jimeng"}))
         assert self._keys([self._c("jimeng_pill", strict=True, relaxed=True, flat=True)], ctx) == {"jimeng_pill"}

tests/test_yuanbao_engine.py

@@ -0,0 +1,94 @@
+"""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