mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-10-14 13:30:55 +02:00
229 lines
9.0 KiB
Python
229 lines
9.0 KiB
Python
"""Regression coverage for image output lifetimes and display fidelity."""
|
|
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image, ImageCms, ImageOps
|
|
|
|
from remove_ai_watermarks import api, image_io, invisible_engine, watermark_registry
|
|
from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
|
|
|
|
|
|
def _source(path, orientation=1, alpha=False, size=(96, 64)):
|
|
width, height = size
|
|
pixels = np.zeros((height, width, 4 if alpha else 3), dtype=np.uint8)
|
|
pixels[:, : width // 3, :3] = (240, 20, 30)
|
|
pixels[: height // 2, width // 3 :, :3] = (10, 200, 40)
|
|
if alpha:
|
|
pixels[:, :, 3] = np.arange(width, dtype=np.uint8)[None, :] + 2 * np.arange(height, dtype=np.uint8)[:, None]
|
|
image = Image.fromarray(pixels)
|
|
exif = Image.Exif()
|
|
exif[274] = orientation
|
|
icc = ImageCms.ImageCmsProfile(ImageCms.createProfile("sRGB")).tobytes()
|
|
image.save(path, exif=exif, icc_profile=icc)
|
|
with Image.open(path) as reopened:
|
|
upright = np.array(ImageOps.exif_transpose(reopened))
|
|
return icc, upright
|
|
|
|
|
|
def _engine(run=None):
|
|
engine = object.__new__(invisible_engine.InvisibleEngine)
|
|
engine._progress_callback = None
|
|
remover = object.__new__(WatermarkRemover)
|
|
remover.configured_profile = "qwen-zimage"
|
|
remover.model_profile = "qwen-zimage"
|
|
remover._qwen_zimage_pipeline = SimpleNamespace(run=run or (lambda source, **kwargs: source))
|
|
engine._remover = remover
|
|
return engine
|
|
|
|
|
|
def _no_visible(monkeypatch):
|
|
monkeypatch.setattr(
|
|
watermark_registry,
|
|
"remove_auto_marks_detailed",
|
|
lambda image, **kwargs: watermark_registry.VisibleRemovalResult(image, ()),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("available", [False, True])
|
|
@pytest.mark.parametrize("in_place", [False, True])
|
|
def test_invisible_batch_passthrough_uses_current_source(tmp_path, monkeypatch, available, in_place):
|
|
source_dir = tmp_path / "in"
|
|
source_dir.mkdir()
|
|
out_dir = source_dir if in_place else tmp_path / "out"
|
|
out_dir.mkdir(exist_ok=True)
|
|
source = source_dir / "image.png"
|
|
icc, expected = _source(source, orientation=3)
|
|
original_bytes = source.read_bytes()
|
|
destination = out_dir / source.name
|
|
if not in_place:
|
|
Image.new("RGB", (96, 64), "blue").save(destination)
|
|
monkeypatch.setattr(invisible_engine, "is_available", lambda: available)
|
|
monkeypatch.setattr(api._SourceEvidence, "has_invisible_target", lambda self: False)
|
|
|
|
summary = api.remove_batch(source_dir, out_dir, mode="invisible")
|
|
|
|
assert summary.processed == 1
|
|
assert summary.failed == 0
|
|
assert summary.items[0].invisible == ("no-signal" if available else "unavailable")
|
|
if in_place:
|
|
assert source.read_bytes() == original_bytes
|
|
with Image.open(destination) as result:
|
|
np.testing.assert_array_equal(np.array(ImageOps.exif_transpose(result)), expected)
|
|
assert result.info["icc_profile"] == icc
|
|
|
|
|
|
@pytest.mark.parametrize("mode", ["visible", "all"])
|
|
@pytest.mark.parametrize("orientation", [1, 2, 3, 4, 6, 8])
|
|
@pytest.mark.parametrize("run_invisible", [False, True])
|
|
def test_batch_display_tags_and_alpha_survive_repeated_in_place_writes(
|
|
tmp_path, monkeypatch, mode, orientation, run_invisible
|
|
):
|
|
source = tmp_path / "image.png"
|
|
icc, expected = _source(source, orientation=orientation, alpha=True)
|
|
_no_visible(monkeypatch)
|
|
monkeypatch.setattr(invisible_engine, "is_available", lambda: True)
|
|
calls = []
|
|
|
|
def run(image, **kwargs):
|
|
calls.append(image.size)
|
|
return image
|
|
|
|
engine = _engine(run)
|
|
for _ in range(2):
|
|
summary = api.remove_batch(tmp_path, tmp_path, mode=mode, force=run_invisible, engine=engine)
|
|
assert summary.failed == 0, summary.errors
|
|
with Image.open(source) as result:
|
|
assert result.info.get("icc_profile") == icc
|
|
np.testing.assert_array_equal(np.array(ImageOps.exif_transpose(result)), expected)
|
|
assert len(calls) == (2 if mode == "all" and run_invisible else 0)
|
|
|
|
|
|
@pytest.mark.parametrize("explicit", [False, True])
|
|
@pytest.mark.parametrize("polish", [False, True])
|
|
def test_invisible_engine_overwrites_original_and_retains_display(tmp_path, explicit, polish):
|
|
source = tmp_path / "source.png"
|
|
icc, expected = _source(source, orientation=6)
|
|
engine = _engine()
|
|
output = engine.remove_watermark(source, source if explicit else None, adaptive_polish=polish)
|
|
assert output == source
|
|
assert output.exists()
|
|
with Image.open(output) as result:
|
|
assert result.info.get("icc_profile") == icc
|
|
assert result.getexif().get(274, 1) == 1
|
|
assert result.size == (64, 96)
|
|
if not polish:
|
|
np.testing.assert_array_equal(np.array(result), expected)
|
|
|
|
|
|
@pytest.mark.parametrize(("tile", "expected"), [(False, (512, 256)), (True, (2048, 1024))])
|
|
def test_tiling_bypasses_resolution_cap(tmp_path, tile, expected):
|
|
source = tmp_path / "source.png"
|
|
_source(source, size=(2048, 1024))
|
|
seen = []
|
|
|
|
def run(image, **kwargs):
|
|
seen.append((image.size, kwargs["tile"]))
|
|
return image
|
|
|
|
_engine(run).remove_watermark(source, tmp_path / "out.png", tile=tile, max_resolution=512)
|
|
assert seen == [(expected, tile)]
|
|
|
|
|
|
def test_invisible_postprocessing_write_failure_raises(tmp_path, monkeypatch):
|
|
source = tmp_path / "source.png"
|
|
_source(source)
|
|
original_writer = image_io.imwrite
|
|
|
|
def write(path, pixels, **kwargs):
|
|
# The model's write succeeds, the final public postprocessing write fails.
|
|
if isinstance(path, Path):
|
|
return False
|
|
return original_writer(path, pixels, **kwargs)
|
|
|
|
monkeypatch.setattr(image_io, "imwrite", write)
|
|
with pytest.raises(OSError, match="write"):
|
|
_engine().remove_watermark(source, tmp_path / "out.png", adaptive_polish=True)
|
|
|
|
|
|
@pytest.mark.parametrize("orientation", [2, 3, 4, 5, 6, 7, 8])
|
|
def test_explicit_upright_writer_does_not_reapply_square_image_orientation(tmp_path, orientation):
|
|
source = tmp_path / "source.png"
|
|
icc, upright = _source(source, orientation=orientation, size=(64, 64))
|
|
assert image_io.imwrite(source, upright[:, :, ::-1], display_tags_from=source, orientation_applied=True)
|
|
with Image.open(source) as result:
|
|
assert result.info.get("icc_profile") == icc
|
|
np.testing.assert_array_equal(np.array(ImageOps.exif_transpose(result)), upright)
|
|
|
|
|
|
@pytest.mark.parametrize("failure", [None, "save", "model"])
|
|
def test_invisible_temporary_input_is_cleaned_after_success_and_failure(tmp_path, monkeypatch, failure):
|
|
import tempfile
|
|
|
|
source = tmp_path / "source.png"
|
|
_source(source)
|
|
original_bytes = source.read_bytes()
|
|
temporary_paths = []
|
|
original_mkstemp = tempfile.mkstemp
|
|
|
|
def make_temp(**kwargs):
|
|
fd, name = original_mkstemp(dir=tmp_path, **kwargs)
|
|
temporary_paths.append(Path(name))
|
|
return fd, name
|
|
|
|
monkeypatch.setattr(tempfile, "mkstemp", make_temp)
|
|
original_save = Image.Image.save
|
|
|
|
def save(image, fp, **kwargs):
|
|
if failure == "save" and Path(fp) in temporary_paths:
|
|
raise OSError("synthetic temp save failure")
|
|
return original_save(image, fp, **kwargs)
|
|
|
|
monkeypatch.setattr(Image.Image, "save", save)
|
|
|
|
def run(image, **kwargs):
|
|
if failure == "model":
|
|
raise OSError("synthetic model failure")
|
|
return image
|
|
|
|
engine = _engine(run)
|
|
destination = tmp_path / "out.png"
|
|
if failure:
|
|
with pytest.raises(OSError, match="synthetic"):
|
|
engine.remove_watermark(source, destination)
|
|
assert not destination.exists()
|
|
else:
|
|
assert engine.remove_watermark(source, destination) == destination
|
|
assert destination.exists()
|
|
assert source.read_bytes() == original_bytes
|
|
assert len(temporary_paths) == 1
|
|
assert not temporary_paths[0].exists()
|
|
|
|
|
|
def test_invisible_batch_preserves_new_model_result(tmp_path, monkeypatch):
|
|
source_dir = tmp_path / "in"
|
|
source_dir.mkdir()
|
|
source = source_dir / "image.png"
|
|
_source(source)
|
|
monkeypatch.setattr(invisible_engine, "is_available", lambda: True)
|
|
engine = _engine(lambda image, **kwargs: Image.new("RGB", image.size, (20, 30, 220)))
|
|
summary = api.remove_batch(source_dir, tmp_path / "out", mode="invisible", force=True, engine=engine)
|
|
assert summary.failed == 0
|
|
assert summary.items[0].invisible == "removed"
|
|
with Image.open(tmp_path / "out" / source.name) as output:
|
|
assert output.getpixel((0, 0)) == (20, 30, 220)
|
|
|
|
|
|
@pytest.mark.parametrize("orientation", [3, 6])
|
|
def test_single_visible_format_conversion_preserves_display_tags(tmp_path, monkeypatch, orientation):
|
|
source = tmp_path / "source.png"
|
|
icc, expected = _source(source, orientation=orientation, size=(64, 64))
|
|
_no_visible(monkeypatch)
|
|
destination = tmp_path / "out.webp"
|
|
api.remove_visible(source, destination, strip_metadata=False)
|
|
with Image.open(destination) as result:
|
|
assert result.info.get("icc_profile") == icc
|
|
np.testing.assert_array_equal(np.array(ImageOps.exif_transpose(result)), expected)
|