Files
remove-ai-watermarks/tests/test_image_io.py
T
test-user 7b47fa9f6a fix(io): Unicode-safe cv2 image IO + un-eat the [gpu] install hint (v0.6.6)
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>
2026-05-27 11:52:48 -07:00

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