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>
80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
"""Tests for Unicode-safe cv2 image IO (issue #17)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from remove_ai_watermarks import image_io
|
|
|
|
if TYPE_CHECKING:
|
|
from pathlib import Path
|
|
|
|
# Non-ASCII filenames that break cv2.imread/imwrite on Windows (issue #17).
|
|
_UNICODE_NAMES = [
|
|
"jimeng-2026-05-27-一面白色的墙.png", # Chinese
|
|
"тест-изображение.png", # Cyrillic
|
|
"café-señor.png", # accented Latin
|
|
]
|
|
|
|
|
|
def _make_bgr() -> np.ndarray:
|
|
img = np.zeros((8, 8, 3), dtype=np.uint8)
|
|
img[2:6, 2:6] = (10, 120, 240) # a BGR block so the round-trip is checkable
|
|
return img
|
|
|
|
|
|
class TestUnicodeRoundTrip:
|
|
def test_write_then_read_preserves_pixels(self, tmp_path: Path) -> None:
|
|
for name in _UNICODE_NAMES:
|
|
path = tmp_path / name
|
|
src = _make_bgr()
|
|
assert image_io.imwrite(path, src) is True
|
|
assert path.exists()
|
|
out = image_io.imread(path)
|
|
assert out is not None
|
|
# PNG is lossless: pixels must match exactly.
|
|
assert np.array_equal(out, src)
|
|
|
|
def test_alpha_round_trip_with_unchanged_flag(self, tmp_path: Path) -> None:
|
|
path = tmp_path / "豆包-alpha.png"
|
|
bgra = np.zeros((8, 8, 4), dtype=np.uint8)
|
|
bgra[..., 3] = 128
|
|
assert image_io.imwrite(path, bgra) is True
|
|
out = image_io.imread(path, cv2.IMREAD_UNCHANGED)
|
|
assert out is not None
|
|
assert out.shape[2] == 4
|
|
assert np.array_equal(out, bgra)
|
|
|
|
def test_reads_file_written_by_raw_cv2(self, tmp_path: Path) -> None:
|
|
# An ASCII file written by plain cv2 must read back identically through
|
|
# the wrapper (decode path is byte-compatible with cv2.imread).
|
|
path = tmp_path / "ascii.png"
|
|
src = _make_bgr()
|
|
cv2.imwrite(str(path), src)
|
|
out = image_io.imread(path)
|
|
assert out is not None
|
|
assert np.array_equal(out, src)
|
|
|
|
|
|
class TestFailureSemantics:
|
|
def test_missing_file_returns_none(self, tmp_path: Path) -> None:
|
|
assert image_io.imread(tmp_path / "does-not-exist-不存在.png") is None
|
|
|
|
def test_empty_file_returns_none(self, tmp_path: Path) -> None:
|
|
path = tmp_path / "empty.png"
|
|
path.write_bytes(b"")
|
|
assert image_io.imread(path) is None
|
|
|
|
def test_undecodable_file_returns_none(self, tmp_path: Path) -> None:
|
|
path = tmp_path / "garbage.png"
|
|
path.write_bytes(b"not an image")
|
|
assert image_io.imread(path) is None
|
|
|
|
def test_imwrite_to_missing_directory_returns_false(self, tmp_path: Path) -> None:
|
|
# An unwritable path must return False (cv2.imwrite contract), not raise.
|
|
path = tmp_path / "no-such-dir" / "out.png"
|
|
assert image_io.imwrite(path, _make_bgr()) is False
|