remove-ai-watermarks/scripts/fill_quality.py

"""Tier B1: measure visible-mark FILL quality against a constructed ground truth.

THE PROBLEM THIS SOLVES
  A real watermarked image has no clean counterpart, so fill quality has only ever been
  eyeballed ("cv2 smears on texture, LaMa is best") and never carried a number. The docs
  state a preference order for `--backend auto` that rests on no measurement.

THE CONSTRUCTION
  Take a VERIFIED-CLEAN corpus image, stamp a known mark onto it using the mark's own
  captured alpha map (the forward model the reverse-alpha work established:
  `stamped = (1-a)*bg + a*white`), then remove it and compare against the TRUE original.
  The original is the answer key that reality never provides.

TWO MEASUREMENT DECISIONS THAT MAKE OR BREAK THIS
  * Score INSIDE the footprint only. The fill touches a tiny corner, so whole-frame PSNR
    sits near 60 dB whatever the backend does and would rank them all "excellent".
  * Report the DAMAGE baseline (stamped vs original) next to the recovery (filled vs
    original). Without it a PSNR is uninterpretable: 30 dB is a triumph if the mark cost
    12 dB and a failure if it cost 29 dB. The honest metric is how much of the gap the
    fill closes.

WHAT IT DOES NOT MEASURE
  Detection. The mark is localized with `force=True`, so a miss cannot contaminate the
  fill numbers -- this isolates the FILL. Detection accuracy is Tier C's job.

DATA SAFETY
  Corpus images are user uploads: read-only, local analysis, output under a gitignored
  data/spaces/ path. No image content is written into the report.

    uv run python scripts/fill_quality.py --n 60
"""

from __future__ import annotations

import argparse
import collections
import glob
import json
import random
import sys
from pathlib import Path
from typing import Any

import cv2
import numpy as np

sys.path.insert(0, str(Path(__file__).parent.parent))
sys.path.insert(0, str(Path(__file__).parent))

from invisible_quality_audit import _ssim  # reuse, do not reimplement a third SSIM

REPO = Path(__file__).resolve().parents[1]
CORPUS = REPO / "data" / "spaces" / "originals"
OUT = REPO / "data" / "spaces" / "_fill_quality.jsonl"

# Text marks: a bundled alpha PNG plus the engine's own corner geometry.
STAMPABLE = ("doubao", "jimeng", "samsung")
# 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))
    return float("inf") if mse == 0 else 10 * float(np.log10(255.0**2 / mse))


def texture_of(box: np.ndarray) -> float:
    """Median Sobel magnitude -- the same texture proxy the pill's flatness gate uses."""
    g = cv2.cvtColor(box, cv2.COLOR_BGR2GRAY) if box.ndim == 3 else box
    gx = cv2.Sobel(g, cv2.CV_32F, 1, 0, ksize=3)
    gy = cv2.Sobel(g, cv2.CV_32F, 0, 1, ksize=3)
    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)


