mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-19 12:07:13 +02:00
129 lines
4.1 KiB
Python
129 lines
4.1 KiB
Python
"""Verified-text manifest and compositor tests without model downloads."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image, PngImagePlugin
|
|
|
|
from remove_ai_watermarks._internal.text_restoration import (
|
|
FIDELITY_BLEND_ALPHA,
|
|
VerifiedTextLine,
|
|
blend_fidelity_anchor,
|
|
load_verified_text_manifest,
|
|
restore_verified_text,
|
|
source_pixel_sha256,
|
|
)
|
|
|
|
|
|
def _manifest(image: Image.Image) -> dict[str, object]:
|
|
return {
|
|
"schema_version": 1,
|
|
"verified": True,
|
|
"source_pixel_sha256": source_pixel_sha256(image),
|
|
"width": image.width,
|
|
"height": image.height,
|
|
"lines": [
|
|
{
|
|
"box": [8, 8, 40, 24],
|
|
"text": "Exact text",
|
|
"script": "alphabetic",
|
|
"angle": 0.0,
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
def test_pixel_hash_ignores_container_metadata(tmp_path) -> None:
|
|
image = Image.new("RGB", (48, 32), (10, 20, 30))
|
|
plain = tmp_path / "plain.png"
|
|
tagged = tmp_path / "tagged.png"
|
|
image.save(plain)
|
|
metadata = PngImagePlugin.PngInfo()
|
|
metadata.add_text("note", "different container bytes")
|
|
image.save(tagged, pnginfo=metadata)
|
|
|
|
with Image.open(plain) as left, Image.open(tagged) as right:
|
|
assert plain.read_bytes() != tagged.read_bytes()
|
|
assert source_pixel_sha256(left) == source_pixel_sha256(right)
|
|
|
|
|
|
def test_verified_manifest_is_bound_to_source_pixels(tmp_path) -> None:
|
|
source = Image.new("RGB", (48, 32), (10, 20, 30))
|
|
path = tmp_path / "lines.json"
|
|
path.write_text(json.dumps(_manifest(source)), encoding="utf-8")
|
|
|
|
loaded = load_verified_text_manifest(path, source)
|
|
|
|
assert loaded.width == 48
|
|
assert loaded.height == 32
|
|
assert loaded.lines == (VerifiedTextLine((8, 8, 40, 24), "Exact text", "alphabetic", 0.0),)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("mutation", "message"),
|
|
[
|
|
({"verified": False}, "verified=true"),
|
|
({"source_pixel_sha256": "0" * 64}, "does not match"),
|
|
({"width": 49}, "dimensions"),
|
|
({"lines": []}, "non-empty"),
|
|
],
|
|
)
|
|
def test_manifest_rejects_unverified_or_unbound_input(tmp_path, mutation, message) -> None:
|
|
source = Image.new("RGB", (48, 32), (10, 20, 30))
|
|
payload = _manifest(source)
|
|
payload.update(mutation)
|
|
path = tmp_path / "lines.json"
|
|
path.write_text(json.dumps(payload), encoding="utf-8")
|
|
|
|
with pytest.raises(ValueError, match=message):
|
|
load_verified_text_manifest(path, source)
|
|
|
|
|
|
def test_fidelity_anchor_uses_the_calibrated_rounding() -> None:
|
|
clean = Image.fromarray(np.array([[[1, 2, 3], [100, 150, 200]]], dtype=np.uint8))
|
|
donor = Image.fromarray(np.array([[[255, 254, 253], [200, 100, 50]]], dtype=np.uint8))
|
|
|
|
result = np.asarray(blend_fidelity_anchor(clean, donor))
|
|
expected = np.rint(
|
|
np.asarray(clean, dtype=np.float32) * (1.0 - FIDELITY_BLEND_ALPHA)
|
|
+ np.asarray(donor, dtype=np.float32) * FIDELITY_BLEND_ALPHA
|
|
).astype(np.uint8)
|
|
|
|
assert np.array_equal(result, expected)
|
|
|
|
|
|
def test_restoration_uses_lama_and_qwen_vae_core(monkeypatch) -> None:
|
|
from remove_ai_watermarks import region_eraser
|
|
|
|
source = np.full((40, 64, 3), 20, dtype=np.uint8)
|
|
source[12:24, 12:44] = 235
|
|
candidate = np.full_like(source, 30)
|
|
candidate[12:24, 12:44] = 150
|
|
donor = np.full_like(source, 40)
|
|
donor[12:24, 12:44] = (210, 220, 230)
|
|
calls: list[np.ndarray] = []
|
|
|
|
def fake_erase(image_bgr, mask):
|
|
calls.append(mask.copy())
|
|
output = image_bgr.copy()
|
|
output[mask > 0] = (30, 30, 30)
|
|
return output
|
|
|
|
monkeypatch.setattr(region_eraser, "lama_available", lambda: True)
|
|
monkeypatch.setattr(region_eraser, "erase_lama", fake_erase)
|
|
|
|
result = restore_verified_text(
|
|
Image.fromarray(source),
|
|
Image.fromarray(candidate),
|
|
Image.fromarray(donor),
|
|
(VerifiedTextLine((8, 8, 48, 28), "Exact text", "alphabetic"),),
|
|
)
|
|
|
|
restored = np.asarray(result)
|
|
assert calls
|
|
assert np.all(restored[16, 20] == donor[16, 20])
|
|
assert np.all(restored[0, 0] == candidate[0, 0])
|