mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-07-12 11:06:33 +02:00
Victor Kuznetsov
1a955b096a
- Replace reverse-alpha removal with localize -> fill (template-free mask + one shared cv2/MI-GAN/big-LaMa fill) for every mark; drops the colour-shift / dark-pit failure modes, version-robust to a moved or re-rendered mark - Separate perception/decision/action: engines report Candidates, a pure decide(candidates, Context) arbiter owns all policy (sensitivity + provenance + pill gate), remove_auto_marks orchestrates -- behavior-preserving (corpus 46/46/92) - Three orthogonal knobs replace --method: --backend cv2|migan|lama, --sensitivity auto|strict|assume-ai, provenance (auto from metadata) - Add high-level api.remove_visible / visible_provenance (lazy top-level re-export); visible --mark auto delegates to it so CLI and library share ONE path - Read+write HEIC/AVIF on the pixel path via pillow-heif; imwrite preserves the input format at max quality (JPEG q100/4:4:4); a no-op copies the original bytes verbatim - Lossless byte-level JPEG metadata strip (no DCT re-encode); consolidate the two remove_ai_metadata into one, delete legacy noai/cleaner + best_auto_mark - Bump 0.13.0 -> 0.14.0 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
106 lines
4.2 KiB
Python
106 lines
4.2 KiB
Python
"""High-level convenience API (remove_visible / visible_provenance) and the lazy
|
|
top-level re-exports."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import remove_ai_watermarks as raiw
|
|
from remove_ai_watermarks import api
|
|
|
|
SAMPLES = Path(__file__).resolve().parents[1] / "data" / "samples"
|
|
DOUBAO = SAMPLES / "doubao-1.png"
|
|
CHATGPT = SAMPLES / "chatgpt-1.png"
|
|
|
|
|
|
class TestTopLevelExports:
|
|
def test_lazy_reexports_resolve(self):
|
|
assert raiw.remove_visible is api.remove_visible
|
|
assert raiw.visible_provenance is api.visible_provenance
|
|
|
|
def test_unknown_attribute_raises(self):
|
|
with pytest.raises(AttributeError):
|
|
_ = raiw.does_not_exist
|
|
|
|
def test_bare_import_is_light(self):
|
|
# importing the package must not pull the heavy cv2/torch stack (PEP 562 lazy).
|
|
# Checked in a FRESH interpreter -- another test in this process may already
|
|
# have imported cv2, so an in-process sys.modules check would be flaky.
|
|
import subprocess
|
|
import sys
|
|
|
|
code = "import remove_ai_watermarks, sys; print(int(any(m in sys.modules for m in ('cv2','torch'))))"
|
|
out = subprocess.run( # noqa: S603 -- fixed sys.executable + literal code, no untrusted input
|
|
[sys.executable, "-c", code], check=True, capture_output=True, text=True
|
|
)
|
|
assert out.stdout.strip() == "0", f"bare import pulled a heavy module: {out.stdout!r}"
|
|
|
|
|
|
class TestRemoveVisibleArray:
|
|
def test_array_no_mark_is_noop_copy(self):
|
|
arr = np.zeros((256, 256, 3), np.uint8)
|
|
result, removed = raiw.remove_visible(arr, backend="cv2")
|
|
assert removed == []
|
|
assert result.shape == arr.shape
|
|
assert np.array_equal(result, arr)
|
|
|
|
def test_array_accepts_knobs(self):
|
|
arr = np.zeros((256, 256, 3), np.uint8)
|
|
result, removed = raiw.remove_visible(arr, sensitivity="assume_ai", backend="cv2")
|
|
assert removed == []
|
|
assert result.shape == arr.shape
|
|
|
|
def test_bad_source_raises(self, tmp_path):
|
|
with pytest.raises(ValueError, match="Could not read image"):
|
|
raiw.remove_visible(tmp_path / "nope.png")
|
|
|
|
|
|
@pytest.mark.skipif(not DOUBAO.exists(), reason="doubao sample not present")
|
|
class TestRemoveVisiblePath:
|
|
def test_path_removes_and_writes(self, tmp_path):
|
|
out = tmp_path / "clean.png"
|
|
result, removed = raiw.remove_visible(DOUBAO, out, backend="cv2")
|
|
assert out.exists()
|
|
assert any("Doubao" in lbl for lbl in removed)
|
|
assert result.shape[2] == 3
|
|
|
|
def test_path_no_output_returns_without_writing(self, tmp_path):
|
|
# output=None returns the array but writes nothing
|
|
result, _ = raiw.remove_visible(DOUBAO, backend="cv2")
|
|
assert result.ndim == 3
|
|
|
|
|
|
class TestNoOpPreservesOriginal:
|
|
def test_no_mark_copies_original_bytes(self, tmp_path):
|
|
# A clean image (no mark) same-format-out must be copied VERBATIM, not
|
|
# re-encoded -- so a no-op never degrades the original ("work with originals").
|
|
import filecmp
|
|
|
|
from PIL import Image
|
|
|
|
src = tmp_path / "clean.jpg"
|
|
Image.fromarray(np.full((40, 40, 3), 120, np.uint8), "RGB").save(src, quality=90)
|
|
out = tmp_path / "clean_out.jpg"
|
|
_, removed = raiw.remove_visible(str(src), str(out), sensitivity="strict", backend="cv2")
|
|
assert removed == []
|
|
assert filecmp.cmp(str(src), str(out), shallow=False) # byte-identical
|
|
|
|
|
|
class TestVisibleProvenance:
|
|
@pytest.mark.skipif(not DOUBAO.exists(), reason="doubao sample not present")
|
|
def test_doubao_tc260_maps_to_bytedance(self):
|
|
prov = raiw.visible_provenance(DOUBAO)
|
|
# TC260 label -> ByteDance family (both doubao and jimeng)
|
|
assert {"doubao", "jimeng"} <= prov
|
|
|
|
@pytest.mark.skipif(not CHATGPT.exists(), reason="chatgpt sample not present")
|
|
def test_openai_image_has_no_visible_vendor(self):
|
|
# OpenAI C2PA is not one of the visible-mark vendors -> empty provenance
|
|
assert raiw.visible_provenance(CHATGPT) == frozenset()
|
|
|
|
def test_unreadable_path_is_empty(self, tmp_path):
|
|
assert raiw.visible_provenance(tmp_path / "missing.png") == frozenset()
|