def clean_sources(n: int, seed: int = 11) -> list[Path]:
    """Corpus images with NO metadata signal and NO mark detection -- verified clean."""
    from remove_ai_watermarks.identify import identify
    from remove_ai_watermarks.image_io import imread
    from remove_ai_watermarks.watermark_registry import detect_marks

    pool = glob.glob(str(CORPUS / "*" / "*"))
    random.Random(seed).shuffle(pool)  # noqa: S311 -- deterministic sampling, not crypto
    out: list[Path] = []
    for p in pool:
        if len(out) >= n:
            break
        path = Path(p)
        try:
            img = imread(str(path))
            if img is None or min(img.shape[:2]) < 400:
                continue
            # check_visible=False: identify would otherwise decode the file a second
            # time and run the very same visible detectors as detect_marks below.
            if identify(path, check_visible=False).signals:
                continue
            if any(d.detected for d in detect_marks(img)):
                continue
        except Exception:  # noqa: S112 -- a bad corpus file just is not a candidate
            continue
        out.append(path)
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=60, help="clean source images")
    ap.add_argument("--out", type=Path, default=OUT)
    a = ap.parse_args()

    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

    backends = ["cv2"] + (["migan"] if migan_available() else []) + (["lama"] if lama_available() else [])
    print(f"backends under test: {backends}")
    print(f"selecting {a.n} verified-clean sources...", flush=True)
    sources = clean_sources(a.n)
    print(f"got {len(sources)}\n", flush=True)

    rows: list[dict[str, Any]] = []
    with open(a.out, "w", encoding="utf-8") as fh:
        for i, src in enumerate(sources, 1):
            base = imread(str(src))
            if base is None:
                continue
            for key in (*STAMPABLE, *SLOT_STAMPABLE):
                st = stamp_any(base, key)
                if st is None:
                    continue
                stamped, (x, y, w, h) = st
                # Score a slightly padded box: the fill dilates, so a tight glyph box
                # would miss damage the backend does just outside the glyph.
                pad = 8
                y0, y1 = max(0, y - pad), min(base.shape[0], y + h + pad)
                x0, x1 = max(0, x - pad), min(base.shape[1], x + w + pad)
                truth = base[y0:y1, x0:x1]
                dmg_psnr = psnr(stamped[y0:y1, x0:x1], truth)
                dmg_ssim = _ssim(
                    cv2.cvtColor(stamped[y0:y1, x0:x1], cv2.COLOR_BGR2GRAY), cv2.cvtColor(truth, cv2.COLOR_BGR2GRAY)
                )
                tex = texture_of(truth)

                mark = get_mark(key)
                loc = mark.localize(stamped, force=True)
                if loc.mask is None:
                    continue
                for backend in backends:
                    try:
                        filled = fill(stamped, loc.mask, backend=backend)
                    except Exception as e:
                        rows.append({"src": src.name, "mark": key, "backend": backend, "error": str(e)[:150]})
                        continue
                    # Invariant: the fill must not touch anything outside its mask.
                    outside = loc.mask == 0
                    untouched = bool(np.array_equal(filled[outside], stamped[outside]))
                    rec = {
                        "src": src.name,
                        "mark": key,
                        "backend": backend,
                        "texture": round(tex, 2),
                        "damage_psnr": round(dmg_psnr, 2),
                        "damage_ssim": round(dmg_ssim, 4),
                        "filled_psnr": round(psnr(filled[y0:y1, x0:x1], truth), 2),
                        "filled_ssim": round(
                            _ssim(
                                cv2.cvtColor(filled[y0:y1, x0:x1], cv2.COLOR_BGR2GRAY),
                                cv2.cvtColor(truth, cv2.COLOR_BGR2GRAY),
                            ),
                            4,
                        ),
                        "outside_mask_untouched": untouched,
                    }
                    rows.append(rec)
                    fh.write(json.dumps(rec) + "\n")
            if i % 10 == 0:
                fh.flush()
                print(f"  {i}/{len(sources)}", flush=True)

    report(rows, a.out)


def report(rows: list[dict[str, Any]], out: Path) -> None:
    good = [r for r in rows if "error" not in r]
    if not good:
        print("no measurements")
        return
    tex = sorted(r["texture"] for r in good)
    t1, t2 = tex[len(tex) // 3], tex[2 * len(tex) // 3]

    def bucket(t: float) -> str:
        return "flat" if t <= t1 else ("mid" if t <= t2 else "textured")

    print(f"\n{'=' * 78}\nFILL QUALITY  n={len(good)} measurements  (texture terciles at {t1:.1f} / {t2:.1f})")
    print(f"{'=' * 78}")
    print("recovered = filled PSNR minus damaged PSNR; how much of the mark's damage the fill undoes\n")
    hdr = f"{'mark':9s} {'backend':8s} {'bg':9s} {'n':>4s} {'damaged':>9s} {'filled':>9s}"
    print(f"{hdr} {'recovered':>10s} {'ssim':>7s}")
    agg: dict[tuple[str, str, str], list[dict[str, Any]]] = collections.defaultdict(list)
    for r in good:
        agg[(r["mark"], r["backend"], bucket(r["texture"]))].append(r)
    for (mark, backend, bg), rs in sorted(agg.items()):
        # MEDIAN, not mean: a fill that reproduces a flat background EXACTLY scores
        # PSNR=inf, and a single inf makes the mean inf -- which reported "+inf" for
        # every flat bucket on the first run and hid the real numbers.
        finite = [r for r in rs if np.isfinite(r["filled_psnr"])]
        perfect = len(rs) - len(finite)
        d = float(np.median([r["damage_psnr"] for r in rs]))
        f = float(np.median([r["filled_psnr"] for r in finite])) if finite else float("inf")
        sv = float(np.median([r["filled_ssim"] for r in rs]))
        # Recovery is the MEDIAN OF THE PER-IMAGE DIFFERENCES, not the difference of the
        # two medians. Those are not the same statistic on skewed data and they can even
        # disagree in SIGN: the first version of this report subtracted medians and
        # produced a backend ranking that a paired comparison did not support.
        rec = float(np.median([r["filled_psnr"] - r["damage_psnr"] for r in finite])) if finite else float("inf")
        tail = f"  ({perfect} exact)" if perfect else ""
        print(f"{mark:9s} {backend:8s} {bg:9s} {len(rs):4d} {d:9.2f} {f:9.2f} {rec:+10.2f} {sv:7.4f}{tail}")

    bad = [r for r in good if not r["outside_mask_untouched"]]
    print(f"\noutside-mask invariant violations: {len(bad)}" + (" <-- BUG" if bad else " (none)"))
    print(f"records: {out}")


if __name__ == "__main__":
    main()