mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-07 06:28:36 +02:00
fix: metadata-strip parity, input robustness, and detection/clash coverage
Bug fixes (each with a regression test): - metadata strip parity across every marker placement: IPTC digitalSourceType in XMP, the Samsung post-EOI trailer, the China TC260 AIGC block in EXIF UserComment, a bare AIGC block in a non-standard APP segment, and the ISOBMFF EXIF path (AIGC + xAI) are all now stripped -- anything a scanner flags, the strip reaches - Samsung genAIType detected when its trailer sits past the 512 KB scan window (file-tail read on large photos) - crashes on edge inputs: Gemini detector on images with a short side < 16px, footprint_mask on a zero-size ndarray, the humanizer on chromatic_shift >= width, and the CLI on unreadable/corrupt/empty input (clean error, not a traceback) - WebP written losslessly (cv2 quality 101), not lossy at 100 - the IPTC digitalSourceType algorithmicMedia (procedural, not trained on sampled data) is no longer flagged as AI-generated, so clean procedural content is not scrubbed - c2pa source-type: compositeWithTrainedAlgorithmicMedia is checked before the bare algorithmicMedia token, so an AI-enhanced composite is not misclassified Detection: - integrity-clash coverage now normalizes ByteDance / Canva / ElevenLabs / Black Forest Labs, so a transplanted manifest next to an independent conflicting stamp is caught; the generic China TC260 AIGC label is attributed to a co-present TC260 vendor, so a legit Doubao image (its own C2PA + TC260 label) does not clash (corpus-validated: 0 new clashes on 5000 carriers) CLI: - batch exits non-zero (with a warning) when any image errors or a GPU-missing SynthID scrub is skipped, and copies the input through so the output dir stays complete -- it used to always exit 0 and could silently drop files Perf: - GeminiEngine reused as a process-wide singleton with a precomputed template ladder: -24% on the identify sparkle path, detection byte-identical Internal: one shared _ai_exif_targets rule set feeds both EXIF scrubbers so their coverage cannot drift; docs synced; maintain.sh hardened so the uv-secure internal teardown crash no longer aborts the gate (still fails on a real finding). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
190dc89d23
commit
a4c901ff39
@@ -573,6 +573,16 @@ class TestMetadataCommand:
|
||||
assert result.exit_code == 0
|
||||
assert "stripped" in result.output
|
||||
|
||||
def test_metadata_remove_in_place(self, runner, tmp_png_with_ai_metadata):
|
||||
"""With ``-o`` omitted, the strip overwrites the source in place (default
|
||||
output_path=None). Previously every test passed an explicit ``-o``."""
|
||||
from remove_ai_watermarks.metadata import has_ai_metadata
|
||||
|
||||
assert has_ai_metadata(tmp_png_with_ai_metadata) # precondition
|
||||
result = runner.invoke(main, ["metadata", str(tmp_png_with_ai_metadata), "--remove"])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert not has_ai_metadata(tmp_png_with_ai_metadata) # source overwritten, AI metadata gone
|
||||
|
||||
|
||||
class TestIdentifyCommand:
|
||||
"""Tests for the 'identify' subcommand."""
|
||||
@@ -600,6 +610,18 @@ class TestIdentifyCommand:
|
||||
assert "AI-generated" in result.output
|
||||
assert "Stable Diffusion" in result.output
|
||||
|
||||
def test_identify_reports_generated_source_kind(self, runner):
|
||||
"""The C2PA trainedAlgorithmicMedia source type sharpens the verdict to
|
||||
'AI-generated (fully synthetic)' at the CLI (the ai_source_kind branch)."""
|
||||
from pathlib import Path
|
||||
|
||||
sample = Path(__file__).resolve().parent.parent / "data" / "samples" / "chatgpt-1.png"
|
||||
if not sample.exists():
|
||||
pytest.skip("chatgpt sample not present")
|
||||
result = runner.invoke(main, ["identify", str(sample), "--no-visible"])
|
||||
assert result.exit_code == 0
|
||||
assert "AI-generated (fully synthetic)" in result.output
|
||||
|
||||
def test_identify_json_is_valid(self, runner, tmp_png_with_ai_metadata):
|
||||
result = runner.invoke(main, ["identify", str(tmp_png_with_ai_metadata), "--no-visible", "--json"])
|
||||
assert result.exit_code == 0
|
||||
@@ -774,6 +796,34 @@ class TestBatchCommand:
|
||||
expected_dir = tmp_path / "input_clean"
|
||||
assert expected_dir.exists()
|
||||
|
||||
def test_batch_errors_exit_nonzero(self, runner, tmp_path):
|
||||
"""Regression: batch used to always exit 0 even when every image errored,
|
||||
hiding failure from a wrapping service. A corrupt image must yield a non-zero
|
||||
exit and an error count."""
|
||||
input_dir = tmp_path / "input"
|
||||
input_dir.mkdir()
|
||||
(input_dir / "corrupt.png").write_bytes(b"this is not a PNG at all" * 50)
|
||||
result = runner.invoke(main, ["batch", str(input_dir), "--mode", "visible"])
|
||||
assert result.exit_code != 0, result.output
|
||||
assert "error" in result.output.lower()
|
||||
|
||||
def test_batch_invisible_gpu_missing_writes_output_and_exits_nonzero(self, runner, tmp_path):
|
||||
"""Regression: batch --mode invisible with a signal-bearing image but no GPU
|
||||
deps used to write NO output for that image and still exit 0, silently dropping
|
||||
the files that most needed processing. It must now copy the input through (so the
|
||||
output dir is complete), warn about the retained SynthID watermark, and exit
|
||||
non-zero -- mirroring the single ``all`` command."""
|
||||
input_dir = _make_batch_dir_with_metadata(tmp_path, count=3) # SD params = invisible signal
|
||||
output_dir = tmp_path / "output"
|
||||
with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["batch", str(input_dir), "-o", str(output_dir), "--mode", "invisible"],
|
||||
)
|
||||
assert result.exit_code != 0, result.output
|
||||
assert "NOT removed" in result.output
|
||||
assert len(list(output_dir.glob("*.png"))) == 3 # every input copied through, none dropped
|
||||
|
||||
|
||||
class TestGpuHintMarkup:
|
||||
"""The GPU-extra install hint must reach the user with the ``[gpu]`` token
|
||||
@@ -887,3 +937,26 @@ def test_visible_backend_runtime_error_exits_cleanly(runner, tmp_path, monkeypat
|
||||
result = runner.invoke(main, ["visible", str(doubao), "-o", str(out), "--backend", "migan"])
|
||||
assert result.exit_code == 1
|
||||
assert not isinstance(result.exception, RuntimeError), "RuntimeError leaked as a traceback"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("name", "content"),
|
||||
[
|
||||
("empty.png", b""),
|
||||
("notimage.jpg", b"plain text, not an image at all " * 20),
|
||||
("truncated.png", b"\x89PNG\r\n\x1a\n" + b"\x00" * 40),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("cmd", [["metadata", "--remove"], ["visible", "--backend", "cv2"]])
|
||||
def test_unreadable_input_exits_cleanly(runner, tmp_path, name, content, cmd):
|
||||
"""Regression: a corrupt / empty / non-image file (real prod uploads include
|
||||
truncated files) must produce a clean 'Error: cannot read/process' + exit 1, NOT a
|
||||
raw PIL.UnidentifiedImageError / OSError / ValueError traceback. Found by the runtime
|
||||
mode fuzz across metadata --remove and visible."""
|
||||
bad = tmp_path / name
|
||||
bad.write_bytes(content)
|
||||
out = tmp_path / "out.png"
|
||||
result = runner.invoke(main, [cmd[0], str(bad), "-o", str(out), *cmd[1:]])
|
||||
assert result.exit_code == 1, result.output
|
||||
assert isinstance(result.exception, SystemExit), f"leaked a raw traceback: {result.exception!r}"
|
||||
assert "Error" in result.output
|
||||
|
||||
@@ -195,3 +195,29 @@ class TestDegenerateAndChannelInputs:
|
||||
bgra = np.zeros((2048, 2048, 4), np.uint8)
|
||||
mask = eng.footprint_mask(bgra, force=True)
|
||||
assert mask is None or mask.shape == (2048, 2048)
|
||||
|
||||
def test_template_match_score_guards_return_zero(self):
|
||||
# Guards return 0.0 (never a false positive) for a mask that cannot hold a
|
||||
# glyph: empty, narrower than min_gw, or shorter than the 4-px floor.
|
||||
assert _template_match_score(np.zeros((0, 5), np.uint8), 1000) == 0.0
|
||||
assert _template_match_score(np.zeros((10, 3), np.uint8), 1000) == 0.0 # width-1 < min_gw
|
||||
assert _template_match_score(np.zeros((3, 200), np.uint8), 1000) == 0.0 # height-1 < 4
|
||||
|
||||
@pytest.mark.parametrize("shape", [(20, 20, 3), (10, 400, 3), (400, 10, 3), (1, 1, 3), (2000, 2000, 3)])
|
||||
def test_locate_box_stays_in_bounds(self, shape):
|
||||
"""locate() must clamp its geometry box inside the image for ANY size/aspect --
|
||||
wide-short, tall-narrow, 1x1, huge -- for both bottom corners (br + bl)."""
|
||||
from remove_ai_watermarks._text_mark_engine import TextMarkEngine
|
||||
from remove_ai_watermarks.doubao_engine import _CONFIG as BR_CONFIG
|
||||
from remove_ai_watermarks.samsung_engine import _CONFIG as BL_CONFIG
|
||||
|
||||
h, w = shape[:2]
|
||||
img = np.zeros(shape, np.uint8)
|
||||
for cfg in (BR_CONFIG, BL_CONFIG):
|
||||
loc = TextMarkEngine(cfg).locate(img)
|
||||
assert loc.x >= 0
|
||||
assert loc.y >= 0
|
||||
assert loc.x + loc.w <= w
|
||||
assert loc.y + loc.h <= h
|
||||
assert loc.w > 0
|
||||
assert loc.h > 0
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from remove_ai_watermarks.humanizer import apply_analog_humanizer, unsharp_mask
|
||||
|
||||
@@ -72,6 +73,15 @@ def test_chromatic_shift_does_not_wrap_opposite_edge():
|
||||
assert result[:, -shift:, 2].min() > 195
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("width", "shift"), [(1, 1), (3, 3), (3, 5), (5, 10)])
|
||||
def test_chromatic_shift_wider_than_image_no_crash(width: int, shift: int):
|
||||
"""Regression: a chromatic_shift >= image width left the edge-replication slices
|
||||
empty and crashed the broadcast (ValueError). The shift must clamp to width-1."""
|
||||
img = np.full((4, width, 3), 120, np.uint8)
|
||||
result = apply_analog_humanizer(img, grain_intensity=0.0, chromatic_shift=shift)
|
||||
assert result.shape == img.shape
|
||||
|
||||
|
||||
def test_unsharp_disabled_returns_unchanged_copy():
|
||||
img = np.full((20, 20, 3), 128, dtype=np.uint8)
|
||||
img[10, 10] = [100, 150, 200]
|
||||
@@ -147,3 +157,19 @@ class TestAdaptivePolish:
|
||||
a = adaptive_polish(soft, reference, seed=7)
|
||||
b = adaptive_polish(soft, reference, seed=7)
|
||||
assert np.array_equal(a, b)
|
||||
|
||||
def test_all_edges_reference_grain_mask_near_zero(self):
|
||||
# An all-high-frequency target: _smooth_grain_mask suppresses edges, so the grain
|
||||
# mask is ~all-zero (grain adds nothing) -- adaptive_polish must still return a
|
||||
# valid same-shape image, not crash on the empty-mask branch.
|
||||
import cv2
|
||||
|
||||
from remove_ai_watermarks.humanizer import _smooth_grain_mask, adaptive_polish
|
||||
|
||||
rng = np.random.default_rng(5)
|
||||
edges = rng.integers(0, 256, (120, 120, 3), dtype=np.uint8) # all high-frequency
|
||||
assert _smooth_grain_mask(edges).mean() < _smooth_grain_mask(np.full((120, 120, 3), 128, np.uint8)).mean()
|
||||
soft = cv2.GaussianBlur(edges, (0, 0), sigmaX=3.0)
|
||||
out = adaptive_polish(soft, edges, seed=0)
|
||||
assert out.shape == soft.shape
|
||||
assert out.dtype == np.uint8
|
||||
|
||||
@@ -926,6 +926,15 @@ class TestVendorOf:
|
||||
assert _vendor_of("a regular photo") is None
|
||||
assert _vendor_of(None) is None
|
||||
|
||||
def test_registered_vendors_normalize(self):
|
||||
# Regression: these registered C2PA vendors returned None, so their claims never
|
||||
# entered clash detection (a coverage hole). They now normalize to one origin.
|
||||
assert _vendor_of("ByteDance (Doubao / Jimeng / Volcano Engine)") == "ByteDance"
|
||||
assert _vendor_of("Dreamina/1.2") == "ByteDance"
|
||||
assert _vendor_of("Canva (Magic Media)") == "Canva"
|
||||
assert _vendor_of("Black Forest Labs (FLUX)") == "Black Forest Labs"
|
||||
assert _vendor_of("Eleven Labs Inc.") == "ElevenLabs"
|
||||
|
||||
|
||||
class TestIntegrityClashesHelper:
|
||||
def test_two_ai_vendors_clash(self):
|
||||
@@ -949,6 +958,31 @@ class TestIntegrityClashesHelper:
|
||||
== []
|
||||
)
|
||||
|
||||
def test_bytedance_c2pa_plus_own_aigc_no_clash(self):
|
||||
# A legit ByteDance/Doubao image carries BOTH a ByteDance C2PA manifest and its
|
||||
# own China TC260 AIGC label. The label is ByteDance's own regulatory stamp, so
|
||||
# it must be attributed to ByteDance and NOT read as a competing origin.
|
||||
assert (
|
||||
_integrity_clashes({"c2pa": "ByteDance", "aigc": "China AIGC (TC260)"}, None, camera_has_ai_marker=True)
|
||||
== []
|
||||
)
|
||||
|
||||
def test_foreign_vendor_plus_aigc_still_clashes(self):
|
||||
# But a NON-Chinese vendor's C2PA next to a China TC260 label names two different
|
||||
# origins -- a laundering tell that must still fire (the generic label stays generic).
|
||||
clashes = _integrity_clashes({"c2pa": "OpenAI", "aigc": "China AIGC (TC260)"}, None, camera_has_ai_marker=True)
|
||||
assert len(clashes) == 1
|
||||
assert "Conflicting AI-origin" in clashes[0]
|
||||
|
||||
def test_bytedance_c2pa_plus_foreign_generator_clashes(self):
|
||||
# Coverage win: a transplanted ByteDance C2PA manifest next to an independent
|
||||
# foreign generator stamp is a laundering tell that went undetected before
|
||||
# ByteDance was added to _vendor_of.
|
||||
clashes = _integrity_clashes({"c2pa": "ByteDance", "exif_generator": "OpenAI"}, None, camera_has_ai_marker=True)
|
||||
assert len(clashes) == 1
|
||||
assert "ByteDance" in clashes[0]
|
||||
assert "OpenAI" in clashes[0]
|
||||
|
||||
def test_manifest_vendor_vs_independent_signal_clashes(self):
|
||||
# A vendor named only inside the manifest still clashes with a genuinely
|
||||
# independent stamp (here an EXIF/XMP generator tag) naming a third vendor.
|
||||
|
||||
@@ -183,6 +183,17 @@ class TestQualityPreservingWrite:
|
||||
assert back is not None
|
||||
assert float(np.abs(img.astype(int) - back.astype(int)).mean()) < 1.0
|
||||
|
||||
def test_webp_written_lossless(self, tmp_path: Path) -> None:
|
||||
# Regression: cv2 WebP quality 1-100 is LOSSY; lossless needs > 100. A
|
||||
# mark-removal .webp re-encode must NOT degrade the untouched pixels, so
|
||||
# a full-frame round-trip of random data must be bit-identical.
|
||||
img = np.random.default_rng(0).integers(0, 256, (80, 80, 3), dtype=np.uint8)
|
||||
p = tmp_path / "x.webp"
|
||||
assert image_io.imwrite(p, img) is True
|
||||
back = image_io.imread(p)
|
||||
assert back is not None
|
||||
assert np.array_equal(back, img), "WebP re-encode was lossy"
|
||||
|
||||
@pytest.mark.skipif(not _heif_writable("HEIF"), reason="no HEIC encoder in this env")
|
||||
def test_heic_write_roundtrips(self, tmp_path: Path) -> None:
|
||||
# cv2 cannot encode HEIC (used to raise); imwrite must route through Pillow.
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from remove_ai_watermarks.invisible_engine import InvisibleEngine, _target_size, is_available
|
||||
@@ -195,3 +196,24 @@ class TestEsrganUpscale:
|
||||
out = InvisibleEngine._esrgan_upscale(self._fake_engine(), img, (512, 341))
|
||||
assert out.size == (512, 341)
|
||||
assert np.array_equal(np.asarray(out), np.asarray(img.resize((512, 341), Image.Resampling.LANCZOS)))
|
||||
|
||||
|
||||
class TestCannyControlImage:
|
||||
"""The ControlNet canny conditioning image builder (pure cv2/numpy; behind the gpu
|
||||
extra since it lives on WatermarkRemover). Skips when torch/diffusers are absent."""
|
||||
|
||||
def test_edge_map_is_3channel_rgb(self):
|
||||
if not is_available():
|
||||
pytest.skip("gpu extra (torch/diffusers) not installed")
|
||||
import numpy as np
|
||||
|
||||
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
|
||||
|
||||
rng = np.random.default_rng(0)
|
||||
img = Image.fromarray(rng.integers(0, 256, (64, 80, 3), dtype=np.uint8))
|
||||
# The method uses no instance state, so call it unbound with a dummy self.
|
||||
out = WatermarkRemover._build_canny_control_image(None, img) # type: ignore[arg-type]
|
||||
arr = np.array(out)
|
||||
assert out.mode == "RGB"
|
||||
assert arr.shape == (64, 80, 3)
|
||||
assert arr.max() <= 255
|
||||
|
||||
+160
-1
@@ -179,6 +179,53 @@ class TestHasAiMetadata:
|
||||
assert np.array_equal(before, after), f"{name}: pixels changed (DCT was re-encoded)"
|
||||
assert not has_ai_metadata(out), f"{name}: AI metadata survived the strip"
|
||||
|
||||
@staticmethod
|
||||
def _xmp_iptc_jpeg(tmp_path: Path, name: str, marker: bytes) -> Path:
|
||||
"""A real (decodable) JPEG carrying the IPTC AI marker in a well-formed APP1
|
||||
XMP segment -- the layout Instagram/Facebook/X and MidJourney/Meta use, where
|
||||
``digitalSourceType`` lives in XMP rather than the APP13 IPTC-IIM record.
|
||||
Synthetic (no corpus): a solid cv2 JPEG with the APP1 spliced in after SOI."""
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
real = tmp_path / f"real-{name}"
|
||||
cv2.imwrite(str(real), np.full((32, 32, 3), 200, np.uint8), [cv2.IMWRITE_JPEG_QUALITY, 100])
|
||||
data = real.read_bytes()
|
||||
xmp = (
|
||||
b"http://ns.adobe.com/xap/1.0/\x00"
|
||||
b'<x:xmpmeta xmlns:x="adobe:ns:meta/"><rdf:RDF><rdf:Description>'
|
||||
b"<Iptc4xmpExt:DigitalSourceType>http://cv.iptc.org/newscodes/digitalsourcetype/"
|
||||
+ marker
|
||||
+ b"</Iptc4xmpExt:DigitalSourceType></rdf:Description></rdf:RDF></x:xmpmeta>"
|
||||
)
|
||||
seg = b"\xff\xe1" + (len(xmp) + 2).to_bytes(2, "big") + xmp
|
||||
path = tmp_path / name
|
||||
path.write_bytes(data[:2] + seg + data[2:]) # splice APP1 right after SOI
|
||||
return path
|
||||
|
||||
@pytest.mark.parametrize("marker", [b"trainedAlgorithmicMedia", b"AISystemUsed"])
|
||||
def test_jpeg_strip_removes_iptc_marker_in_xmp(self, tmp_path: Path, marker: bytes):
|
||||
"""Regression: the lossless JPEG strip must drop an AI-bearing APP1 XMP packet
|
||||
when the AI signal is an IPTC ``digitalSourceType`` / 2025.1 field, not only a
|
||||
C2PA or China-AIGC token. Before the fix these survived because the APP1 branch
|
||||
of ``_jpeg_app_carries_ai`` checked only c2pa + AIGC markers, leaving the
|
||||
Instagram/MidJourney/Meta 'Made with AI' XMP intact. Pixels stay bit-identical."""
|
||||
import numpy as np
|
||||
|
||||
from remove_ai_watermarks import image_io
|
||||
from remove_ai_watermarks.metadata import remove_ai_metadata
|
||||
|
||||
src = self._xmp_iptc_jpeg(tmp_path, "iptc-xmp.jpg", marker)
|
||||
assert has_ai_metadata(src) # detected before
|
||||
before = image_io.imread(str(src))
|
||||
out = tmp_path / "clean.jpg"
|
||||
remove_ai_metadata(src, out)
|
||||
after = image_io.imread(str(out))
|
||||
assert before is not None
|
||||
assert after is not None
|
||||
assert np.array_equal(before, after), "pixels changed: the DCT scan was re-encoded"
|
||||
assert not has_ai_metadata(out), "IPTC AI marker in XMP survived the strip"
|
||||
|
||||
|
||||
class TestC2paMarkerIn:
|
||||
"""The C2PA presence check requires a JUMBF wrapper or the C2PA uuid box, so
|
||||
@@ -235,6 +282,59 @@ class TestSamsungGenai:
|
||||
p = self._samsung_jpeg(tmp_path, "stray.jpg", b'some other blob "genAIType":1 elsewhere')
|
||||
assert samsung_genai(p) is None
|
||||
|
||||
def test_remove_strips_post_eoi_trailer(self, tmp_path: Path):
|
||||
"""Regression: Galaxy AI appends ``PhotoEditor_Re_Edit_Data`` as a trailer AFTER
|
||||
the JPEG EOI, so the verbatim scan copy in the lossless strip carried it through
|
||||
(survived 0/8 on the corpus). The strip now truncates a Samsung AI trailer at EOI;
|
||||
pixels stay bit-identical and a non-Samsung trailer is preserved."""
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from remove_ai_watermarks import image_io
|
||||
from remove_ai_watermarks.metadata import remove_ai_metadata
|
||||
|
||||
real = tmp_path / "real.jpg"
|
||||
cv2.imwrite(str(real), np.full((32, 32, 3), 180, np.uint8), [cv2.IMWRITE_JPEG_QUALITY, 100])
|
||||
src = tmp_path / "galaxy.jpg"
|
||||
src.write_bytes(real.read_bytes() + b'PhotoEditor_Re_Edit_Data{"genAIType":1}')
|
||||
assert samsung_genai(src) == 1 # detected before
|
||||
before = image_io.imread(str(src))
|
||||
out = tmp_path / "clean.jpg"
|
||||
remove_ai_metadata(src, out)
|
||||
after = image_io.imread(str(out))
|
||||
assert before is not None
|
||||
assert after is not None
|
||||
assert np.array_equal(before, after), "pixels changed: the DCT scan was re-encoded"
|
||||
assert samsung_genai(out) is None, "Samsung genAIType trailer survived the strip"
|
||||
|
||||
def test_non_samsung_trailer_preserved(self, tmp_path: Path):
|
||||
"""A benign post-EOI trailer (e.g. an MPF block) must NOT be truncated."""
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from remove_ai_watermarks.metadata import _strip_samsung_trailer
|
||||
|
||||
real = tmp_path / "real.jpg"
|
||||
cv2.imwrite(str(real), np.full((16, 16, 3), 90, np.uint8))
|
||||
tail = real.read_bytes() + b"MPF-benign-trailer-bytes"
|
||||
assert _strip_samsung_trailer(tail) == tail
|
||||
|
||||
def test_detects_trailer_past_scan_window(self, tmp_path: Path):
|
||||
"""Regression: the marker is a trailer AFTER the JPEG EOI, so on a multi-MB
|
||||
photo it sits past the 512 KB quick-scan window. Detection must read the file
|
||||
tail too, else it disagrees with removal (which reads the whole file). A random
|
||||
1400x1400 q100 JPEG exceeds 512 KB; the marker is only in its post-EOI tail."""
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
real = tmp_path / "big.jpg"
|
||||
big = np.random.default_rng(0).integers(0, 256, (1400, 1400, 3), dtype=np.uint8)
|
||||
cv2.imwrite(str(real), big, [cv2.IMWRITE_JPEG_QUALITY, 100])
|
||||
assert real.stat().st_size > 512 * 1024 # trailer will be past the quick-scan window
|
||||
p = tmp_path / "galaxy_big.jpg"
|
||||
p.write_bytes(real.read_bytes() + b'PhotoEditor_Re_Edit_Data{"genAIType":1}')
|
||||
assert samsung_genai(p) == 1
|
||||
|
||||
def test_clean_image_is_none(self, tmp_clean_png):
|
||||
assert samsung_genai(tmp_clean_png) is None
|
||||
|
||||
@@ -300,7 +400,6 @@ class TestGetAiMetadataRealSample:
|
||||
[
|
||||
b"trainedAlgorithmicMedia",
|
||||
b"compositeSynthetic",
|
||||
b"algorithmicMedia",
|
||||
b"compositeWithTrainedAlgorithmicMedia",
|
||||
],
|
||||
)
|
||||
@@ -311,6 +410,21 @@ def test_has_ai_metadata_detects_each_iptc_marker(tmp_path: Path, marker: bytes)
|
||||
assert has_ai_metadata(path)
|
||||
|
||||
|
||||
def test_bare_algorithmic_media_not_flagged_ai(tmp_path: Path):
|
||||
"""Regression: the IPTC ``algorithmicMedia`` digitalSourceType is PROCEDURAL (an
|
||||
algorithm not trained on sampled data), NOT AI/ML generation. It must NOT be flagged
|
||||
-- flagging it made identify assert is_ai=high + has_invisible_target=True, which
|
||||
would trigger a diffusion scrub of clean procedural content. It is a distinct token
|
||||
from ``trainedAlgorithmicMedia``, so real 'Made with AI' labels are unaffected."""
|
||||
path = tmp_path / "proc.jpg"
|
||||
path.write_bytes(
|
||||
b"\xff\xd8\xff\xe1<x:xmpmeta><Iptc4xmpExt:DigitalSourceType>"
|
||||
b"http://cv.iptc.org/newscodes/digitalsourcetype/algorithmicMedia"
|
||||
b"</Iptc4xmpExt:DigitalSourceType></x:xmpmeta>\xff\xd9"
|
||||
)
|
||||
assert not has_ai_metadata(path)
|
||||
|
||||
|
||||
# ── SynthID-source detection (metadata proxy) ────────────────────────
|
||||
|
||||
|
||||
@@ -920,6 +1034,18 @@ class TestAIGCLabel:
|
||||
def test_has_ai_metadata_detects_raw_json_exif_form(self, tmp_path: Path):
|
||||
assert has_ai_metadata(self._aigc_exif_jpeg(tmp_path))
|
||||
|
||||
def test_remove_strips_raw_json_exif_form(self, tmp_path: Path):
|
||||
"""Regression: the TC260 AIGC ``{"AIGC":{...}}`` block Doubao embeds in EXIF
|
||||
UserComment must be scrubbed on removal. Before the fix it survived because
|
||||
``_scrub_ai_exif`` only touched Software/Make/Artist/ImageDescription in the
|
||||
0th IFD, never UserComment in the Exif sub-IFD."""
|
||||
from remove_ai_watermarks.metadata import aigc_label, remove_ai_metadata
|
||||
|
||||
out = tmp_path / "clean.jpg"
|
||||
remove_ai_metadata(self._aigc_exif_jpeg(tmp_path), out)
|
||||
assert aigc_label(out) is None
|
||||
assert not has_ai_metadata(out)
|
||||
|
||||
def _aigc_bare_jpeg(self, tmp_path: Path, producer: str = "00119144030008867405X210002") -> Path:
|
||||
"""Some China-served generators glue the TC260 label straight to its JSON
|
||||
as a bare ``AIGC{...}`` blob inside a JPEG APP segment (no ``"AIGC":``
|
||||
@@ -944,6 +1070,18 @@ class TestAIGCLabel:
|
||||
def test_has_ai_metadata_detects_bare_aigc_jpeg_form(self, tmp_path: Path):
|
||||
assert has_ai_metadata(self._aigc_bare_jpeg(tmp_path))
|
||||
|
||||
def test_remove_strips_bare_aigc_jpeg_form(self, tmp_path: Path):
|
||||
"""Regression: a bare ``AIGC{...}`` blob in a non-standard JPEG APP segment
|
||||
(APP9 here) is detected by aigc_label, so removal must drop that segment too.
|
||||
Before the fix ``_jpeg_app_carries_ai`` only inspected APP11/APP1-XMP/APP13, so
|
||||
the blob survived the lossless strip (detection<->removal parity break)."""
|
||||
from remove_ai_watermarks.metadata import aigc_label, remove_ai_metadata
|
||||
|
||||
out = tmp_path / "clean.jpg"
|
||||
remove_ai_metadata(self._aigc_bare_jpeg(tmp_path), out)
|
||||
assert aigc_label(out) is None
|
||||
assert not has_ai_metadata(out)
|
||||
|
||||
def test_bare_aigc_without_tc260_field_ignored(self, tmp_path: Path):
|
||||
"""A bare ``AIGC{...}`` blob with no TC260 field must not false-positive."""
|
||||
from remove_ai_watermarks.metadata import aigc_label
|
||||
@@ -1192,6 +1330,27 @@ class TestLateProvenanceBox:
|
||||
p.write_bytes(b"\x89PNG\r\n\x1a\n not an isobmff file")
|
||||
assert scan_c2pa_region(p) == b""
|
||||
|
||||
def test_scan_c2pa_region_reads_largesize_uuid(self, tmp_path: Path):
|
||||
"""A 64-bit largesize (size32 == 1) uuid box must be walked and collected."""
|
||||
import struct
|
||||
|
||||
from remove_ai_watermarks.noai.isobmff import scan_c2pa_region
|
||||
|
||||
payload = b"LARGESIZE-C2PA-MANIFEST"
|
||||
total = 16 + len(payload) # 4 (size32=1) + 4 (type) + 8 (largesize) + payload
|
||||
uuid_box = struct.pack(">I", 1) + b"uuid" + struct.pack(">Q", total) + payload
|
||||
p = tmp_path / "large.mp4"
|
||||
p.write_bytes(_MP4_FTYP + uuid_box)
|
||||
assert payload in scan_c2pa_region(p)
|
||||
|
||||
def test_scan_c2pa_region_caps_at_max_total(self, tmp_path: Path):
|
||||
"""The collected payload is bounded by ``max_total`` (never unbounded)."""
|
||||
from remove_ai_watermarks.noai.isobmff import scan_c2pa_region
|
||||
|
||||
p = tmp_path / "big.mp4"
|
||||
p.write_bytes(_MP4_FTYP + _box(b"uuid", b"A" * 5000))
|
||||
assert len(scan_c2pa_region(p, max_total=1000)) <= 1000
|
||||
|
||||
def test_front_placed_manifest_still_detected(self, tmp_path: Path):
|
||||
# Regression: a faststart MP4 (manifest before mdat) is unaffected.
|
||||
from remove_ai_watermarks.metadata import C2PA_UUID
|
||||
|
||||
@@ -321,6 +321,18 @@ class TestC2PADigitalSourceType:
|
||||
assert "compositeWithTrainedAlgorithmicMedia" in info["source_type"]
|
||||
assert "synthid_watermark" in info # AI-enhanced + OpenAI issuer
|
||||
|
||||
def test_composite_and_bare_algorithmic_cooccur_is_ai(self):
|
||||
"""Regression: a manifest carrying BOTH ``compositeWithTrainedAlgorithmicMedia``
|
||||
(AI-enhanced) and a bare procedural ``algorithmicMedia`` token must classify as
|
||||
AI-enhanced. Before the reorder the bare-token elif fired first and returned
|
||||
non-AI, dropping the composite AI signal (a false negative)."""
|
||||
from remove_ai_watermarks.noai.c2pa import _populate_registry_fields
|
||||
|
||||
info: dict = {}
|
||||
_populate_registry_fields(b"x compositeWithTrainedAlgorithmicMedia x algorithmicMedia x", info)
|
||||
assert info.get("ai_source_kind") == "enhanced"
|
||||
assert "compositeWithTrainedAlgorithmicMedia" in info["source_type"]
|
||||
|
||||
|
||||
# ── ISOBMFF (AVIF / HEIF / JPEG-XL container stripping) ──────────────
|
||||
|
||||
@@ -400,6 +412,33 @@ class TestISOBMFF:
|
||||
assert ifd[piexif.ImageIFD.Software].strip() == b""
|
||||
assert ifd[piexif.ImageIFD.Make] == b"Canon"
|
||||
|
||||
def test_blank_aigc_block_in_exif(self):
|
||||
"""Parity with the JPEG path: the China TC260 ``{"AIGC":{...}}`` block in EXIF
|
||||
ImageDescription must be blanked on the ISOBMFF path too -- ``blank_ai_exif_tokens``
|
||||
is the ONLY EXIF scrubber for HEIC/AVIF (``_scrub_ai_exif`` never runs there)."""
|
||||
import piexif
|
||||
|
||||
aigc = b'{"AIGC":{"Label":"1","ContentProducer":"00119144030008867405X210002","ProduceID":"abc"}}'
|
||||
data = self._avif_with_exif({piexif.ImageIFD.ImageDescription: aigc, piexif.ImageIFD.Make: b"Canon"})
|
||||
out, blanked = blank_ai_exif_tokens(data)
|
||||
assert blanked >= 1
|
||||
assert len(out) == len(data) # same length -> box sizes / iloc stay valid
|
||||
assert b'"AIGC"' not in out # TC260 block destroyed
|
||||
assert b"Canon" in out # camera tag preserved
|
||||
|
||||
def test_blank_xai_signature_pair_in_exif(self):
|
||||
"""Parity: the xAI/Grok ``Signature:`` blob + UUID ``Artist`` pair in EXIF is
|
||||
dropped together on the ISOBMFF path too."""
|
||||
import piexif
|
||||
|
||||
sig = b"Signature: " + b"A" * 80
|
||||
art = b"12345678-1234-1234-1234-123456789012"
|
||||
data = self._avif_with_exif({piexif.ImageIFD.ImageDescription: sig, piexif.ImageIFD.Artist: art})
|
||||
out, blanked = blank_ai_exif_tokens(data)
|
||||
assert blanked == 2 # both the signature and the UUID artist
|
||||
assert len(out) == len(data)
|
||||
assert b"Signature: AAAA" not in out
|
||||
|
||||
def test_blank_leaves_clean_exif_untouched(self):
|
||||
import piexif
|
||||
|
||||
@@ -414,6 +453,105 @@ class TestISOBMFF:
|
||||
assert out == FTYP + b"\x00\x00\x00\x0cmdat" + b"pixels!!"
|
||||
|
||||
|
||||
class TestIterTopLevelBoxes:
|
||||
"""The box walker's three size encodings and its underflow/overflow guards."""
|
||||
|
||||
def test_64bit_largesize(self):
|
||||
from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes
|
||||
|
||||
# size32 == 1 -> a 64-bit largesize follows the type; total box length = 24.
|
||||
box = struct.pack(">I", 1) + b"uuid" + struct.pack(">Q", 24) + b"payload!"
|
||||
boxes = list(_iter_top_level_boxes(box))
|
||||
assert len(boxes) == 1
|
||||
start, end, btype, payload_off = boxes[0]
|
||||
assert (start, end, btype, payload_off) == (0, 24, b"uuid", 16)
|
||||
|
||||
def test_size0_runs_to_eof(self):
|
||||
from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes
|
||||
|
||||
box = struct.pack(">I", 0) + b"mdat" + b"tail-to-eof"
|
||||
boxes = list(_iter_top_level_boxes(box))
|
||||
assert len(boxes) == 1
|
||||
start, end, btype, payload_off = boxes[0]
|
||||
assert (start, end, btype, payload_off) == (0, len(box), b"mdat", 8)
|
||||
|
||||
def test_underflow_size_stops_safely(self):
|
||||
from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes
|
||||
|
||||
# size (4) < the 8-byte header -> the guard returns without yielding a box.
|
||||
assert list(_iter_top_level_boxes(struct.pack(">I", 4) + b"ftyp" + b"more")) == []
|
||||
|
||||
def test_overflow_size_stops_safely(self):
|
||||
from remove_ai_watermarks.noai.isobmff import _iter_top_level_boxes
|
||||
|
||||
# size claims 999 but the buffer is far shorter -> guard returns, no partial box.
|
||||
assert list(_iter_top_level_boxes(struct.pack(">I", 999) + b"uuid" + b"x")) == []
|
||||
|
||||
|
||||
class TestBlankAiXmpPackets:
|
||||
"""XMP-packet blanking: same-length overwrite only for AI-marked packets, and only
|
||||
when the packet is fully delimited."""
|
||||
|
||||
AIMARK = b"trainedAlgorithmicMedia"
|
||||
|
||||
def test_ai_packet_blanked_same_length(self):
|
||||
from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets
|
||||
|
||||
packet = b'<?xpacket begin="x"?><x:xmpmeta>' + self.AIMARK + b'</x:xmpmeta><?xpacket end="w"?>'
|
||||
data = b"boxhdr" + packet + b"tail"
|
||||
out, n = blank_ai_xmp_packets(data)
|
||||
assert n == 1
|
||||
assert len(out) == len(data) # same length -> iloc offsets stay valid
|
||||
assert self.AIMARK not in out
|
||||
assert b"boxhdr" in out
|
||||
assert b"tail" in out
|
||||
|
||||
def test_clean_packet_left_intact(self):
|
||||
from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets
|
||||
|
||||
packet = b'<?xpacket begin="x"?><x:xmpmeta>plain copyright</x:xmpmeta><?xpacket end="w"?>'
|
||||
out, n = blank_ai_xmp_packets(packet)
|
||||
assert n == 0
|
||||
assert out == packet
|
||||
|
||||
def test_missing_end_delimiter_not_blanked(self):
|
||||
from remove_ai_watermarks.noai.isobmff import blank_ai_xmp_packets
|
||||
|
||||
# No <?xpacket end?> -> the packet regex cannot match, so it is left unchanged.
|
||||
data = b'<?xpacket begin="x"?><x:xmpmeta>' + self.AIMARK + b"</x:xmpmeta>"
|
||||
out, n = blank_ai_xmp_packets(data)
|
||||
assert n == 0
|
||||
assert out == data
|
||||
|
||||
|
||||
class TestC2paBufferScans:
|
||||
"""The shared buffer-scan helpers (used by both the PNG caBX parser and the
|
||||
format-agnostic binary scan). Data-driven off the registries so they stay valid
|
||||
as vendors are added."""
|
||||
|
||||
def test_soft_binding_vendors_in(self):
|
||||
from remove_ai_watermarks.noai.c2pa import C2PA_SOFT_BINDINGS, soft_binding_vendors_in
|
||||
|
||||
sig, name = next(iter(C2PA_SOFT_BINDINGS.items()))
|
||||
assert name in soft_binding_vendors_in(b"...manifest..." + sig + b"...tail...")
|
||||
assert soft_binding_vendors_in(b"") == []
|
||||
assert soft_binding_vendors_in(b"no soft-binding assertion here") == []
|
||||
|
||||
def test_synthid_vendors_in_requires_synthid_issuer(self):
|
||||
from remove_ai_watermarks.noai.c2pa import C2PA_ISSUERS, SYNTHID_C2PA_ISSUERS, synthid_vendors_in
|
||||
|
||||
syn_sig = next(s for s in C2PA_ISSUERS if s in SYNTHID_C2PA_ISSUERS)
|
||||
non_sig = next(s for s in C2PA_ISSUERS if s not in SYNTHID_C2PA_ISSUERS)
|
||||
assert C2PA_ISSUERS[syn_sig] in synthid_vendors_in(b"x" + syn_sig + b"x")
|
||||
# an issuer that does NOT pair SynthID with C2PA must not be reported as one
|
||||
assert C2PA_ISSUERS[non_sig] not in synthid_vendors_in(b"x" + non_sig + b"x")
|
||||
|
||||
def test_synthid_verdict_format(self):
|
||||
from remove_ai_watermarks.noai.c2pa import synthid_verdict
|
||||
|
||||
assert synthid_verdict("Google LLC") == "likely present (Google LLC embeds SynthID with C2PA)"
|
||||
|
||||
|
||||
class TestC2PAInvalidSignature:
|
||||
"""A .png file that is not actually PNG-signed must read as clean, not crash."""
|
||||
|
||||
|
||||
@@ -8,6 +8,40 @@ import pytest
|
||||
from remove_ai_watermarks.region_eraser import boxes_to_mask, erase, lama_available, migan_available
|
||||
|
||||
|
||||
class TestPaddedCropBox:
|
||||
"""The padded bounding box that bounds the learned backends' ONNX working set."""
|
||||
|
||||
def test_empty_mask_returns_none(self):
|
||||
from remove_ai_watermarks.region_eraser import _padded_crop_box
|
||||
|
||||
assert _padded_crop_box(np.zeros((100, 100), np.uint8), 100, 100, pad_frac=0.1, pad_min=8) is None
|
||||
|
||||
def test_pad_min_dominates_and_clamps_at_border(self):
|
||||
from remove_ai_watermarks.region_eraser import _padded_crop_box
|
||||
|
||||
mask = np.zeros((100, 100), np.uint8)
|
||||
mask[0:5, 0:5] = 255 # 5-px mark in the top-left corner
|
||||
# pad = max(8, int(0.1*5)) = 8; x0 clamps to 0 (not -8), x1 = min(100, 4+1+8) = 13.
|
||||
assert _padded_crop_box(mask, 100, 100, pad_frac=0.1, pad_min=8) == (0, 0, 13, 13)
|
||||
|
||||
def test_pad_frac_dominates_for_large_mark(self):
|
||||
from remove_ai_watermarks.region_eraser import _padded_crop_box
|
||||
|
||||
mask = np.zeros((400, 400), np.uint8)
|
||||
mask[100:300, 100:300] = 255 # 200-px span
|
||||
# pad = max(8, int(0.2*200)) = 40; box = (100-40, .., 299+1+40, ..).
|
||||
assert _padded_crop_box(mask, 400, 400, pad_frac=0.2, pad_min=8) == (60, 60, 340, 340)
|
||||
|
||||
def test_clamps_at_far_border(self):
|
||||
from remove_ai_watermarks.region_eraser import _padded_crop_box
|
||||
|
||||
mask = np.zeros((50, 60), np.uint8)
|
||||
mask[45:50, 55:60] = 255 # bottom-right corner
|
||||
_x0, _y0, x1, y1 = _padded_crop_box(mask, 50, 60, pad_frac=0.1, pad_min=8)
|
||||
assert x1 == 60 # clamped to w, no overflow
|
||||
assert y1 == 50 # clamped to h, no overflow
|
||||
|
||||
|
||||
class TestBoxesToMask:
|
||||
def test_mask_set_inside_box(self):
|
||||
mask = boxes_to_mask((100, 100), [(10, 20, 30, 40)], dilate=0)
|
||||
|
||||
@@ -49,6 +49,28 @@ class TestScan:
|
||||
dets = reg.detect_marks(np.zeros((256, 256, 3), np.uint8), include_explicit=False)
|
||||
assert not any(d.detected for d in dets)
|
||||
|
||||
@pytest.mark.parametrize("shape", [(1, 1, 3), (8, 8, 3), (15, 15, 3), (12, 300, 3), (300, 10, 3)])
|
||||
def test_tiny_image_no_crash(self, shape):
|
||||
"""Regression: an image whose short side is < 16 px (below the Gemini template
|
||||
floor) must yield no detection, not crash. detect_marks/remove_auto_marks are
|
||||
the public visible/all/batch path; a tiny thumbnail in a batch used to take the
|
||||
whole auto pass down with an IndexError (empty candidate list dereference)."""
|
||||
img = np.full(shape, 100, np.uint8)
|
||||
assert not any(d.detected for d in reg.detect_marks(img, include_explicit=False))
|
||||
result, removed = reg.remove_auto_marks(img, backend="cv2")
|
||||
assert removed == []
|
||||
assert result.shape == img.shape
|
||||
|
||||
@pytest.mark.parametrize("shape", [(0, 5), (5, 0), (0, 5, 4), (0, 0)])
|
||||
def test_forced_remove_on_empty_array_no_crash(self, shape):
|
||||
"""Regression: footprint_mask ran to_bgr (cvtColor) before any size check, so a
|
||||
forced remove on a zero-size ndarray crashed (cv2.error on an empty Mat). detect
|
||||
already guarded this; footprint_mask must too. Covers the text + gemini engines."""
|
||||
empty = np.zeros(shape, np.uint8)
|
||||
for key in ("doubao", "jimeng", "samsung", "gemini"):
|
||||
_result, mask = reg.get_mark(key).remove(empty, force=True)
|
||||
assert mask is None
|
||||
|
||||
|
||||
class TestBackendResolution:
|
||||
def test_auto_resolves_to_available_backend(self):
|
||||
|
||||
Reference in New Issue
Block a user