mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-06-10 12:53:56 +02:00
Victor Kuznetsov
2fcd00ced0
Nine findings from a high-effort project-wide review, fixed and verified (571 passed, ruff/pyright clean): Correctness: - all/batch now remove Doubao/Jimeng/Samsung visible text marks: the visible step routes through the registry (new cli._remove_visible_auto) instead of a hardcoded GeminiEngine, so they no longer leave the wordmark intact. - batch always reads the original source (dropped the out_path-reuse that re-processed already-cleaned outputs on a re-run). - img2img_runner only retries the diffusion call on the deprecated-callback TypeError; any other TypeError now propagates instead of double-running. - gemini detect/remove and the reverse-alpha engines normalize channels via a new image_io.to_bgr, fixing a grayscale/BGRA crash in the FP-gate path. - _png_late_metadata advances its cursor by the clamped length, so a malformed chunk length no longer aborts the late AI-label scan. Cleanup / efficiency: - Consolidate the ~90%-identical Doubao/Jimeng/Samsung engines into a shared config-driven _text_mark_engine.TextMarkEngine base; each engine is now a thin subclass (TextMarkConfig + test shims). Behavior is byte-exact (the three engine test suites pass unchanged). Registry adapters collapse to one _text_mark(...) row each. Gemini stays a separate engine. - scan_head is memoized per (path, size, mtime), so identify() reads the file head once instead of ~8 times. - invisible_engine post-processing decodes/encodes the output once (chained in memory) instead of 2-4 times across stages. - Remove the orphaned get_model_id_for_profile (+ CONTROLNET_PROFILE); derive the --strength help from the strength constants (strength_default_help) so it cannot drift; share the --pipeline/--strength click options; simplify the retired --auto resolver. Net -835 lines. Tests added for the registry-routed visible pass, to_bgr, the polish/model/guidance wiring, and strength_default_help. CLAUDE.md updated for the new base module, the engine/registry changes, image_io.to_bgr, and the scan_head cache. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
106 lines
3.7 KiB
Python
106 lines
3.7 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 TestToBgr:
|
|
def test_grayscale_2d_promoted_to_bgr(self) -> None:
|
|
gray = np.full((4, 5), 120, dtype=np.uint8)
|
|
out = image_io.to_bgr(gray)
|
|
assert out.shape == (4, 5, 3)
|
|
# GRAY2BGR replicates the channel, so all three match the source.
|
|
assert np.array_equal(out[..., 0], gray)
|
|
assert np.array_equal(out[..., 0], out[..., 2])
|
|
|
|
def test_single_channel_3d_promoted(self) -> None:
|
|
gray = np.full((4, 5, 1), 7, dtype=np.uint8)
|
|
assert image_io.to_bgr(gray).shape == (4, 5, 3)
|
|
|
|
def test_bgra_dropped_to_bgr(self) -> None:
|
|
bgra = np.zeros((4, 5, 4), dtype=np.uint8)
|
|
bgra[..., :3] = (10, 120, 240)
|
|
out = image_io.to_bgr(bgra)
|
|
assert out.shape == (4, 5, 3)
|
|
assert np.array_equal(out, bgra[..., :3])
|
|
|
|
def test_bgr_returned_unchanged(self) -> None:
|
|
bgr = _make_bgr()
|
|
out = image_io.to_bgr(bgr)
|
|
assert out is bgr # 3-channel: no copy
|
|
|
|
|
|
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
|