Improve visible-mark evaluation and Samsung masking

scripts/detector_response.py

@@ -82,7 +82,9 @@ REPO = Path(__file__).resolve().parents[1]
 sys.path.insert(0, str(REPO / "src"))
 sys.path.insert(0, str(REPO / "scripts"))
 
-from fill_quality import SLOT_STAMPABLE, STAMPABLE, clean_sources, stamp_any, texture_of  # noqa: E402
+from fill_quality import clean_sources, stamp_any, texture_of  # noqa: E402
+
+from remove_ai_watermarks import watermark_registry as wr  # noqa: E402
 
 OUT = REPO / ".local-eval" / "detector-response.jsonl"
 
@@ -91,7 +93,7 @@ OUT = REPO / ".local-eval" / "detector-response.jsonl"
 # front-end was introduced for) and above it (a larger re-render, or an upscaled upload).
 SIZES = (0.6, 0.8, 1.0, 1.3)
 ALPHAS = (0.3, 0.5, 0.75, 1.0)
-MARKS = (*STAMPABLE, *SLOT_STAMPABLE)
+MARKS = tuple(wr.mark_keys())
 
 
 def aspect_of(image: np.ndarray) -> str:

scripts/fill_quality.py

@@ -53,25 +53,6 @@ from invisible_quality_audit import _ssim  # noqa: E402 -- path selects this wor
 CORPUS = REPO / ".local-eval" / "originals"
 OUT = REPO / ".local-eval" / "fill-quality.jsonl"
 
-# Text marks: a bundled alpha PNG plus the engine's own corner geometry.
-STAMPABLE = (
-    "doubao",
-    "jimeng",
-    "qwen",
-    "kling",
-    "yuanbao",
-    "samsung",
-    "runninghub",
-    "baidu",
-    "liblib",
-    "microsoft",
-)
-# Marks with no bundled alpha asset. Both expose a default footprint via
-# `footprint_mask(force=True)`, so they are stamped by fitting their own alpha source
-# into that slot: Gemini's alpha is derived from its background captures, the pill's
-# from its synthetic font-rendered silhouette.
-SLOT_STAMPABLE = ("gemini", "jimeng_pill")
-
 
 def psnr(a: np.ndarray, b: np.ndarray) -> float:
     mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
@@ -86,122 +67,13 @@ def texture_of(box: np.ndarray) -> float:
     return float(np.median(cv2.magnitude(gx, gy)))
 
 
