From d391cc42d1435dfd930d7388b5790039672c5d13 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Thu, 3 Sep 2026 11:05:59 -0700 Subject: [PATCH] Add post-fill validation for visible watermarks --- README.md | 3 + docs/cli.md | 7 + docs/known-limitations.md | 14 ++ docs/module-internals.md | 33 ++- docs/python-api.md | 19 +- src/remove_ai_watermarks/__init__.py | 7 + src/remove_ai_watermarks/api.py | 53 ++++- src/remove_ai_watermarks/cli.py | 16 +- .../watermark_registry.py | 204 ++++++++++++++++-- tests/test_api.py | 20 ++ tests/test_cli.py | 36 ++++ tests/test_watermark_registry.py | 115 ++++++++++ 12 files changed, 498 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 4f1133a..ee16b0a 100644 --- a/README.md +++ b/README.md @@ -401,6 +401,9 @@ import remove_ai_watermarks as raiw result, removed = raiw.remove_visible("watermarked.png", "clean.png") print(removed) +report = raiw.remove_visible_detailed("watermarked.png", "clean.png") +print(report.status) # cleaned | partial | unvalidated | no_watermark + openai_synthid = raiw.verify_openai_synthid("image.png", acknowledge_upload=True) print(openai_synthid.status) diff --git a/docs/cli.md b/docs/cli.md index ab059f7..a924686 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -151,8 +151,15 @@ The default behavior: - removes every detected match, except the weakly detected Jimeng label pill, which needs corroboration (see [supported signals](supported-signals.md)); - selects the best installed fill backend; +- checks the filled region once without changing it or retrying the fill; - strips AI metadata from the output. +The automatic command reports `Removed and validated` when the detector no +longer accepts an overlapping mark, `Partial` when an overlapping residual is +still detected, and `Post-removal validation unavailable` when the check fails. +The latter two still write the once-filled output and return success; they are +result-quality statuses, not requests for an automatic second edit. + Use a specific mark: ```bash diff --git a/docs/known-limitations.md b/docs/known-limitations.md index dbd4265..fac0ef4 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -36,6 +36,20 @@ have to be reconstructed. No backend can recover detail that is completely hidden by an opaque mark. A successful detection therefore does not guarantee a visually perfect fill. +### Post-fill validation is not a visual-quality oracle + +Automatic visible removal checks each result once with the same detector and +reports `cleaned`, `partial`, or `unvalidated`. `partial` means an accepted +detection still overlaps the region that was filled. This can catch a residual, +but it cannot prove that reconstructed texture, text, or structure looks natural: +the validator measures mark evidence, not perceptual quality, and it is not +independent from the detector that selected the mark. + +The check is deliberately read-only. It does not enlarge the mask or trigger a +second fill. Mask quality and backend choice must be improved before the single +fill; a post-fill retry can compound blur or invented detail without removing +the detector signal. + ### Automatic detection covers registered variants only The registry contains vendor and locale specific templates. A redesigned mark, diff --git a/docs/module-internals.md b/docs/module-internals.md index 7b17202..804e6df 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -89,6 +89,8 @@ Regression coverage: points: - `remove_visible` +- `remove_visible_detailed`, returning a `VisibleRemovalResult` with per-mark + post-fill validation - `visible_provenance` and the image pipeline that the `all` and `batch` commands are thin wrappers @@ -121,9 +123,9 @@ The package root exposes all of them lazily through [`__getattr__`](../src/remove_ai_watermarks/__init__.py), keeping a plain package import free of the heavier image and model imports. -For path inputs, `remove_visible` reads provenance metadata, preserves alpha, -and optionally writes and strips metadata. Array inputs are treated as BGR -arrays and have no file provenance or separate alpha plane. +For path inputs, both visible-removal entry points read provenance metadata, +preserve alpha, and optionally write and strip metadata. Array inputs are +treated as BGR arrays and have no file provenance or separate alpha plane. When no visible mark is removed, a same-format path copy preserves the original bytes. `write_noop=False` leaves the requested output path untouched instead. @@ -1173,6 +1175,31 @@ absent or unmapped producer still falls back to the ByteDance pair. `remove_auto_marks` removes every selected mark, not only the strongest one. This matters for images that carry marks in more than one corner. +The automatic path has four stages: perception, decision, action, then read-only +validation. The first action reuses the exact detection produced by perception, +so the common one-mark path pays for one pre-fill and one post-fill detector pass +rather than two pre-fill passes plus validation. Later co-firing marks retain the +historical re-detection on the progressively edited image instead of acting on +stale geometry. + +After one fill, the same detector runs at the same resolved trust level. A +post-fill detection counts as a residual only if its region overlaps the actual +mask bounding box. The per-mark status is `cleaned`, `partial`, or `unvalidated`; +the aggregate also has `no_watermark`. Validation never expands the mask or +retries the fill. Because it uses the same detector, this is a consistency check, +not an independent oracle of visual quality. + +Runtime was measured on 2026-09-03 with 12 tracked public visible fixtures, five +warmed repetitions each, explicit `cv2`, alternating old/new order, and no +concurrent benchmark load. Across 60 paired runs, the old path's median was +373.9 ms and the validated path's median was 379.2 ms; the median paired delta +was -0.9% (interquartile range -3.6% to +1.5%) and the aggregate delta was +-1.5%. All output pixels and removed-label lists matched the old path. These +figures show no measurable regression on that corpus, not a cross-machine +latency guarantee. A separate exact snapshot check kept every strict/relaxed +verdict, confidence, and region unchanged for all 12 detectors on all 12 +fixtures. + Regression coverage: - [`test_watermark_registry.py`](../tests/test_watermark_registry.py) diff --git a/docs/python-api.md b/docs/python-api.md index d9241c5..c84aa89 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -50,9 +50,26 @@ The function returns: An empty `removed` list means that no registered visible mark was selected. It does not prove the image has no metadata or invisible watermark. +Use the detailed form when the caller needs to distinguish a validated fill +from a residual or a validator failure: + +```python +report = raiw.remove_visible_detailed("watermarked.png", "clean.png") +print(report.status) # "no_watermark" | "cleaned" | "partial" | "unvalidated" +for mark in report.marks: + print(mark.label, mark.status, mark.confidence_before, mark.confidence_after) +``` + +Each selected mark is localized and filled once. The same detector then checks +the result without changing the mask or running another fill. `partial` means +that detector still accepts a region overlapping the mask that was filled; +`unvalidated` means the check failed, not that the mark remains. The legacy +`remove_visible` tuple API performs the same check but intentionally discards +the detailed status. + ### Path input -For a path input, `remove_visible`: +For a path input, `remove_visible` and `remove_visible_detailed`: - reads metadata provenance for the default `auto` sensitivity; - preserves a separate alpha channel; diff --git a/src/remove_ai_watermarks/__init__.py b/src/remove_ai_watermarks/__init__.py index b4320ef..b966573 100644 --- a/src/remove_ai_watermarks/__init__.py +++ b/src/remove_ai_watermarks/__init__.py @@ -5,6 +5,7 @@ High-level API (lazy, so ``import remove_ai_watermarks`` stays cheap):: import remove_ai_watermarks as raiw raiw.remove_visible("in.png", "out.png") # clean a file (provenance auto) result, removed = raiw.remove_visible(bgr_array) # array -> array + report = raiw.remove_visible_detailed(bgr_array) # cleaned / partial / unvalidated raiw.visible_provenance("in.png") # -> frozenset of confirmed vendors raiw.identify_video("in.mp4") # -> VideoProvenanceReport raiw.inspect_video_metadata("in.mp4") # -> VideoMetadataReport @@ -42,6 +43,7 @@ __all__ = [ "OpenAIProvenanceError", "OpenAISynthIDDetection", "RemoveAllResult", + "VisibleRemovalResult", "__version__", "identify_video", "inspect_video_metadata", @@ -53,6 +55,7 @@ __all__ = [ "remove_video_metadata", "remove_video_visible", "remove_visible", + "remove_visible_detailed", "verify_openai_synthid", "visible_provenance", ] @@ -63,9 +66,11 @@ if TYPE_CHECKING: InvisibleOptions, MetadataStripIncomplete, RemoveAllResult, + VisibleRemovalResult, remove_all, remove_batch, remove_visible, + remove_visible_detailed, visible_provenance, ) from remove_ai_watermarks.openai_provenance import ( @@ -92,9 +97,11 @@ def __getattr__(name: str) -> object: "InvisibleOptions", "MetadataStripIncomplete", "RemoveAllResult", + "VisibleRemovalResult", "remove_all", "remove_batch", "remove_visible", + "remove_visible_detailed", "visible_provenance", ): from remove_ai_watermarks import api diff --git a/src/remove_ai_watermarks/api.py b/src/remove_ai_watermarks/api.py index 7e20f13..ce869d2 100644 --- a/src/remove_ai_watermarks/api.py +++ b/src/remove_ai_watermarks/api.py @@ -8,6 +8,7 @@ up metadata provenance, or preserve the alpha channel by hand: import remove_ai_watermarks as raiw raiw.remove_visible("in.png", "out.png") # path -> file, provenance auto result, removed = raiw.remove_visible(bgr_array) # array -> array + report = raiw.remove_visible_detailed(bgr_array) # includes post-fill status raiw.remove_visible("shot.png", "out.png", sensitivity="strict") raiw.visible_provenance("in.png") # -> frozenset({"gemini"}) @@ -26,7 +27,16 @@ if TYPE_CHECKING: from numpy.typing import NDArray - from remove_ai_watermarks.watermark_registry import Backend, Sensitivity + from remove_ai_watermarks.watermark_registry import Backend, Sensitivity, VisibleRemovalResult + + +def __getattr__(name: str) -> object: + """Lazily expose the detailed visible result type without loading image code.""" + if name == "VisibleRemovalResult": + from remove_ai_watermarks.watermark_registry import VisibleRemovalResult + + return VisibleRemovalResult + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") @dataclass(frozen=True) @@ -182,13 +192,42 @@ def remove_visible( output path untouched, so a caller that treats "no mark" as "produce nothing" (the CLI ``visible`` no-mark contract) does not clobber a pre-existing file at that path. """ + report = remove_visible_detailed( + source, + output, + sensitivity=sensitivity, + backend=backend, + strip_metadata=strip_metadata, + write_noop=write_noop, + ) + return report.image, report.labels + + +def remove_visible_detailed( + source: str | Path | NDArray[Any], + output: str | Path | None = None, + *, + sensitivity: Sensitivity = "auto", + backend: Backend = "auto", + strip_metadata: bool = True, + write_noop: bool = True, +) -> VisibleRemovalResult: + """Remove known visible marks and report the read-only post-fill check. + + Unlike :func:`remove_visible`, this returns one :class:`VisibleRemovalResult`. + Its aggregate ``status`` is ``no_watermark``, ``cleaned``, ``partial``, or + ``unvalidated``. Each mark also records the before/after detector confidence, + filled mask box, backend, and any overlapping residual region. + + The validator never edits the image: every selected mask is filled exactly once, + then the same detector checks the result. All input, output, metadata, alpha, and + no-op semantics otherwise match :func:`remove_visible`. + """ from remove_ai_watermarks import watermark_registry - # Reject a removed sensitivity loudly; `Sensitivity` is a Literal and not enforced - # at runtime, so a 0.15 caller would otherwise get `auto` behavior in silence. watermark_registry.validate_sensitivity(sensitivity) loaded = _load_visible_input(source) - result, removed = watermark_registry.remove_auto_marks( + report = watermark_registry.remove_auto_marks_detailed( loaded.bgr, sensitivity=sensitivity, provenance=loaded.provenance, @@ -197,13 +236,13 @@ def remove_visible( if output is not None: _write_visible_result( loaded, - result, - removed, + report.image, + report.labels, output, strip_metadata=strip_metadata, write_noop=write_noop, ) - return result, removed + return report # ── The three-stage image pipeline (visible -> invisible -> metadata) ── diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index be976d3..8dd2565 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -574,7 +574,7 @@ def _run_visible_auto( t0 = time.monotonic() try: with console.status("Detecting & removing visible marks..."): - result, removed = api.remove_visible( + report = api.remove_visible_detailed( str(source), str(output), sensitivity=sensitivity, @@ -593,13 +593,25 @@ def _run_visible_auto( raise SystemExit(1) from e elapsed = time.monotonic() - t0 + result, removed = report.image, report.labels h, w = result.shape[:2] console.print(f" Input: {source.name} ({w}x{h})") if not removed: # write_noop=False means nothing was written, so a pre-existing output is intact. console.print(f" No known visible mark detected. Checked: {', '.join(watermark_registry.mark_keys())}.") _no_visible_mark_exit(source) - console.print(f" Removed: {', '.join(removed)}") + if report.status == "partial": + residuals = [mark.label for mark in report.marks if mark.status == "partial"] + console.print( + f" Partial: fill completed for {', '.join(removed)}; " + f"overlapping residual still detected for {', '.join(residuals)}." + ) + elif report.status == "unvalidated": + unavailable = [mark.label for mark in report.marks if mark.status == "unvalidated"] + console.print(f" Filled: {', '.join(removed)}") + console.print(f" Post-removal validation unavailable for: {', '.join(unavailable)}.") + else: + console.print(f" Removed and validated: {', '.join(removed)}") size_kb = output.stat().st_size / 1024 console.print(f" Saved: {output} ({size_kb:.0f} KB, {elapsed:.2f}s)") diff --git a/src/remove_ai_watermarks/watermark_registry.py b/src/remove_ai_watermarks/watermark_registry.py index b53ce35..a1d548c 100644 --- a/src/remove_ai_watermarks/watermark_registry.py +++ b/src/remove_ai_watermarks/watermark_registry.py @@ -5,10 +5,11 @@ A single catalog that ties each known visible mark to (a) where it usually sits, registry detects every known mark in its usual place and removes the ones present. -**Localize -> fill.** A known mark is removed by LOCALIZING it (a version-robust +**Localize -> fill -> validate.** A known mark is removed by LOCALIZING it (a version-robust detector that returns a binary footprint MASK) and then handing that mask to ONE shared, swappable fill backend (``region_eraser``: cv2 Telea/NS, -MI-GAN, or big-LaMa). No mark carries a reverse-alpha step any more: the old +MI-GAN, or big-LaMa). The same detector then checks the result without changing +the mask or running a second fill. No mark carries a reverse-alpha step any more: the old ``original = (wm - a*logo)/(1-a)`` recovery depended on a fixed captured alpha map at a fixed position, broke whenever a vendor re-rendered or moved its mark, and was not color-lossless even with the right map (it amplifies quantization/JPEG-chroma @@ -199,6 +200,13 @@ class Candidate: detected_strict: bool detected_relaxed: bool features: dict[str, float] # generic; both construction sites always supply it (empty when none) + # Retain the exact perception outputs for the action stage. The first selected + # mark is still being localized on the unchanged input, so throwing these away + # and re-detecting it paid for a duplicate scan. Reusing that scan makes room for + # the read-only post-removal validation without adding another detector pass to + # the common one-mark path. Optional for policy-only Candidates in tests. + strict_detection: MarkDetection | None = field(default=None, compare=False, repr=False) + relaxed_detection: MarkDetection | None = field(default=None, compare=False, repr=False) @dataclass(frozen=True) @@ -211,6 +219,55 @@ class Decision: relax: bool +VisibleRemovalStatus = Literal["no_watermark", "cleaned", "partial", "unvalidated"] +MarkValidationStatus = Literal["cleaned", "partial", "unvalidated"] + + +@dataclass(frozen=True) +class MarkRemovalResult: + """One visible-mark fill and its read-only post-removal validation. + + ``partial`` means the same detector still accepts an overlapping region after + the one fill. It is useful evidence about this engine, not an independent oracle. + Validation never expands the mask and never runs a second fill. + """ + + key: str + label: str + location: str + region: Region + mask_bbox: Region + backend: Literal["cv2", "migan", "lama"] + confidence_before: float + confidence_after: float | None + status: MarkValidationStatus + residual_region: Region | None = None + + +@dataclass(frozen=True) +class VisibleRemovalResult: + """Detailed result of one automatic visible-mark pass.""" + + image: NDArray[Any] = field(compare=False, repr=False) + marks: tuple[MarkRemovalResult, ...] + + @property + def labels(self) -> list[str]: + """Labels whose masks were filled, preserving the legacy list contract.""" + return [mark.label for mark in self.marks] + + @property + def status(self) -> VisibleRemovalStatus: + """Aggregate status, with a known residual taking precedence over unknown.""" + if not self.marks: + return "no_watermark" + if any(mark.status == "partial" for mark in self.marks): + return "partial" + if any(mark.status == "unvalidated" for mark in self.marks): + return "unvalidated" + return "cleaned" + + @dataclass(frozen=True) class KnownMark: """A known visible watermark: where it lives, how to find and mask it. @@ -882,7 +939,17 @@ def _build_candidates(image: NDArray[Any]) -> list[Candidate]: continue strict, relaxed = m.detect_both(image) feats = m.features(image) if (strict.detected or relaxed.detected) else {} - cands.append(Candidate(m.key, m.label, strict.detected, relaxed.detected, feats)) + cands.append( + Candidate( + m.key, + m.label, + strict.detected, + relaxed.detected, + feats, + strict_detection=strict, + relaxed_detection=relaxed, + ) + ) return cands @@ -913,14 +980,31 @@ def decide(candidates: list[Candidate], context: Context) -> list[Decision]: return fired -def remove_auto_marks( +def _mask_bbox(mask: NDArray[Any]) -> Region | None: + """Smallest ``(x, y, w, h)`` box containing a mask, or None when empty.""" + import cv2 + + x, y, w, h = cv2.boundingRect(mask) + return (x, y, w, h) if w > 0 and h > 0 else None + + +def _regions_overlap(left: Region, right: Region) -> bool: + """Whether two positive-area ``(x, y, w, h)`` regions overlap.""" + lx, ly, lw, lh = left + rx, ry, rw, rh = right + if min(lw, lh, rw, rh) <= 0: + return False + return lx < rx + rw and rx < lx + lw and ly < ry + rh and ry < ly + lh + + +def remove_auto_marks_detailed( image: NDArray[Any], *, sensitivity: Sensitivity = "auto", provenance: frozenset[str] = frozenset(), backend: Backend = "auto", -) -> tuple[NDArray[Any], list[str]]: - """Remove EVERY decided ``in_auto`` mark in one pass, chaining the result. +) -> VisibleRemovalResult: + """Fill every decided mark once, then validate each result without editing it. The three stages are separated: PERCEPTION (:func:`_build_candidates` -- engines report what they see, no policy), DECISION (:func:`decide` -- the pure arbiter over @@ -932,16 +1016,104 @@ def remove_auto_marks( Three orthogonal knobs: ``sensitivity`` (how hard to trust a borderline mark -- see :data:`Sensitivity`), ``provenance`` (vendor keys external metadata confirms, the evidence that drives ``auto``; the TC260 label maps to ``jimeng``/``doubao``), - and ``backend`` (the shared fill). Returns ``(result, [labels removed])``; empty - means nothing fired.""" + and ``backend`` (the shared fill). + + Each successful action has a fourth, read-only stage: the same mark detector runs + once on the filled image. An accepted region counts as a residual only when it + overlaps the mask that was actually filled. The result is ``partial`` in that case, + ``unvalidated`` if the detector fails, and ``cleaned`` otherwise. This is a + self-consistency check, not an independent removal oracle; it never changes the + mask, retries the fill, or hides the candidate image. + + The first selected mark reuses the exact detection from perception because its + input is still unchanged. Later marks retain the historical re-detection on the + progressively modified image, so overlapping/co-firing marks cannot act on stale + evidence.""" context = Context(sensitivity=sensitivity, provenance=provenance) result = image - labels: list[str] = [] + removals: list[MarkRemovalResult] = [] + resolved_backend: Literal["cv2", "migan", "lama"] | None = None for d in decide(_build_candidates(image), context): - result, region = get_mark(d.candidate.key).remove(result, backend=backend, provenance=d.relax, force=False) - # Only report the mark as removed when a fill actually happened: remove() returns - # a None region when the localized mask came back empty, and reporting it anyway - # would claim a removal that left the pixels unchanged. - if region is not None: - labels.append(d.candidate.label) - return result, labels + mark = get_mark(d.candidate.key) + cached_detection = None + if not removals: + cached_detection = d.candidate.relaxed_detection if d.relax else d.candidate.strict_detection + localization = mark.localize( + result, + provenance=d.relax, + force=False, + detection=cached_detection, + ) + mask = localization.mask + if mask is None: + continue + bbox = _mask_bbox(mask) + if bbox is None: + continue + + if resolved_backend is None: + resolved_backend = resolve_backend(backend) + region = localization.region + confidence_before = localization.confidence + result = fill(result, mask, backend=resolved_backend) + del localization, mask + + confidence_after: float | None = None + residual_region: Region | None = None + validation_status: MarkValidationStatus + try: + after = mark.detect(result, provenance=d.relax) + confidence_after = after.confidence + if after.detected and _regions_overlap(after.region, bbox): + validation_status = "partial" + residual_region = after.region + else: + validation_status = "cleaned" + except Exception as exc: + validation_status = "unvalidated" + logger.warning( + "Post-removal validation failed for %s after %s fill in mask %s: %s", + mark.key, + resolved_backend, + bbox, + exc, + ) + + removals.append( + MarkRemovalResult( + key=mark.key, + label=mark.label, + location=mark.location, + region=region, + mask_bbox=bbox, + backend=resolved_backend, + confidence_before=confidence_before, + confidence_after=confidence_after, + status=validation_status, + residual_region=residual_region, + ) + ) + + return VisibleRemovalResult(result, tuple(removals)) + + +def remove_auto_marks( + image: NDArray[Any], + *, + sensitivity: Sensitivity = "auto", + provenance: frozenset[str] = frozenset(), + backend: Backend = "auto", +) -> tuple[NDArray[Any], list[str]]: + """Compatibility wrapper returning the filled image and removed labels. + + The action still performs the read-only post-removal check. Call + :func:`remove_auto_marks_detailed` when its ``cleaned`` / ``partial`` / + ``unvalidated`` result is needed. + """ + report = remove_auto_marks_detailed( + image, + sensitivity=sensitivity, + provenance=provenance, + backend=backend, + ) + return report.image, report.labels diff --git a/tests/test_api.py b/tests/test_api.py index 5cbfca4..5faccbb 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -21,6 +21,8 @@ class TestTopLevelExports: from remove_ai_watermarks import openai_provenance assert raiw.remove_visible is api.remove_visible + assert raiw.remove_visible_detailed is api.remove_visible_detailed + assert raiw.VisibleRemovalResult is api.VisibleRemovalResult assert raiw.visible_provenance is api.visible_provenance assert raiw.verify_openai_synthid is openai_provenance.verify_openai_synthid assert raiw.OpenAIProvenanceError is openai_provenance.OpenAIProvenanceError @@ -64,6 +66,16 @@ class TestRemoveVisibleArray: assert removed == [] assert result.shape == arr.shape + def test_detailed_no_mark_result_is_explicit(self): + arr = np.zeros((256, 256, 3), np.uint8) + + report = raiw.remove_visible_detailed(arr, backend="cv2") + + assert report.status == "no_watermark" + assert report.labels == [] + assert report.marks == () + assert np.array_equal(report.image, arr) + def test_bad_source_raises(self, tmp_path): with pytest.raises(ValueError, match="Could not read image"): raiw.remove_visible(tmp_path / "nope.png") @@ -83,6 +95,14 @@ class TestRemoveVisiblePath: result, _ = raiw.remove_visible(DOUBAO, backend="cv2") assert result.ndim == 3 + def test_detailed_result_reports_post_removal_status(self): + report = raiw.remove_visible_detailed(DOUBAO, backend="cv2") + + assert report.status == "cleaned" + assert any("Doubao" in label for label in report.labels) + assert len(report.marks) == 1 + assert report.marks[0].confidence_after is not None + class TestNoOpPreservesOriginal: def test_no_mark_copies_original_bytes(self, tmp_path): diff --git a/tests/test_cli.py b/tests/test_cli.py index ad1f00a..97f818b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -156,6 +156,42 @@ class TestVisibleCommand: assert not output.exists() assert "all" in result.output + @pytest.mark.parametrize( + ("status", "expected"), + [ + ("partial", "overlapping residual still detected"), + ("unvalidated", "Post-removal validation unavailable"), + ], + ) + def test_visible_auto_reports_non_clean_validation_status(self, runner, sample_png, tmp_path, status, expected): + from remove_ai_watermarks.watermark_registry import MarkRemovalResult, VisibleRemovalResult + + output = tmp_path / "clean.png" + image = cv2.imread(str(sample_png), cv2.IMREAD_COLOR) + mark = MarkRemovalResult( + key="gemini", + label="Gemini visible watermark", + location="bottom-right", + region=(10, 10, 20, 20), + mask_bbox=(10, 10, 20, 20), + backend="cv2", + confidence_before=0.9, + confidence_after=0.6 if status == "partial" else None, + status=status, + residual_region=(12, 12, 10, 10) if status == "partial" else None, + ) + + def _fake_detailed(_source, destination, **_kwargs): + assert cv2.imwrite(str(destination), image) + return VisibleRemovalResult(image, (mark,)) + + with patch("remove_ai_watermarks.api.remove_visible_detailed", side_effect=_fake_detailed): + result = runner.invoke(main, ["visible", str(sample_png), "-o", str(output)]) + + assert result.exit_code == 0, result.output + assert expected in result.output + assert "Removed: Gemini visible watermark" not in result.output + def test_visible_basic(self, runner, sample_png, tmp_path): output = tmp_path / "clean.png" result = runner.invoke( diff --git a/tests/test_watermark_registry.py b/tests/test_watermark_registry.py index 67ce8c4..7c4e737 100644 --- a/tests/test_watermark_registry.py +++ b/tests/test_watermark_registry.py @@ -185,6 +185,121 @@ class TestFill: assert np.array_equal(out, img) +class TestPostRemovalValidation: + """A post-check reports the first fill; it never edits the image again.""" + + @staticmethod + def _install_single_mark( + monkeypatch: pytest.MonkeyPatch, + *, + residual: bool, + residual_region: reg.Region = (3, 3, 2, 2), + validation_error: bool = False, + ) -> dict[str, int]: + calls = {"scan": 0, "fill": 0, "mask": 0} + before = reg.MarkDetection("test", "Test mark", "top-left", True, 0.9, (2, 2, 4, 4)) + after = reg.MarkDetection("test", "Test mark", "top-left", residual, 0.7, residual_region) + + def detect_both(_image): + calls["scan"] += 1 + return before, before + + def detect_after(_image, *, provenance=False): + del provenance + calls["scan"] += 1 + if validation_error: + raise RuntimeError("validator failed") + return after + + def mask(_image, *, force=False, detection=None): + del force + calls["mask"] += 1 + # The perception result must reach the mask builder. Calling `_detect` + # again before the fill would discard the optimization that pays for the + # post-removal check. + assert detection is before + result = np.zeros((10, 10), np.uint8) + result[2:6, 2:6] = 255 + return result + + mark = reg.KnownMark( + key="test", + label="Test mark", + location="top-left", + in_auto=True, + product="test", + label_regime=None, + platform="Test", + _detect=detect_after, + _mask=mask, + _detect_both=detect_both, + ) + + def fill(image, selected_mask, *, backend="auto"): + del backend + calls["fill"] += 1 + result = image.copy() + result[selected_mask > 0] = 123 + return result + + monkeypatch.setattr(reg, "_REGISTRY", (mark,)) + monkeypatch.setattr(reg, "_PRODUCT_OF", {"test": "test"}) + monkeypatch.setattr(reg, "fill", fill) + return calls + + def test_cleaned_result_reuses_detection_and_fills_once(self, monkeypatch: pytest.MonkeyPatch): + calls = self._install_single_mark(monkeypatch, residual=False) + + report = reg.remove_auto_marks_detailed(np.zeros((10, 10, 3), np.uint8), backend="cv2") + + assert report.status == "cleaned" + assert report.labels == ["Test mark"] + assert report.marks[0].status == "cleaned" + assert report.marks[0].mask_bbox == (2, 2, 4, 4) + assert report.marks[0].confidence_before == 0.9 + assert report.marks[0].confidence_after == 0.7 + assert report.marks[0].residual_region is None + assert calls == {"scan": 2, "fill": 1, "mask": 1} + + def test_overlapping_residual_is_partial_without_retry(self, monkeypatch: pytest.MonkeyPatch): + calls = self._install_single_mark(monkeypatch, residual=True) + + report = reg.remove_auto_marks_detailed(np.zeros((10, 10, 3), np.uint8), backend="cv2") + + assert report.status == "partial" + assert report.marks[0].status == "partial" + assert report.marks[0].residual_region == (3, 3, 2, 2) + assert calls["fill"] == 1 + + def test_nonoverlapping_detection_is_not_this_removals_residual(self, monkeypatch: pytest.MonkeyPatch): + self._install_single_mark(monkeypatch, residual=True, residual_region=(8, 8, 2, 2)) + + report = reg.remove_auto_marks_detailed(np.zeros((10, 10, 3), np.uint8), backend="cv2") + + assert report.status == "cleaned" + assert report.marks[0].residual_region is None + + def test_validator_failure_is_not_reported_as_cleaned(self, monkeypatch: pytest.MonkeyPatch): + calls = self._install_single_mark(monkeypatch, residual=False, validation_error=True) + + report = reg.remove_auto_marks_detailed(np.zeros((10, 10, 3), np.uint8), backend="cv2") + + assert report.status == "unvalidated" + assert report.marks[0].status == "unvalidated" + assert report.marks[0].confidence_after is None + assert calls["fill"] == 1 + + def test_no_decision_is_no_watermark(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(reg, "_REGISTRY", ()) + monkeypatch.setattr(reg, "_PRODUCT_OF", {}) + + report = reg.remove_auto_marks_detailed(np.zeros((10, 10, 3), np.uint8), backend="cv2") + + assert report.status == "no_watermark" + assert report.labels == [] + assert report.marks == () + + class TestProvenanceGate: """The Gemini trust gate relaxes from GEMINI_SPARKLE_TRUST_CONF to _GEMINI_PROVENANCE_MIN_CONF when provenance confirms Google; tested