From cbb515578468c73961cf9ff7af541ebe5ac935c8 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 4 Sep 2026 20:59:09 -0700 Subject: [PATCH] Expose visible validation in pipeline results --- README.md | 7 ++ docs/cli.md | 9 +++ docs/module-internals.md | 14 +++- docs/python-api.md | 18 ++++- src/remove_ai_watermarks/__init__.py | 3 + src/remove_ai_watermarks/api.py | 100 ++++++++++++++++++++------- src/remove_ai_watermarks/cli.py | 45 ++++++++---- tests/test_api.py | 100 +++++++++++++++++++++++++++ tests/test_cli.py | 78 +++++++++++++++++++-- 9 files changed, 329 insertions(+), 45 deletions(-) diff --git a/README.md b/README.md index ee16b0a..69f199e 100644 --- a/README.md +++ b/README.md @@ -404,6 +404,13 @@ print(removed) report = raiw.remove_visible_detailed("watermarked.png", "clean.png") print(report.status) # cleaned | partial | unvalidated | no_watermark +complete = raiw.remove_all("watermarked.png", "clean.png") +print(complete.visible_status, complete.visible_marks) + +images = raiw.remove_batch("images", "images_clean", mode="visible") +for item in images.items: + print(item.source, item.visible_status, item.visible_marks) + 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 a924686..ebf5861 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -643,6 +643,10 @@ The command runs: 3. AI metadata stripping. The visible options and diffusion options are also available on `all`. +After the visible stage, `all` reports whether the once-filled result was +validated as cleaned, still produced an overlapping residual (`Partial`), or +could not be validated. These statuses do not trigger another fill and do not +change the exit code by themselves. When the `qwen-zimage` extra is unavailable, `all` still writes the result of the visible and metadata stages, prints a prominent warning, and exits with code @@ -680,6 +684,11 @@ The invisible and full modes accept the same main diffusion controls as their single image counterparts. Run `batch --help` for the authoritative option list. +For `visible` and `all` modes, the final summary includes counts for +`no watermark`, `cleaned`, `partial`, and `unvalidated`. Partial and unavailable +validation results are printed explicitly, but the once-filled outputs are kept +and the batch does not retry them automatically. + ## Exit behavior The CLI uses nonzero exit codes for meaningful incomplete outcomes, including diff --git a/docs/module-internals.md b/docs/module-internals.md index a5d743d..41a16b2 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -97,8 +97,12 @@ and the image pipeline that the `all` and `batch` commands are thin wrappers over: - `remove_all`, returning a `RemoveAllResult` after the visible, invisible, and - metadata stages -- `remove_batch`, returning a `BatchSummary` for one directory and one mode + metadata stages. Its legacy `visible_label` remains, while `visible_status` + and `visible_marks` retain the aggregate and per-mark post-fill result +- `remove_batch`, returning a `BatchSummary` for one directory and one mode. + `items` contains one lightweight `BatchItemResult` per source, including the + same visible status and per-mark records for visible/all modes without keeping + image arrays alive across the directory - `InvisibleOptions`, the invisible stage's knobs as one immutable value. Engine knobs only, under the engine's own names and defaults, so a bare `InvisibleOptions()` behaves exactly like calling the engine with no arguments. @@ -117,7 +121,11 @@ over: - `MetadataStripIncomplete`, raised before any write when AI metadata survives `remove_all` reports progress as `(stage, detail)` pairs of stable tokens, not -prose the caller has to parse back. +prose the caller has to parse back. Both pipeline entry points call +`remove_auto_marks_detailed` directly; carrying validation into their result +objects adds no detector or fill pass. The original positional constructors for +`RemoveAllResult` and `BatchSummary` remain valid through defaulted appended +fields. The package root exposes all of them lazily through [`__getattr__`](../src/remove_ai_watermarks/__init__.py), keeping a plain package diff --git a/docs/python-api.md b/docs/python-api.md index c84aa89..40790bc 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -123,10 +123,17 @@ import remove_ai_watermarks as raiw result = raiw.remove_all("input.png", "clean.png") # -> RemoveAllResult print(result.output) # the path written print(result.visible_label) # the marks removed, or None +print(result.visible_status) # "no_watermark" | "cleaned" | "partial" | "unvalidated" +for mark in result.visible_marks: + print(mark.label, mark.status, mark.confidence_before, mark.confidence_after) print(result.invisible) # "removed" | "no-signal" | "unavailable" ``` -`invisible` is the field to check. `"unavailable"` means the GPU extra is not +The visible status and per-mark records come from the same detector/fill/check +pass as `remove_visible_detailed`; `remove_all` does not run another detector or +fill. `partial` and `unvalidated` are quality statuses, not automatic retries. + +`invisible` is also important. `"unavailable"` means the GPU extra is not installed, so the output *looks* processed but still carries the watermark; `"no-signal"` means the scrub was deliberately skipped because nothing was locally detectable, which is a successful run. @@ -198,10 +205,19 @@ end the run: summary = raiw.remove_batch("in_dir", "out_dir", mode="visible") # -> BatchSummary print(summary.processed, summary.failed, summary.errors) print(summary.invisible_unavailable) # outputs that still carry the watermark +for item in summary.items: # -> BatchItemResult + print(item.source, item.output, item.visible_status, item.error) + for mark in item.visible_marks: + print(mark.label, mark.status) ``` `mode` is `all`, `visible`, `invisible`, or `metadata`. Pass a constructed `InvisibleEngine` as `engine` to load the model once for the whole directory. +Visible and all items contain the aggregate visible status and lightweight +per-mark records; image arrays are not retained in the batch summary. Other +modes leave those two fields unset and empty. The original `BatchSummary` +aggregate fields and `RemoveAllResult.visible_label` remain available for +existing callers. ## Classify a photograph from pixels diff --git a/src/remove_ai_watermarks/__init__.py b/src/remove_ai_watermarks/__init__.py index b966573..0fdc74f 100644 --- a/src/remove_ai_watermarks/__init__.py +++ b/src/remove_ai_watermarks/__init__.py @@ -37,6 +37,7 @@ _warnings.filterwarnings("ignore", message=r".*ImageProcessorFast.*") __version__ = "0.37.1" __all__ = [ + "BatchItemResult", "BatchSummary", "InvisibleOptions", "MetadataStripIncomplete", @@ -62,6 +63,7 @@ __all__ = [ if TYPE_CHECKING: from remove_ai_watermarks.api import ( + BatchItemResult, BatchSummary, InvisibleOptions, MetadataStripIncomplete, @@ -93,6 +95,7 @@ def __getattr__(name: str) -> object: """Lazily resolve the high-level API (PEP 562), so the heavy imports (cv2, the metadata/identify stack) load only when a caller actually reaches for them.""" if name in ( + "BatchItemResult", "BatchSummary", "InvisibleOptions", "MetadataStripIncomplete", diff --git a/src/remove_ai_watermarks/api.py b/src/remove_ai_watermarks/api.py index ce869d2..614adef 100644 --- a/src/remove_ai_watermarks/api.py +++ b/src/remove_ai_watermarks/api.py @@ -27,7 +27,13 @@ if TYPE_CHECKING: from numpy.typing import NDArray - from remove_ai_watermarks.watermark_registry import Backend, Sensitivity, VisibleRemovalResult + from remove_ai_watermarks.watermark_registry import ( + Backend, + MarkRemovalResult, + Sensitivity, + VisibleRemovalResult, + VisibleRemovalStatus, + ) def __getattr__(name: str) -> object: @@ -295,11 +301,19 @@ InvisibleOutcome = Literal["removed", "no-signal", "unavailable"] @dataclass(frozen=True) class RemoveAllResult: - """Outcome of :func:`remove_all` -- what each stage actually did.""" + """Outcome of :func:`remove_all`, including visible validation details. + + ``visible_label`` preserves the original compact contract. ``visible_status`` and + ``visible_marks`` carry the aggregate and per-mark results from the same pass. + """ output: Path visible_label: str | None # the removed mark(s), or None when nothing fired invisible: InvisibleOutcome + # Optional defaults keep direct construction with the original three positional + # fields compatible. Results returned by remove_all always populate both fields. + visible_status: VisibleRemovalStatus | None = None + visible_marks: tuple[MarkRemovalResult, ...] = () class MetadataStripIncomplete(RuntimeError): @@ -437,6 +451,10 @@ def remove_all( prose: the caller owns the wording. The invisible stage reports its :data:`InvisibleOutcome`, plus a ``strength=`` line before it runs. + The returned ``visible_status`` and ``visible_marks`` preserve the detailed + post-fill result from stage one. They do not add another detector or fill pass; + ``visible_label`` remains the compact compatibility field. + Raises :class:`MetadataStripIncomplete` before writing anything when AI metadata survives, and ``OSError`` when the output cannot be written. """ @@ -468,12 +486,13 @@ def remove_all( staged = Path(tmp_name) try: # ── 1. Visible marks ── - result, removed = watermark_registry.remove_auto_marks( + visible = watermark_registry.remove_auto_marks_detailed( image, sensitivity=sensitivity, provenance=evidence.visible_provenance(), backend=backend, ) + result, removed = visible.image, visible.labels visible_label = ", ".join(removed) if removed else None say("visible", visible_label or "") if not image_io.write_bgr_with_alpha(staged, result, alpha): @@ -502,7 +521,7 @@ def remove_all( finally: if staged.exists(): staged.unlink() - return RemoveAllResult(out, visible_label, outcome) + return RemoveAllResult(out, visible_label, outcome, visible.status, visible.marks) def _run_invisible( @@ -580,9 +599,26 @@ def _run_invisible( BatchMode = Literal["all", "visible", "metadata", "invisible"] +@dataclass(frozen=True, slots=True) +class BatchItemResult: + """Outcome and visible validation details for one batch source.""" + + source: Path + output: Path | None + mode: BatchMode + visible_status: VisibleRemovalStatus | None = None + visible_marks: tuple[MarkRemovalResult, ...] = () + invisible: InvisibleOutcome | None = None + error: str | None = None + + @dataclass(frozen=True) class BatchSummary: - """Per-directory outcome of :func:`remove_batch`.""" + """Per-directory outcome of :func:`remove_batch`. + + ``items`` keeps each file's stage outcomes and visible validation details without + retaining its image array in memory. The original aggregate fields remain intact. + """ processed: int failed: int @@ -590,6 +626,8 @@ class BatchSummary: # missing. Non-empty means the outputs LOOK processed but still carry it. invisible_unavailable: list[Path] errors: list[tuple[Path, str]] + # Optional default preserves the original four-positional-field constructor. + items: tuple[BatchItemResult, ...] = () def remove_batch( @@ -611,6 +649,11 @@ def remove_batch( directory. ``force`` and ``engine`` are threaded straight through, so a caller that passes a constructed ``InvisibleEngine`` loads the model once for the whole run. + ``items`` on the returned summary carries one :class:`BatchItemResult` per source. + Visible and all modes populate its ``visible_status`` and ``visible_marks`` from the + same once-filled pass; other modes and failures before that pass leave them None and + empty, respectively. + ``progress`` receives ``(path, stage, detail)``. Every image ends with exactly one terminal stage -- ``done`` or ``failed`` -- whatever the mode does in between, so a caller driving a progress bar can advance on that alone. @@ -627,23 +670,27 @@ def remove_batch( processed = failed = 0 unavailable: list[Path] = [] errors: list[tuple[Path, str]] = [] + items: list[BatchItemResult] = [] for img_path in sorted(p for p in src_dir.iterdir() if is_supported_format(p)): out_path = out_dir / img_path.name try: - outcome = _run_batch_one(img_path, out_path, mode, backend, sensitivity, invisible, force, engine, say) + result = _run_batch_one(img_path, out_path, mode, backend, sensitivity, invisible, force, engine, say) except Exception as exc: failed += 1 - errors.append((img_path, str(exc))) - say(img_path, "failed", str(exc)) + message = str(exc) + errors.append((img_path, message)) + items.append(BatchItemResult(img_path, None, mode, error=message)) + say(img_path, "failed", message) continue processed += 1 - if outcome == "unavailable": + items.append(result) + if result.invisible == "unavailable": unavailable.append(img_path) # Exactly one terminal event per image, in every mode: a caller's progress bar # advances on this and nothing else. Keying it off a mode-specific stage line # left `visible` and `metadata` runs sitting at 0% for the whole batch. - say(img_path, "done", outcome or "") - return BatchSummary(processed, failed, unavailable, errors) + say(img_path, "done", result.invisible or "") + return BatchSummary(processed, failed, unavailable, errors, tuple(items)) def _run_batch_one( @@ -656,8 +703,8 @@ def _run_batch_one( force: bool, engine: Any | None, say: Callable[[Path, str, str], None], -) -> InvisibleOutcome | None: - """One image, one mode. Returns the invisible outcome when that stage ran.""" +) -> BatchItemResult: + """Run one image in one mode and return its stage details.""" from remove_ai_watermarks import image_io from remove_ai_watermarks.metadata import strip_and_verify @@ -672,7 +719,14 @@ def _run_batch_one( engine=engine, progress=lambda stage, detail: say(img_path, stage, detail), ) - return result.invisible + return BatchItemResult( + img_path, + out_path, + mode, + result.visible_status, + result.visible_marks, + result.invisible, + ) # NOTE on holders in batch: `out_path` may EQUAL `img_path` (nothing forbids # `-o `), and the visible stage rewrites it. So each stage that @@ -689,24 +743,22 @@ def _run_batch_one( # re-processed as if it were the input. from remove_ai_watermarks import watermark_registry - image, alpha = image_io.read_bgr_and_alpha(img_path) - if image is None: - raise ValueError(f"Could not read image: {img_path}") - result, _ = watermark_registry.remove_auto_marks( - image, + loaded = _load_visible_input(img_path) + visible = watermark_registry.remove_auto_marks_detailed( + loaded.bgr, sensitivity=sensitivity, - provenance=visible_provenance(img_path), + provenance=loaded.provenance, backend=backend, ) - if not image_io.write_bgr_with_alpha(out_path, result, alpha): + if not image_io.write_bgr_with_alpha(out_path, visible.image, loaded.alpha): raise OSError(f"failed to write output (is the destination writable?): {out_path}") - return None + return BatchItemResult(img_path, out_path, mode, visible.status, visible.marks) if mode == "metadata": _, leftover = strip_and_verify(img_path, out_path) if leftover: raise MetadataStripIncomplete(set(leftover)) - return None + return BatchItemResult(img_path, out_path, mode) # invisible-only: no preceding visible pass, so out_path does not exist yet, and # the input IS the pristine original for both the gate and the vendor probe. @@ -726,4 +778,4 @@ def _run_batch_one( src_bgr, src_alpha = image_io.read_bgr_and_alpha(img_path) if src_bgr is None or not image_io.write_bgr_with_alpha(out_path, src_bgr, src_alpha): raise OSError(f"failed to copy input through to output: {out_path}") - return outcome + return BatchItemResult(img_path, out_path, mode, invisible=outcome) diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 8dd2565..eb2cf1b 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -560,6 +560,21 @@ def main(ctx: click.Context, verbose: bool) -> None: # ── Visible (Gemini) watermark removal ── +def _print_visible_validation(status: str | None, marks: tuple[Any, ...], labels: str) -> None: + """Render one visible pass consistently across ``visible`` and ``all``.""" + if status == "partial": + residuals = [mark.label for mark in marks if mark.status == "partial"] + console.print( + f" Partial: fill completed for {labels}; overlapping residual still detected for {', '.join(residuals)}." + ) + elif status == "unvalidated": + unavailable = [mark.label for mark in marks if mark.status == "unvalidated"] + console.print(f" Filled: {labels}") + console.print(f" Post-removal validation unavailable for: {', '.join(unavailable)}.") + elif status == "cleaned": + console.print(f" Removed and validated: {labels}") + + def _run_visible_auto( source: Path, output: Path, @@ -600,18 +615,7 @@ def _run_visible_auto( # 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) - 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)}") + _print_visible_validation(report.status, report.marks, ", ".join(removed)) size_kb = output.stat().st_size / 1024 console.print(f" Saved: {output} ({size_kb:.0f} KB, {elapsed:.2f}s)") @@ -1724,6 +1728,8 @@ def cmd_all( console.print(f" Error: {e}") raise SystemExit(1) from e + _print_visible_validation(outcome.visible_status, outcome.visible_marks, outcome.visible_label or "") + # ── Done ── elapsed = time.monotonic() - t0 size_kb = output.stat().st_size / 1024 @@ -1901,6 +1907,21 @@ def cmd_batch( console.print(f"\n {processed} processed" + (f" {errors} errors" if errors else "")) + status_order = ("no_watermark", "cleaned", "partial", "unvalidated") + counts = dict.fromkeys(status_order, 0) + for item in summary.items: + if item.visible_status is not None: + counts[item.visible_status] += 1 + if sum(counts.values()): + details = ", ".join(f"{count} {status.replace('_', ' ')}" for status, count in counts.items() if count) + console.print(f" Visible validation: {details}") + partial = counts["partial"] + unvalidated = counts["unvalidated"] + if partial: + console.print(f" Warning: overlapping visible residual detected in {partial} output(s).") + if unvalidated: + console.print(f" Post-removal validation unavailable for {unvalidated} output(s).") + if synthid_skipped_count: # Mirror the single `all` command: a silently retained SynthID watermark is the # #1 "it didn't work" report, so make the skipped scrub impossible to miss. diff --git a/tests/test_api.py b/tests/test_api.py index 5faccbb..e051d2d 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -20,6 +20,7 @@ class TestTopLevelExports: def test_lazy_reexports_resolve(self): from remove_ai_watermarks import openai_provenance + assert raiw.BatchItemResult is api.BatchItemResult assert raiw.remove_visible is api.remove_visible assert raiw.remove_visible_detailed is api.remove_visible_detailed assert raiw.VisibleRemovalResult is api.VisibleRemovalResult @@ -327,6 +328,52 @@ class TestRemoveAllLibrary: monkeypatch.setattr(invisible_engine, "is_available", lambda: True) return FakeEngine() + def test_propagates_the_detailed_visible_result(self, monkeypatch, tmp_path): + from remove_ai_watermarks import image_io, invisible_engine, watermark_registry + + image = np.full((96, 96, 3), 120, np.uint8) + source = tmp_path / "input.png" + assert image_io.imwrite(source, image) + mark = watermark_registry.MarkRemovalResult( + key="gemini", + label="Gemini visible watermark", + location="bottom-right", + region=(60, 60, 20, 20), + mask_bbox=(58, 58, 24, 24), + backend="cv2", + confidence_before=0.9, + confidence_after=0.7, + status="partial", + residual_region=(62, 62, 16, 16), + ) + visible = watermark_registry.VisibleRemovalResult(image, (mark,)) + detailed_calls: list[int] = [] + + def fake_detailed(*_args, **_kwargs): + detailed_calls.append(1) + return visible + + monkeypatch.setattr(watermark_registry, "remove_auto_marks_detailed", fake_detailed) + monkeypatch.setattr( + watermark_registry, + "remove_auto_marks", + lambda *_a, **_k: pytest.fail("remove_all discarded the detailed visible result"), + ) + monkeypatch.setattr(invisible_engine, "is_available", lambda: False) + + result = api.remove_all(source, tmp_path / "output.png", backend="cv2") + + assert result.visible_label == mark.label + assert result.visible_status == "partial" + assert result.visible_marks == (mark,) + assert detailed_calls == [1] + + def test_legacy_result_constructor_remains_valid(self, tmp_path): + result = api.RemoveAllResult(tmp_path / "output.png", None, "no-signal") + + assert result.visible_status is None + assert result.visible_marks == () + @pytest.mark.skipif(not DOUBAO.exists(), reason="doubao sample not present") def test_runs_all_three_stages_and_reports_each(self, monkeypatch, tmp_path): from remove_ai_watermarks import api @@ -418,6 +465,59 @@ class TestRemoveAllLibrary: class TestRemoveBatchLibrary: + @pytest.mark.parametrize("mode", ["visible", "all"]) + def test_records_each_visible_validation_result(self, monkeypatch, tmp_path, mode): + from remove_ai_watermarks import image_io, invisible_engine, watermark_registry + + src = tmp_path / "in" + src.mkdir() + image = np.full((96, 96, 3), 120, np.uint8) + assert image_io.imwrite(src / "input.png", image) + mark = watermark_registry.MarkRemovalResult( + key="gemini", + label="Gemini visible watermark", + location="bottom-right", + region=(60, 60, 20, 20), + mask_bbox=(58, 58, 24, 24), + backend="cv2", + confidence_before=0.9, + confidence_after=None, + status="unvalidated", + ) + visible = watermark_registry.VisibleRemovalResult(image, (mark,)) + detailed_calls: list[int] = [] + + def fake_detailed(*_args, **_kwargs): + detailed_calls.append(1) + return visible + + monkeypatch.setattr(watermark_registry, "remove_auto_marks_detailed", fake_detailed) + monkeypatch.setattr( + watermark_registry, + "remove_auto_marks", + lambda *_a, **_k: pytest.fail("batch discarded the detailed visible result"), + ) + monkeypatch.setattr(invisible_engine, "is_available", lambda: False) + + summary = api.remove_batch(src, tmp_path / "out", mode=mode, backend="cv2") + + assert summary.processed == 1 + assert summary.failed == 0 + assert len(summary.items) == 1 + item = summary.items[0] + assert item.source == src / "input.png" + assert item.output == tmp_path / "out" / "input.png" + assert item.mode == mode + assert item.visible_status == "unvalidated" + assert item.visible_marks == (mark,) + assert item.error is None + assert detailed_calls == [1] + + def test_legacy_summary_constructor_remains_valid(self): + summary = api.BatchSummary(0, 0, [], []) + + assert summary.items == () + @pytest.mark.skipif(not DOUBAO.exists(), reason="doubao sample not present") def test_visible_mode_writes_every_image(self, tmp_path): import shutil diff --git a/tests/test_cli.py b/tests/test_cli.py index 97f818b..69c3803 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -538,6 +538,35 @@ class TestAllCommand: assert result.exit_code == 0, result.output assert output.exists() + def test_all_reports_a_partial_visible_validation(self, runner, sample_png, tmp_path): + from remove_ai_watermarks.api import RemoveAllResult + from remove_ai_watermarks.watermark_registry import MarkRemovalResult + + output = tmp_path / "clean.png" + 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.7, + status="partial", + residual_region=(12, 12, 10, 10), + ) + + def fake_remove_all(_source, destination, **_kwargs): + destination.write_bytes(sample_png.read_bytes()) + return RemoveAllResult(destination, mark.label, "no-signal", "partial", (mark,)) + + with patch("remove_ai_watermarks.api.remove_all", side_effect=fake_remove_all): + result = runner.invoke(main, ["all", str(sample_png), "-o", str(output)]) + + assert result.exit_code == 0, result.output + assert "overlapping residual still detected" in result.output + assert mark.label in result.output + def test_all_cpu_offload_flows_to_engine(self, runner, sample_png, tmp_path): mock_cls, _mock_engine = _mock_invisible_engine() output = tmp_path / "clean.png" @@ -578,22 +607,24 @@ class TestAllCommand: result = runner.invoke(main, ["all", "/nonexistent/file.png"]) assert result.exit_code != 0 - def test_all_visible_step_uses_registry(self, runner, sample_png, tmp_path): + def test_all_visible_step_uses_detailed_registry(self, runner, sample_png, tmp_path): """Regression (#1): the `all` visible step must route through the registry - (remove_auto_marks), so Doubao/Jimeng/Samsung/pill marks are handled -- not - just the Gemini sparkle via a hardcoded GeminiEngine.""" + auto detector, so Doubao/Jimeng/Samsung/pill marks are handled -- not just the + Gemini sparkle via a hardcoded GeminiEngine -- and validation is preserved.""" + from remove_ai_watermarks.watermark_registry import VisibleRemovalResult + mock_cls, _mock_engine = _mock_invisible_engine() output = tmp_path / "clean.png" def _fake_remove_auto(image, **kwargs): - return image, [] + return VisibleRemovalResult(image, ()) with ( patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True), patch( - "remove_ai_watermarks.watermark_registry.remove_auto_marks", side_effect=_fake_remove_auto + "remove_ai_watermarks.watermark_registry.remove_auto_marks_detailed", side_effect=_fake_remove_auto ) as mock_auto, ): result = runner.invoke(main, ["all", str(sample_png), "-o", str(output), "--force"]) @@ -919,6 +950,43 @@ class TestBatchCommand: assert output_dir.exists() assert len(list(output_dir.glob("*.png"))) == 3 + def test_batch_reports_visible_validation_counts(self, runner, tmp_path): + from remove_ai_watermarks.api import BatchItemResult, BatchSummary + from remove_ai_watermarks.watermark_registry import MarkRemovalResult + + input_dir = _make_batch_dir(tmp_path, count=1) + source = input_dir / "img_0.png" + output_dir = tmp_path / "output" + 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=None, + status="unvalidated", + ) + item = BatchItemResult( + source, + output_dir / source.name, + "visible", + "unvalidated", + (mark,), + ) + summary = BatchSummary(1, 0, [], [], (item,)) + + with patch("remove_ai_watermarks.api.remove_batch", return_value=summary): + result = runner.invoke( + main, + ["batch", str(input_dir), "-o", str(output_dir), "--mode", "visible"], + ) + + assert result.exit_code == 0, result.output + assert "Visible validation: 1 unvalidated" in result.output + assert "validation unavailable for 1 output" in result.output + def test_batch_metadata_mode(self, runner, tmp_path): input_dir = _make_batch_dir_with_metadata(tmp_path) output_dir = tmp_path / "output"