Add verified text restoration

tests/test_api.py

@@ -245,6 +245,7 @@ class TestInvisibleOptionsMirrorTheEngine:
             tile=True,
             tile_size=768,
             tile_overlap=64,
+            text_manifest=tmp_path / "verified-lines.json",
         )
         seen: dict[str, object] = {}
 

tests/test_invisible_engine.py

@@ -43,6 +43,76 @@ class TestInvisibleEngineInit:
         assert engine._preload_kwargs == {"global_only": True}
 
 
+class TestVerifiedTextMode:
+    """The experimental mode must fail before loading models on unmeasured inputs."""
+
+    @staticmethod
+    def _engine(profile: str = "qwen-zimage") -> InvisibleEngine:
+        engine = object.__new__(InvisibleEngine)
+        engine._progress_callback = None
+        engine._remover = SimpleNamespace(model_profile=profile)
+        return engine
+
+    def test_rejects_incompatible_pipeline_options(self, tmp_path):
+        import pytest
+
+        manifest = tmp_path / "manifest.json"
+        manifest.write_text("{}", encoding="utf-8")
+        cases = (
+            ("sdxl-zimage", {}, "qwen-zimage"),
+            ("qwen-zimage", {"max_resolution": 1024}, "max-resolution 0"),
+            ("qwen-zimage", {"tile": True}, "not calibrated"),
+            ("qwen-zimage", {"humanize": 1.0}, "humanize=0"),
+            ("qwen-zimage", {"adaptive_polish": True}, "polish disabled"),
+        )
+        for profile, kwargs, message in cases:
+            with pytest.raises(ValueError, match=message):
+                self._engine(profile).remove_watermark(
+                    tmp_path / "unused.png",
+                    text_manifest=manifest,
+                    **kwargs,
+                )
+
+    def test_loads_and_forwards_verified_manifest(self, tmp_path, monkeypatch):
+        import json
+
+        from remove_ai_watermarks import region_eraser
+        from remove_ai_watermarks._internal.text_restoration import source_pixel_sha256
+
+        source = tmp_path / "source.png"
+        output = tmp_path / "output.png"
+        image = Image.new("RGB", (48, 32), (10, 20, 30))
+        image.save(source)
+        manifest = tmp_path / "manifest.json"
+        manifest.write_text(
+            json.dumps(
+                {
+                    "schema_version": 1,
+                    "verified": True,
+                    "source_pixel_sha256": source_pixel_sha256(image),
+                    "width": 48,
+                    "height": 32,
+                    "lines": [{"box": [8, 8, 40, 24], "text": "Exact", "script": "alphabetic"}],
+                }
+            ),
+            encoding="utf-8",
+        )
+        seen = {}
+
+        def fake_remove(**kwargs):
+            seen.update(kwargs)
+            Image.open(kwargs["image_path"]).save(kwargs["output_path"])
+            return kwargs["output_path"]
+
+        engine = self._engine()
+        engine._remover.remove_watermark = fake_remove
+        monkeypatch.setattr(region_eraser, "lama_available", lambda: True)
+
+        engine.remove_watermark(source, output, text_manifest=manifest)
+
+        assert seen["text_manifest"].lines[0].text == "Exact"
+
+
 class TestNativeOutputSize:
     """Model-side latent-grid rounding must not change the public output size."""
 

tests/test_qwen_zimage_pipeline.py

