mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-06-05 02:28:00 +02:00
7b47fa9f6a
Two CLI/IO robustness bugs surfaced by issues #17 and #19. #17 -- non-ASCII image paths (Chinese/Cyrillic/accented) failed on Windows: cv2.imread/imwrite use the platform ANSI code-page API, so the decode came back empty with a "can't open/read file" warning. New image_io.imread/imwrite route through np.fromfile+cv2.imdecode / cv2.imencode+tofile (Unicode-safe, byte- identical output, cv2.imread None-semantics preserved); all 8 cv2 read/write call sites now go through it. Behavior-neutral on macOS/Linux (already accept UTF-8 paths), so the fix is correct-by-construction for the Windows-only bug. #19 (incidental) -- rich parsed the "[gpu]" in the GPU-extra install hint as a style tag and dropped it, so the printed command was the un-installable "pip install 'remove-ai-watermarks'". Escaped as \[gpu] at both call sites. Tests: test_image_io.py (non-ASCII round-trip, alpha, missing/empty/garbage semantics); test_cli.py::TestGpuHintMarkup (install hint keeps the extra). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
75 lines
2.5 KiB
Python
75 lines
2.5 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
|