-def engine_for(mark_key: str) -> Any:
-    """The TextMarkEngine instance for a mark key (the registry does not expose one)."""
-    import importlib
-
-    mod = importlib.import_module(f"remove_ai_watermarks.{mark_key}_engine")
-    cls = next(
-        obj
-        for name, obj in vars(mod).items()
-        # Exclude the imported base class: it takes a `config` arg, the subclasses do not.
-        if isinstance(obj, type) and name.endswith("Engine") and name != "TextMarkEngine"
-    )
-    return cls()
-
-
-def composite(image: np.ndarray, alpha: np.ndarray, x: int, y: int) -> np.ndarray:
-    """The single forward model every stamp uses: ``stamped = (1-a)*bg + a*white``.
-
-    ``alpha`` is float and clipped to [0,1] here, so a caller may pre-multiply it by an
-    opacity factor without worrying about overflow. Shared by ``stamp``/``stamp_slot``
-    here and by ``scripts/detector_response.py`` (which imports it) so the model is
-    written ONCE.
-    """
-    gh, gw = alpha.shape[:2]
-    out = image.copy()
-    roi = out[y : y + gh, x : x + gw].astype(np.float32)
-    a3 = np.clip(alpha, 0.0, 1.0)[..., None]
-    out[y : y + gh, x : x + gw] = np.clip(roi * (1 - a3) + 255.0 * a3, 0, 255).astype(np.uint8)
-    return out
-
-
-def stamp(
-    image: np.ndarray, mark_key: str, *, size_mult: float = 1.0, alpha_mult: float = 1.0
-) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
-    """Composite a mark's alpha glyph into its canonical corner. Returns (stamped, bbox).
-
-    ``size_mult`` scales the glyph box off the engine's nominal geometry and ``alpha_mult``
-    its opacity -- both default to 1.0 (the geometry the engine assumes), and
-    `scripts/detector_response.py` sweeps them to build response curves.
-    """
-    from remove_ai_watermarks._text_mark_engine import load_alpha_template
-
-    engine = engine_for(mark_key)
-    cfg = engine.config
-    alpha = load_alpha_template(cfg.asset_name)
-    if alpha is None:
-        return None
-
-    loc = engine.locate(image)
-    base = engine.scale_base(image)
-    gw = max(cfg.min_gw, int(cfg.alpha_width_frac * base * size_mult))
-    gh = max(4, int(cfg.alpha_height_frac * base * size_mult))
-    if gw < 8 or gh < 4 or gw > loc.w or gh > loc.h:
-        return None
-
-    a = cv2.resize(alpha, (gw, gh), interpolation=cv2.INTER_AREA).astype(np.float32) * alpha_mult
-    x = loc.x + (loc.w - gw) // 2
-    y = loc.y + (loc.h - gh) // 2
-    return composite(image, a, x, y), (x, y, gw, gh)
-
-
-def _slot_alpha(mark_key: str) -> np.ndarray | None:
-    """The alpha source for a mark that ships no bundled alpha PNG."""
-    if mark_key == "gemini":
-        from remove_ai_watermarks.gemini_engine import _shared_engine
-
-        # _shared_engine is the package's own lru_cache singleton: constructing
-        # GeminiEngine() per image reloads its captures, recomputes both alpha maps and
-        # rebuilds the whole template ladder -- the exact cost that singleton exists to
-        # avoid. Private attribute access is deliberate (no accessor exists); a
-        # measurement script may reach in, product code must not.
-        return np.asarray(_shared_engine()._alpha_large, dtype=np.float32)
-    if mark_key == "jimeng_pill":
-        from remove_ai_watermarks._text_mark_engine import load_alpha_template
-
-        return load_alpha_template("jimeng_pill.png")
-    return None
-
-
-def stamp_slot(
-    image: np.ndarray, mark_key: str, *, size_mult: float = 1.0, alpha_mult: float = 1.0
-) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
-    """Stamp a mark into its OWN default footprint slot (the `--no-detect` geometry).
-
-    Used for gemini and the pill, which have no bundled alpha asset with corner
-    fractions. The mark is fitted into the middle of its slot so the fill has to
-    recover the same kind of region it would in production. ``size_mult``/``alpha_mult``
-    match ``stamp`` (both default to 1.0).
-    """
-    from remove_ai_watermarks.watermark_registry import get_mark
-
-    alpha = _slot_alpha(mark_key)
-    if alpha is None:
-        return None
-    loc = get_mark(mark_key).localize(image, force=True)
-    if loc.mask is None:
-        return None
-    ys, xs = np.nonzero(loc.mask)
-    if ys.size == 0:
-        return None
-    y0, y1, x0, x1 = int(ys.min()), int(ys.max()), int(xs.min()), int(xs.max())
-    bw, bh = x1 - x0 + 1, y1 - y0 + 1
-    gw, gh = max(8, int(bw * 0.7 * size_mult)), max(8, int(bh * 0.7 * size_mult))
-    a = cv2.resize(alpha, (gw, gh), interpolation=cv2.INTER_AREA).astype(np.float32) * alpha_mult
-    x, y = x0 + (bw - gw) // 2, y0 + (bh - gh) // 2
-    # A large size_mult can grow the glyph past its slot and push x/y negative; guard so
-    # composite never writes out of bounds (numpy would wrap a negative index silently).
-    if x < 0 or y < 0 or x + gw > image.shape[1] or y + gh > image.shape[0]:
-        return None
-    return composite(image, a, x, y), (x, y, gw, gh)
-
-
 def stamp_any(
     image: np.ndarray, mark_key: str, *, size_mult: float = 1.0, alpha_mult: float = 1.0
 ) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