@@ -579,6 +579,78 @@ def test_cli_qwen_zimage_keeps_profile_postprocess_default(tmp_image_path, monke
     assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is True
 
 
+def test_cli_forwards_verified_text_manifest(tmp_image_path, tmp_path, monkeypatch):
+    from remove_ai_watermarks import cli
+
+    manifest = tmp_path / "manifest.json"
+    manifest.write_text("{}", encoding="utf-8")
+    mock_engine = MagicMock()
+    mock_engine.remove_watermark.return_value = tmp_image_path
+    monkeypatch.setattr("remove_ai_watermarks.invisible_engine.is_available", lambda: True)
+    monkeypatch.setattr("remove_ai_watermarks.invisible_engine.InvisibleEngine", MagicMock(return_value=mock_engine))
+
+    result = CliRunner().invoke(
+        cli.main,
+        ["invisible", str(tmp_image_path), "--text-manifest", str(manifest), "--force"],
+    )
+
+    assert result.exit_code == 0, result.output
+    assert mock_engine.remove_watermark.call_args.kwargs["text_manifest"] == manifest
+
+
+def test_cli_reports_verified_text_manifest_errors(tmp_image_path, tmp_path, monkeypatch):
+    from remove_ai_watermarks import cli
+
+    manifest = tmp_path / "manifest.json"
+    manifest.write_text("{}", encoding="utf-8")
+    mock_engine = MagicMock()
+    mock_engine.remove_watermark.side_effect = ValueError("manifest pixels do not match")
+    monkeypatch.setattr("remove_ai_watermarks.invisible_engine.is_available", lambda: True)
+    monkeypatch.setattr("remove_ai_watermarks.invisible_engine.InvisibleEngine", MagicMock(return_value=mock_engine))
+
+    result = CliRunner().invoke(
+        cli.main,
+        ["invisible", str(tmp_image_path), "--text-manifest", str(manifest), "--force"],
+    )
+
+    assert result.exit_code == 1
+    assert "manifest pixels do not match" in result.output
+
+
+def test_no_face_path_still_runs_verified_text_restoration(monkeypatch):
+    from remove_ai_watermarks._internal import qwen_zimage_pipeline, text_restoration
+    from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
+    from remove_ai_watermarks._internal.text_restoration import VerifiedTextLine, VerifiedTextManifest
+
+    pipeline = object.__new__(QwenZImagePipeline)
+    pipeline.device = "cuda"
+    pipeline.progress_callback = None
+    source = Image.new("RGB", (32, 32), (10, 20, 30))
+    donor = Image.new("RGB", (32, 32), (40, 50, 60))
+    global_result = Image.new("RGB", (32, 32), (70, 80, 90))
+    anchor = Image.new("RGB", (32, 32), (100, 110, 120))
+    restored = Image.new("RGB", (32, 32), (130, 140, 150))
+    pipeline._qwen_vae_roundtrip = MagicMock(return_value=donor)
+    pipeline._run_global = MagicMock(return_value=global_result)
+    monkeypatch.setattr(qwen_zimage_pipeline, "detect_faces", lambda _image: [])
+    blend = MagicMock(return_value=anchor)
+    restore = MagicMock(return_value=restored)
+    monkeypatch.setattr(text_restoration, "blend_fidelity_anchor", blend)
+    monkeypatch.setattr(text_restoration, "restore_verified_text", restore)
+    manifest = VerifiedTextManifest(
+        "0" * 64,
+        32,
+        32,
+        (VerifiedTextLine((4, 4, 20, 16), "Exact", "alphabetic"),),
+    )
+
+    result = pipeline.run(source, strength=0.1, seed=0, text_manifest=manifest)
+
+    assert result is restored
+    blend.assert_called_once_with(global_result, donor)
+    restore.assert_called_once_with(source, anchor, donor, manifest.lines)
+
+
 def test_watermark_remover_dispatches_to_full_pipeline(tmp_path, monkeypatch):
     from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
 
@@ -601,6 +673,7 @@ def test_watermark_remover_dispatches_to_full_pipeline(tmp_path, monkeypatch):
     _, kwargs = runtime.run.call_args
     assert kwargs["strength"] == pytest.approx(0.084)
     assert kwargs["seed"] == 0
+    assert kwargs["text_manifest"] is None
     assert output.exists()
 
 

tests/test_selective_text_restoration.py

@@ -8,6 +8,8 @@ from pathlib import Path
 import numpy as np
 import pytest
 
+from remove_ai_watermarks._internal import text_restoration
+
 SCRIPT = Path(__file__).parents[1] / "scripts" / "selective_text_restoration.py"
 SPEC = importlib.util.spec_from_file_location("selective_text_restoration", SCRIPT)
 assert SPEC is not None
@@ -87,7 +89,7 @@ def test_fresh_silhouette_uses_new_color_instead_of_source_pixels() -> None:
     mask = np.zeros((9, 9), dtype=np.uint8)
     mask[3:6, 3:6] = 255
 
-    result = module.composite_fresh_silhouette(background, mask, (220, 180, 40), feather=0)
+    result = text_restoration.composite_fresh_silhouette(background, mask, (220, 180, 40), feather=0)
 
     assert np.all(result[3:6, 3:6] == (220, 180, 40))
     np.testing.assert_array_equal(result[0, 0], background[0, 0])
@@ -99,7 +101,7 @@ def test_fresh_silhouette_antialiasing_softens_binary_edges() -> None:
     mask = np.zeros((9, 9), dtype=np.uint8)
     mask[3:6, 3:6] = 255
 
-    result = module.composite_fresh_silhouette(background, mask, (220, 180, 40), feather=1.0)
+    result = text_restoration.composite_fresh_silhouette(background, mask, (220, 180, 40), feather=1.0)
 
     assert np.all(result[3, 3] > background[3, 3])
     assert np.all(result[3, 3] < (220, 180, 40))
@@ -113,7 +115,7 @@ def test_reconstructed_glyphs_keep_exact_donor_core_and_fresh_edge() -> None:
     mask = np.zeros((9, 9), dtype=np.uint8)
     mask[3:6, 3:6] = 255
 
-    fresh_edge = module.composite_fresh_silhouette(background, mask, (220, 180, 40))
+    fresh_edge = text_restoration.composite_fresh_silhouette(background, mask, (220, 180, 40))
     result = module.composite_reconstructed_glyphs(donor, fresh_edge, mask, feather=0.5)
 
     np.testing.assert_array_equal(result[3:6, 3:6], donor[3:6, 3:6])
@@ -163,13 +165,15 @@ def test_detect_line_boxes_fails_closed_on_count_mismatch() -> None:
 
 
 def test_residual_mask_is_limited_to_original_glyph_positions(monkeypatch) -> None:
+    from remove_ai_watermarks._internal import text_restoration
+
     background = np.zeros((8, 8, 3), dtype=np.uint8)
     original = np.zeros((8, 8), dtype=np.uint8)
     original[3, 3] = 255
     detected = np.zeros((8, 8), dtype=np.uint8)
     detected[3, 3] = 255
     detected[6, 6] = 255
-    monkeypatch.setattr(module, "foreground_mask", lambda _image, _box: detected)
+    monkeypatch.setattr(text_restoration, "_foreground_mask", lambda _image, _box: detected)
 
     residual = module.residual_glyph_mask(background, original, (0, 0, 8, 8))
 

tests/test_text_restoration.py

@@ -0,0 +1,128 @@
+"""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])