mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-06-05 10:38:00 +02:00
Victor Kuznetsov
5d0e6c3a65
Apply fixes from a full-repo review (code, tests, docs). Security / correctness: - Clamp attacker-controlled PNG/caBX chunk lengths to the remaining file size in metadata.py and noai/c2pa.py (a malformed length no longer drives a multi-GB read); skipped chunks seek instead of read. - noai/isobmff.strip_c2pa_boxes is now fail-safe on a malformed box: return the original bytes with a warning instead of silently truncating the tail, so metadata --remove can no longer emit a corrupt file. - doubao_engine._fixed_alpha_map clamps the glyph box to the image (no crash on degenerate width-vs-height). - watermark_remover._run_region_hires gates the phaseCorrelate offset on response and magnitude (a spurious shift no longer garbles text) and drops the generator after a CPU fallback (no MPS/CPU device mismatch). Robustness: - gemini_engine, doubao_engine, region_eraser normalize grayscale and RGBA inputs to BGR at the engine entry points. - image_io.imwrite returns False on an unwritable path (matches cv2). - invisible_engine guards a None imread result before use. - trustmark_detector._decoder uses a double-checked threading lock. - ctrlregen.tiling.tile_positions raises on overlap >= tile. - humanizer chromatic shift no longer wraps opposite-edge pixels. - identify OpenAI caveat keyed on the normalized vendor, not a substring. - Remove the dead "visible --detect-threshold" CLI option. - publish.yml verifies the release tag matches the package version. Docs: - README strength 0.05 to 0.10; .env.example HF_TOKEN marked optional; doubao_capture README updated to reverse-alpha-only; CLAUDE.md synced with the new behaviors and the batch command. Tests: new test_security_clamp.py for the read clamp and isobmff fail-safe; erase CLI coverage; integrity-clash rule 2 end-to-end; multi-tag EXIF survival and cross-format strip guards; channel/size, tiling, humanizer, and imwrite regressions. Full suite 493 passed, 2 skipped; ruff and pyright src/ clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
93 lines
3.5 KiB
Python
93 lines
3.5 KiB
Python
"""Tests for the universal region eraser."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from remove_ai_watermarks.region_eraser import boxes_to_mask, erase, lama_available
|
|
|
|
|
|
class TestBoxesToMask:
|
|
def test_mask_set_inside_box(self):
|
|
mask = boxes_to_mask((100, 100), [(10, 20, 30, 40)], dilate=0)
|
|
assert mask[25, 15] == 255 # inside
|
|
assert mask[0, 0] == 0 # outside
|
|
assert mask.shape == (100, 100)
|
|
|
|
def test_multiple_boxes(self):
|
|
mask = boxes_to_mask((100, 100), [(0, 0, 10, 10), (90, 90, 10, 10)], dilate=0)
|
|
assert mask[5, 5] == 255
|
|
assert mask[95, 95] == 255
|
|
assert mask[50, 50] == 0
|
|
|
|
def test_dilate_grows_mask(self):
|
|
m0 = boxes_to_mask((100, 100), [(40, 40, 10, 10)], dilate=0)
|
|
m5 = boxes_to_mask((100, 100), [(40, 40, 10, 10)], dilate=5)
|
|
assert m5.sum() > m0.sum()
|
|
|
|
def test_box_clipped_to_bounds(self):
|
|
# box partly outside the image must not raise and stays in-bounds
|
|
mask = boxes_to_mask((50, 50), [(40, 40, 100, 100)], dilate=0)
|
|
assert mask[45, 45] == 255
|
|
|
|
|
|
class TestEraseCv2:
|
|
def _image_with_logo(self) -> tuple[np.ndarray, tuple[int, int, int, int]]:
|
|
img = np.full((200, 200, 3), 120, np.uint8) # flat gray background
|
|
box = (140, 160, 50, 30)
|
|
x, y, w, h = box
|
|
img[y : y + h, x : x + w] = (255, 255, 255) # bright "logo"
|
|
return img, box
|
|
|
|
def test_erase_changes_region(self):
|
|
img, box = self._image_with_logo()
|
|
out = erase(img, boxes=[box], backend="cv2")
|
|
x, y, w, h = box
|
|
# on a flat background the logo region should be repainted near gray
|
|
region = out[y : y + h, x : x + w]
|
|
assert abs(float(region.mean()) - 120) < 20
|
|
assert not np.array_equal(out, img)
|
|
|
|
def test_pixels_outside_box_untouched(self):
|
|
img, box = self._image_with_logo()
|
|
out = erase(img, boxes=[box], backend="cv2", dilate=0)
|
|
# a far corner must be identical
|
|
assert np.array_equal(img[:50, :50], out[:50, :50])
|
|
|
|
def test_no_boxes_returns_copy(self):
|
|
img = np.full((100, 100, 3), 50, np.uint8)
|
|
out = erase(img, boxes=[], backend="cv2")
|
|
assert np.array_equal(img, out)
|
|
|
|
def test_empty_mask_returns_copy(self):
|
|
img = np.full((100, 100, 3), 50, np.uint8)
|
|
out = erase(img, mask=np.zeros((100, 100), np.uint8), backend="cv2")
|
|
assert np.array_equal(img, out)
|
|
|
|
|
|
class TestNonBgrInputs:
|
|
"""cv2.inpaint rejects 4-channel BGRA and 2D-only entry points must work."""
|
|
|
|
def test_grayscale_2d_does_not_raise(self):
|
|
gray = np.full((100, 100), 120, np.uint8)
|
|
out = erase(gray, boxes=[(40, 40, 20, 20)], backend="cv2")
|
|
assert out.shape == gray.shape
|
|
|
|
def test_bgra_preserves_alpha_and_does_not_raise(self):
|
|
bgra = np.full((100, 100, 4), 120, np.uint8)
|
|
bgra[..., 3] = 200 # opaque-ish alpha plane
|
|
out = erase(bgra, boxes=[(40, 40, 20, 20)], backend="cv2", dilate=0)
|
|
assert out.shape == bgra.shape
|
|
# alpha plane is carried through unchanged
|
|
assert np.array_equal(out[..., 3], bgra[..., 3])
|
|
|
|
|
|
class TestLamaBackend:
|
|
def test_lama_raises_when_unavailable(self):
|
|
img = np.full((100, 100, 3), 50, np.uint8)
|
|
if lama_available():
|
|
pytest.skip("onnxruntime installed; cannot test the unavailable path")
|
|
with pytest.raises(RuntimeError, match="onnxruntime"):
|
|
erase(img, boxes=[(10, 10, 20, 20)], backend="lama")
|