-    fn = stamp if mark_key in STAMPABLE else stamp_slot
-    return fn(image, mark_key, size_mult=size_mult, alpha_mult=alpha_mult)
+    """Stamp via the same forward model that produces the verified gallery."""
+    from render_visible_examples import stamp_image_mark
+
+    return stamp_image_mark(mark_key, image, size_mult=size_mult, alpha_mult=alpha_mult)
 
 
 def clean_sources(n: int, seed: int = 11) -> list[Path]:
@@ -241,7 +113,7 @@ def main() -> None:
 
     from remove_ai_watermarks.image_io import imread
     from remove_ai_watermarks.region_eraser import lama_available, migan_available
-    from remove_ai_watermarks.watermark_registry import fill, get_mark
+    from remove_ai_watermarks.watermark_registry import fill, get_mark, mark_keys
 
     backends = ["cv2"] + (["migan"] if migan_available() else []) + (["lama"] if lama_available() else [])
     print(f"backends under test: {backends}")
@@ -255,7 +127,7 @@ def main() -> None:
             base = imread(str(src))
             if base is None:
                 continue
-            for key in (*STAMPABLE, *SLOT_STAMPABLE):
+            for key in mark_keys():
                 st = stamp_any(base, key)
                 if st is None:
                     continue

scripts/render_visible_examples.py

@@ -5,8 +5,9 @@ sample of everything it supports and a canary test can hold both sides to it
 (mark registered without example; engine regressed on its canonical example).
 
 Every example is SYNTHETIC: a deterministic generated base photo with the mark's
-own committed silhouette composited at the engine's measured geometry. The
-silhouettes are repository-owned font renders, not copied vendor raster assets.
+own committed alpha or silhouette composited at the engine's measured geometry.
+Some alpha assets were solved from controlled provider captures; the other shapes
+are repository-owned renders rather than copied vendor rasters.
 
 Regenerate with:
     uv run python scripts/render_visible_examples.py
@@ -16,37 +17,31 @@ from __future__ import annotations
 
 import sys
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 import cv2
 import numpy as np
 
-if TYPE_CHECKING:
-    from collections.abc import Callable
-
 _ROOT = Path(__file__).resolve().parents[1]
 sys.path.insert(0, str(_ROOT / "src"))
 
 from remove_ai_watermarks import watermark_registry as wr  # noqa: E402
 from remove_ai_watermarks.image_io import imread  # noqa: E402
 
-ASSETS = _ROOT / "src" / "remove_ai_watermarks" / "assets"
 OUT = _ROOT / "data" / "fixtures" / "visible"
 
 # Composite strength per key: glyph target luma for the light-overlay class.
-# Kling is a thin light-gray run (not near-white); Samsung is a faint overlay
-# expressed by SCALING ITS ALPHA to 0.38 toward full white (see below); everything
-# else is the bold near-white house style.
+# Kling is a thin light-gray run. Samsung's solved asset already stores its faint
+# measured opacity and was solved against pure white, so it must use 255 without a
+# second attenuation. Everything else uses the bold near-white house style.
 _STRENGTH: dict[str, int] = {
     "kling": 208,
+    "samsung": 255,
 }
 _DEFAULT_STRENGTH = 238
 
 # Base geometry: one canonical size per mark where the engine's size modes or
 # calibrated orientation matter, otherwise a plain 3:2 landscape frame.
