mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-10-21 16:21:42 +02:00
Expose visible validation in pipeline results
This commit is contained in:
9 files changed
+329
-45
No files matched your search
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+17
-1
@@ -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
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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=<value>`` 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 <the input directory>`), 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)
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
+73
-5
@@ -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"
|
||||
|
||||
Reference in new issue
Block a user