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:
Victor Kuznetsov
2026-07-13 10:49:24 +03:00
co-authored by Claude Opus 4.8
parent 190dc89d23
commit a4c901ff39
25 changed files with 894 additions and 97 deletions
+73
View File
@@ -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
+26
View File
@@ -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
+26
View File
@@ -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
+34
View File
@@ -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.
+11
View File
@@ -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.
+22
View File
@@ -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
View File
@@ -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
+138
View File
@@ -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."""
+34
View File
@@ -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)
+22
View File
@@ -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):