-# Samsung keeps the larger base: its overlay is faint (peak alpha ~0.38) and
-# the real marks live on ~2958px phone photos -- at 1536 the example falls to 0.39,
-# just under the engine's 0.40 gate.
+# Samsung keeps the larger base because real marks live on ~2958px phone photos.
 _SIZE: dict[str, tuple[int, int]] = {
     "qwen": (1536, 1536),
     "liblib": (1152, 1536),
@@ -73,22 +68,31 @@ def base_photo(w: int, h: int, seed: int = 7) -> np.ndarray:
 
 
 def _glyph_asset(name: str) -> np.ndarray:
-    at = imread(str(ASSETS / name), cv2.IMREAD_GRAYSCALE)
+    from remove_ai_watermarks._text_mark_engine import load_alpha_template
+
+    at = load_alpha_template(name)
     if at is None:
         raise RuntimeError(f"missing silhouette asset: {name}")
-    return at.astype(np.float32) / 255.0
+    return at
 
 
 def _composite_light(base: np.ndarray, alpha: np.ndarray, x: int, y: int, strength: int) -> np.ndarray:
     out = base.copy()
     h, w = alpha.shape[:2]
     roi = out[y : y + h, x : x + w].astype(np.float32)
+    alpha = np.clip(alpha, 0.0, 1.0)
     a3 = alpha[:, :, None] if alpha.ndim == 2 else alpha
     out[y : y + h, x : x + w] = np.clip(roi * (1 - a3) + strength * a3, 0, 255).astype(np.uint8)
     return out
 
 
-def _text_mark_example(key: str, base: np.ndarray) -> np.ndarray:
+def _stamp_text_mark(
+    key: str,
+    base: np.ndarray,
+    *,
+    size_mult: float,
+    alpha_mult: float,
+) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
     engine = wr._engine(key)  # the generator drives the engine's own config
     cfg = engine.config
     h, w = base.shape[:2]
@@ -96,22 +100,27 @@ def _text_mark_example(key: str, base: np.ndarray) -> np.ndarray:
         # Opaque white pill with dark text/sparkle holes, at the measured inset.
         at = _glyph_asset("microsoft_alpha.png")
         long_side = max(w, h)
-        pw = int(0.152 * long_side)
+        pw = int(0.152 * long_side * size_mult)
         ph = max(4, int(pw / (at.shape[1] / at.shape[0])))
         pad = int(0.010 * long_side)
+        if pw > w - pad or ph > h - pad:
+            return None
         pill = cv2.resize(at, (pw, ph))
         x, y = w - pad - pw, pad
         out = base.copy()
         bright = (pill > 0.6)[:, :, None]
-        out[y : y + ph, x : x + pw] = np.where(bright, 245.0, 46.0).astype(np.uint8)
-        return out
+        target = np.where(bright, 245.0, 46.0).astype(np.float32)
+        blend = float(np.clip(alpha_mult, 0.0, 1.0))
+        roi = out[y : y + ph, x : x + pw].astype(np.float32)
+        out[y : y + ph, x : x + pw] = np.clip(roi * (1 - blend) + target * blend, 0, 255).astype(np.uint8)
+        return out, (x, y, pw, ph)
     base_dim = {"short": min(w, h), "width": w, "long": max(w, h)}[cfg.scale_basis]
     # Size the glyph ON a ladder rung: the continuous front ends sweep only the
     # configured rungs, and a glyph sized between rungs collapses the NCC (the
     # comb-collapse qwen's own two-rung ladder exists to avoid).
     rung = max(cfg.ladder) if cfg.detect_frontend != "binary" else 1.0
-    gw = int(cfg.alpha_width_frac * base_dim * rung)
-    gh = max(4, int(cfg.alpha_height_frac * base_dim * rung))
+    gw = max(cfg.min_gw, int(cfg.alpha_width_frac * base_dim * rung * size_mult))
+    gh = max(4, int(cfg.alpha_height_frac * base_dim * rung * size_mult))
     loc = engine.locate(base)
     # Corner-hugging placement: the yuanbao/runninghub anchor gates demote a match
     # that does not hug the corner, and the real marks sit flush on the box's
@@ -124,39 +133,67 @@ def _text_mark_example(key: str, base: np.ndarray) -> np.ndarray:
         x = loc.x
     y = loc.y if cfg.corner in ("tl", "tr") else loc.y + loc.h - gh
     x, y = max(0, x), max(0, y)
+    if x + gw > w or y + gh > h:
+        return None
     at = _glyph_asset(f"{key}_alpha.png")
     alpha = cv2.resize(at, (gw, gh))
-    if key == "samsung":  # faint overlay: peak alpha 0.38 toward FULL white
-        alpha = alpha * 0.38
-    return _composite_light(base, alpha, x, y, _STRENGTH.get(key, _DEFAULT_STRENGTH))
+    alpha = alpha * alpha_mult
+    return _composite_light(base, alpha, x, y, _STRENGTH.get(key, _DEFAULT_STRENGTH)), (x, y, gw, gh)
 
 
-def _gemini_example(base: np.ndarray) -> np.ndarray:
-    from remove_ai_watermarks.gemini_engine import GeminiEngine, get_watermark_config, get_watermark_size
+def _stamp_gemini(
+    base: np.ndarray,
+    *,
+    size_mult: float,
+    alpha_mult: float,
+) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
+    from remove_ai_watermarks.gemini_engine import _shared_engine, get_watermark_config, get_watermark_size
 
     h, w = base.shape[:2]
-    eng = GeminiEngine()
+    eng = _shared_engine()
     size = get_watermark_size(w, h)
     alpha = eng.get_alpha_map(size)
     cfg = get_watermark_config(w, h)
     x, y = cfg.get_position(w, h)
-    return _composite_light(base, alpha.astype(np.float32), x, y, 255)
+    ah, aw = alpha.shape[:2]
+    gw, gh = max(8, int(aw * size_mult)), max(8, int(ah * size_mult))
+    x, y = x + aw - gw, y + ah - gh
+    if x < 0 or y < 0 or x + gw > w or y + gh > h:
+        return None
+    resized = cv2.resize(alpha, (gw, gh), interpolation=cv2.INTER_AREA).astype(np.float32) * alpha_mult
+    return _composite_light(base, resized, x, y, 255), (x, y, gw, gh)
 
 
-def _pill_example(base: np.ndarray) -> np.ndarray:
+def _stamp_pill(
+    base: np.ndarray,
+    *,
+    size_mult: float,
+    alpha_mult: float,
+) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
     h, w = base.shape[:2]
     at = _glyph_asset("jimeng_pill.png")
-    pw = max(24, int(0.161 * w))
+    pw = max(24, int(0.161 * w * size_mult))
     ph = max(8, int(pw * at.shape[0] / at.shape[1]))
     x, y = int(0.03 * w), int(0.03 * h)
-    alpha = cv2.resize(at, (pw, ph))
-    return _composite_light(base, alpha, x, y, 232)
+    if x + pw > w or y + ph > h:
+        return None
+    alpha = cv2.resize(at, (pw, ph)) * alpha_mult
+    return _composite_light(base, alpha, x, y, 232), (x, y, pw, ph)
 
 
-_BUILDERS: dict[str, Callable[[np.ndarray], np.ndarray]] = {
-    "gemini": _gemini_example,
-    "jimeng_pill": _pill_example,
-}
+def stamp_image_mark(
+    key: str,
+    base: np.ndarray,
+    *,
+    size_mult: float = 1.0,
+    alpha_mult: float = 1.0,
+) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
+    """Composite one image mark using the gallery's verified forward model."""
+    if key == "gemini":
+        return _stamp_gemini(base, size_mult=size_mult, alpha_mult=alpha_mult)
+    if key == "jimeng_pill":
+        return _stamp_pill(base, size_mult=size_mult, alpha_mult=alpha_mult)
+    return _stamp_text_mark(key, base, size_mult=size_mult, alpha_mult=alpha_mult)
 
 
 def build(
@@ -177,8 +214,10 @@ def build_pair(
     """Return the clean synthetic source and the matching marked image."""
     size = size or _SIZE.get(key, (1536, 1152))
     clean = base_photo(*size, seed=seed)
-    builder = _BUILDERS.get(key)
-    marked = builder(clean) if builder is not None else _text_mark_example(key, clean)
+    stamped = stamp_image_mark(key, clean)
+    if stamped is None:
+        raise RuntimeError(f"could not stamp {key} at its nominal geometry")
+    marked, _box = stamped
     return clean, marked