From 7b47fa9f6a038ee9a9d5d4753dc193413eb3a1e2 Mon Sep 17 00:00:00 2001 From: test-user Date: Wed, 27 May 2026 11:52:48 -0700 Subject: [PATCH] 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) --- CLAUDE.md | 2 + pyproject.toml | 2 +- src/remove_ai_watermarks/__init__.py | 2 +- src/remove_ai_watermarks/cli.py | 15 ++-- src/remove_ai_watermarks/doubao_engine.py | 4 +- src/remove_ai_watermarks/gemini_engine.py | 4 +- src/remove_ai_watermarks/image_io.py | 66 +++++++++++++++++ src/remove_ai_watermarks/invisible_engine.py | 12 ++- .../invisible_watermark.py | 5 +- tests/test_cli.py | 18 +++++ tests/test_image_io.py | 74 +++++++++++++++++++ uv.lock | 2 +- 12 files changed, 189 insertions(+), 17 deletions(-) create mode 100644 src/remove_ai_watermarks/image_io.py create mode 100644 tests/test_image_io.py diff --git a/CLAUDE.md b/CLAUDE.md index ac6248a..8601b54 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -37,6 +37,7 @@ You are a **principal Python engineer** maintaining a CLI tool and library for r - `invisible_watermark.py` — `detect_invisible_watermark(path)` decodes the OPEN DWT-DCT watermarks (public decoder, no key) embedded by Stable Diffusion / SDXL / FLUX via the `imwatermark` library. Known fixed patterns (verified against upstream source) live in `_BITS_48` (SDXL 48-bit, FLUX.2 48-bit) and `_SD1_STRING` ("StableDiffusionV1", SD 1.x/2.x). Optional dep (extra `detect`); returns None when absent. The `detect` extra pulls **torch** transitively (invisible-watermark declares torch a hard dep, and `WatermarkDecoder` eagerly imports `rivaGan` -> `torch` at import time), so detection needs torch present even though dwtDct runs CPU-only on cv2/numpy/pywavelets — no GPU and no separate `gpu` extra required. **Unlike SynthID this is locally detectable**, but the watermark is fragile (does not survive JPEG re-encode/resize — verified gone after JPEG q90), so it confirms origin only on pristine files. Add new known patterns here. The file carries a top-of-module pyright pragma because imwatermark/cv2 ship no type stubs. - `trustmark_detector.py` — `detect_trustmark(path)` decodes the OPEN, keyless **Adobe TrustMark** watermark (the soft binding behind Adobe Durable Content Credentials, `alg` `com.adobe.trustmark.P`) via the optional `trustmark` package (extra `trustmark`; pulls torch, downloads model weights on first use). Mirrors `invisible_watermark.py` (lazy singleton, top-of-module pyright pragma, returns None when absent). It detects *provenance*, not AI origin as such (TrustMark also marks human-authored content), so `identify` lists it as a watermark without setting `is_ai_generated`. Other soft-binding vendors (Digimarc/Imatag/Steg.AI/...) have no public decoder — they are only *named* via the `C2PA_SOFT_BINDINGS` scan, not decoded. - `face_protector.py` — YOLO detect + soft-blend pattern; mirror this for any "protect region during diffusion" features +- `image_io.py` — Unicode-safe cv2 IO (issue #17). `imread(path, flags=None)` / `imwrite(path, img)` wrap `np.fromfile`+`cv2.imdecode` / `cv2.imencode`+`tofile` so non-ASCII paths work on Windows -- bare `cv2.imread`/`cv2.imwrite` use the platform ANSI code-page API there and fail (empty decode + `can't open/read file`) on Chinese/Cyrillic/accented filenames. `imread` keeps `cv2.imread` semantics (defaults to `IMREAD_COLOR`, returns `None` on missing/empty/undecodable). **Every cv2 file read/write in the package routes through here; do not call `cv2.imread`/`cv2.imwrite` directly.** macOS/Linux already accept UTF-8 paths, so it is behavior-neutral there (the bug only reproduces on Windows). cv2/numpy are imported lazily inside the functions, so the module is cheap to import in a bare env. ### Doubao clean-reverse-alpha distillation (researched 2026-05-26, NOT yet shipped) @@ -61,6 +62,7 @@ Who embeds what, and whether it is locally detectable (so we know which gaps are - `invisible` pipeline processes at **native resolution by default** (`max_resolution=0`), matching the hosted raiw.cc backend (fal fast-sdxl, no pre-downscale). The old forced downscale-to-1024 -> upscale-back round-trip was the main quality loss (issue #10) and is gone; at strength ~0.05 SDXL img2img does not need the ~1024 downscale. `--max-resolution N` re-introduces an opt-in long-side cap purely to bound GPU/MPS memory on very large inputs (it reintroduces the lossy round-trip). For huge images that OOM at native, tile-based diffusion is still the proper long-term fix. **Concrete MPS data point (verified 2026-05-25 on a 1254x1254 gpt-image SDXL run, fp32, 20 GB MPS ceiling):** native res OOMs at the *UNet* step (peak ~17 GiB), not only the VAE decode, and the auto-fallback in `img2img_runner` reloads on CPU and finishes (slow, ~13 min) -- the output is still weight-identical and defeats SynthID, so "looks hung/crashed" on Mac is usually this CPU fallback, not a pipeline error. Adding `enable_vae_tiling()` alone does NOT prevent it (the peak is the UNet, not the VAE). The fast Mac workarounds are fp16 on MPS (roughly halves memory) or `--max-resolution` to cap the long side; neither is wired as the default. The native-vs-downscale decision lives in the pure helper `invisible_engine._target_size(w, h, max_resolution)` (returns `None` for native, a clamped target tuple otherwise) so it is unit-tested (`tests/test_invisible_engine.py::TestTargetSize`, the #10/#15 regression guard) without loading the model -- keep that logic in the helper, don't re-inline it. - Pyright first run is slow (2-3 min) due to ML deps (torch/diffusers/transformers stubs); full-project `uv run pyright` can stall for many minutes — scope it to changed files. - `ultralytics` monkey-patches `PIL.Image.open` and tries to autoload `pi_heif`. When `pi_heif` is missing, opening files raises `ModuleNotFoundError`, not `UnidentifiedImageError`. Code that opens user-supplied or unknown-format files should `except Exception`, not just `OSError`/`UnidentifiedImageError`. +- **rich `console.print` parses `[word]` as a style tag and silently drops unknown ones.** A literal bracketed token in a print string disappears: `pip install 'remove-ai-watermarks[gpu]'` rendered as `...remove-ai-watermarks'` (the `[gpu]` extra eaten), which sent users a broken install command (surfaced via #19). Escape the literal bracket as `\[gpu]` (in a normal Python string that is `"\\[gpu]"`) in any rich string carrying user-facing brackets. Regression-guarded by `tests/test_cli.py::TestGpuHintMarkup`. - Metadata detection for AVIF/HEIF/JPEG-XL relies on a binary scan for `C2PA_UUID` + `IPTC_AI_MARKERS`, plus EXIF `Software` / XMP `CreatorTool` generator tags via `metadata.exif_generator` (validated with synthesized AVIF/JPEG fixtures + an XMP raw-scan fixture). C2PA removal in those containers is implemented via `noai/isobmff.py` (top-level ``uuid`` / ``jumb`` box stripper, no re-encoding), which now also drops a top-level XMP ``uuid`` box that carries an AI label (matched by AI-marker content, not by the XMP UUID, so byte-order-robust) and covers MP4/MOV/M4V/M4A by content sniff. **Non-ISOBMFF audio/video removal is via ffmpeg** (`_FFMPEG_STRIP_EXTS` -> `_strip_with_ffmpeg`): WebM/Matroska (EBML), MP3 (ID3), WAV/FLAC/OGG (RIFF/Vorbis) are stripped losslessly with `ffmpeg -map_metadata -1 -map_chapters -1 -c copy` (codec data untouched). Requires ffmpeg on PATH; raises `RuntimeError` if absent or if ffmpeg can't parse the file. Verified end-to-end (a real ffmpeg-made WAV/MP3 with a `title=Suno AI` tag -> tag gone, audio bytes preserved). **Still NOT built (deliberate):** EXIF/XMP stored as *items inside the ``meta`` box* (typical for AVIF/HEIF images) needs meta-box surgery (iinf/iloc edit + mdat splice) with corruption risk -- exiftool would do it but is a non-installed binary dep, so it stays a documented gap. **Audio watermark DETECTION (Resemble PerTh) was evaluated and NOT built (2026-05-26):** `resemble-perth`'s `PerthImplicitWatermarker.get_watermark()` returns a raw bit-array with **no presence/confidence flag** (clean audio decodes to arbitrary bits too), so reliably distinguishing watermarked-from-clean needs either Resemble's fixed payload or a confidence API -- neither is public, and there's no real Resemble sample to calibrate against. Same wall-class as the SynthID pixel detector: the decode exists, reliable presence-detection does not. (perth's top-level `PerthImplicitWatermarker` is also gated to None unless `librosa` is importable.) - **SynthID detection is metadata-only.** There is no reliable *local* detector of the SynthID *pixel* watermark — Google's decoder is proprietary, no public spec or API (only a waitlisted portal). Authoritative confirmation: Google DeepMind's own paper "SynthID-Image: Image watermarking at internet scale" (Gowal et al., arXiv:2510.09263) states the verification service is restricted to "trusted testers" and does not release detector weights or a reproducible algorithm — so a local pixel detector is infeasible by design, not just unbuilt. https://arxiv.org/abs/2510.09263 We detect SynthID by its C2PA companion (`synthid_source` / `SYNTHID_C2PA_ISSUERS`), which is reliable while the manifest is intact but says nothing once C2PA is stripped. **Surface-dependent blind spot (verified 2026-05-24):** the same Google model emits different metadata per surface -- the Gemini *app* wraps outputs in Google C2PA, but the *API/playground* (AI Studio, Nano Banana / gemini-2.5-flash-image) emits the SynthID *pixel* watermark (confirmed via the Gemini-app oracle) + the visible sparkle but **no C2PA/IPTC at all**, so `synthid_source` returns None despite SynthID being present. Only the pixel oracle or the visible-sparkle detector catches those. (Meta AI is another surface mismatch: it writes the IPTC `digitalSourceType=trainedAlgorithmicMedia` marker, not C2PA and not SynthID.) Google→SynthID is long-standing; OpenAI→SynthID is confirmed by OpenAI's Help Center (ChatGPT/Codex/API "include both C2PA metadata and SynthID watermarks", updated 2026-05-21) but time-gated (pre-rollout OpenAI images carry C2PA without SynthID), so the OpenAI verdict is hedged "likely". Oracles: Gemini app "Verify with SynthID" (Google), openai.com/verify (OpenAI). The spectral phase-coherence approach from `github.com/aloshdenny/reverse-SynthID` was evaluated (May 2026) and **does not work for real-content detection**: on its own shipped codebook + validation set, watermarked and cleaned images were indistinguishable (conf within noise, cleaned often higher); it only fires on pure-black 1024x1024 reference images at exact resolution (the controlled case it was calibrated on). The README's "90% / conf=0.91" reproduces only in that lab condition. Do not build a production detector on it; if revisited, it is experimental/diagnostic only and needs a per-resolution, per-model reference corpus. A from-scratch gpt-image pilot (2026-05-24) confirmed this independently: 5 independent solid-black gpt-image outputs share a near-identical fixed signature (pairwise residual correlation **0.92**, avg-template retains 97% energy), so the watermark/carrier IS strongly present and consistent on flat content — but the carrier frequencies extracted from it do NOT discriminate real content (carrier-to-random ratio: cleaned 1.86 > watermarked 1.53; a non-gpt-image image scored highest at 3.67). The signature drowns in content texture. Net: a perfectly consistent solid-color signature still yields no real-content pixel detector with magnitude/carrier methods. A corpus discrimination test (2026-05-24, `scripts/synthid_pixel_probe.py`, raw zero-mean residual NCC) independently re-confirms this: at matched resolution, SynthID positives do NOT cluster apart from negatives (within-Gemini 0.07; at 1024 px pos-vs-neg >= pos-vs-pos). The only high correlations were near-duplicate *content* (5 ChatGPT renders of one prompt at ~0.92, while a distinct ChatGPT image scored ~0 against them) — content, not a carrier. The probe is solid-fills-only and EXPERIMENTAL/DIAGNOSTIC; do not use it on real content. **Correction (deeper re-examination 2026-05-25):** the carrier IS real on solid fills — the earlier "no carrier" was a *method* artifact of using spatial / FFT-magnitude NCC, which can't see it. The carrier is a fixed *phase* at specific low frequencies, so the right metric is **per-bin phase coherence**. On 8 white `gemini-2.5-flash-image` fills (generated via the reverse-SynthID trick: identity-edit prompt "Recreate this image exactly as it is" on a synthetic pure-white PNG — this bypasses the recitation block that rejects text prompts for pure colors), phase coherence at the white carriers `(0,±7..±12,±20..±23)` = **0.86** vs **0.31** random; single-image leave-one-out phase-match **+0.83** vs real photos **-0.24**. (Black `2.5-flash` fills clip to std≈0 — SynthID can't push values below 0, so no carrier in black; the repo's dark carriers come from nano-banana-pro.) **But it does not generalize:** (a) carriers are model-version + resolution + color specific — the repo's v4 codebook (built for `gemini-3.1-flash-image-preview` + `nano-banana-pro-preview`) scores ~0.527 on my 2.5-flash white fills, indistinguishable from negatives (~0.50), i.e. carriers shift across model versions and need a per-model codebook; (b) on real content (30 `2.5-flash` images) the carrier collapses — set phase coherence at carriers 0.37 ≈ random 0.42, and the repo's v4 detector gives content 0.518 ≈ negatives 0.504 (no separation; a faint +0.24 single-image lean is likely a brightness confound). Net: the spectral/phase approach is a real *controlled-fill* characterizer, NOT an arbitrary-real-content detector, and is brittle to model version. Metadata proxy + visible sparkle + online oracles remain the ceiling for real content. - **External AI-vs-real classifier models are out of scope (decided 2026-05-24).** Generic HuggingFace detectors (`Organika/sdxl-detector` Swin Transformer, `umm-maybe/AI-image-detector`, and fine-tunes) exist and report ~0.98 on their *own* SDXL-vs-real validation sets, but they are per-generator and the model cards themselves note degraded accuracy off-distribution; they are untested on gpt-image / Gemini Nano Banana (the metadata-stripped surfaces we care about), and our own light SDXL pass would likely defeat them the same way it defeats SynthID. Detection here stays local + signal-based (metadata + visible sparkle); do not add a bundled classifier dependency. diff --git a/pyproject.toml b/pyproject.toml index d274737..8b3d762 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "remove-ai-watermarks" -version = "0.6.5" +version = "0.6.6" description = "Remove visible and invisible AI watermarks from images (Gemini / Nano Banana, ChatGPT, Stable Diffusion)" readme = "README.md" requires-python = ">=3.10" diff --git a/src/remove_ai_watermarks/__init__.py b/src/remove_ai_watermarks/__init__.py index 0a6209e..5ea2d25 100644 --- a/src/remove_ai_watermarks/__init__.py +++ b/src/remove_ai_watermarks/__init__.py @@ -1,3 +1,3 @@ """Remove-AI-Watermarks: Unified tool for removing visible and invisible AI watermarks.""" -__version__ = "0.6.5" +__version__ = "0.6.6" diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index e4506dd..435848e 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -85,7 +85,9 @@ def _read_bgr_and_alpha(path: Path) -> tuple[np.ndarray | None, np.ndarray | Non """ import cv2 - image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED) + from remove_ai_watermarks import image_io + + image = image_io.imread(path, cv2.IMREAD_UNCHANGED) if image is None: return None, None if image.ndim == 2: @@ -109,11 +111,12 @@ def _write_bgr_with_alpha( forced to 0 inside that bbox (expanded by ``pad`` px) so the watermark area becomes fully transparent in the saved file. """ - import cv2 import numpy as np + from remove_ai_watermarks import image_io + if alpha is None or path.suffix.lower() not in _ALPHA_FORMATS: - cv2.imwrite(str(path), bgr) + image_io.imwrite(path, bgr) return alpha_out = alpha @@ -127,7 +130,7 @@ def _write_bgr_with_alpha( alpha_out[y0:y1, x0:x1] = 0 bgra = np.dstack([bgr, alpha_out]) - cv2.imwrite(str(path), bgra) + image_io.imwrite(path, bgra) def _run_doubao_if_selected( @@ -481,7 +484,7 @@ def cmd_invisible( if not invisible_available(): console.print( "[red]Error:[/] GPU dependencies not installed.\n" - " Install them with: [bold]pip install 'remove-ai-watermarks[gpu]'[/]" + " Install them with: [bold]pip install 'remove-ai-watermarks\\[gpu]'[/]" ) raise SystemExit(1) @@ -744,7 +747,7 @@ def cmd_all( if not invisible_available(): console.print( " [yellow]⚠[/] Skipped — GPU dependencies not installed.\n" - " Install them with: [bold]pip install 'remove-ai-watermarks[gpu]'[/]" + " Install them with: [bold]pip install 'remove-ai-watermarks\\[gpu]'[/]" ) else: from remove_ai_watermarks.invisible_engine import InvisibleEngine diff --git a/src/remove_ai_watermarks/doubao_engine.py b/src/remove_ai_watermarks/doubao_engine.py index 19a27a7..82c6833 100644 --- a/src/remove_ai_watermarks/doubao_engine.py +++ b/src/remove_ai_watermarks/doubao_engine.py @@ -239,7 +239,9 @@ class DoubaoEngine: def load_image_bgr(path: str | Path) -> NDArray: """Read an image as BGR ndarray (helper for scripts/tests).""" - img = cv2.imread(str(path), cv2.IMREAD_COLOR) + from remove_ai_watermarks import image_io + + img = image_io.imread(path, cv2.IMREAD_COLOR) if img is None: raise FileNotFoundError(f"Failed to read image: {path}") return img diff --git a/src/remove_ai_watermarks/gemini_engine.py b/src/remove_ai_watermarks/gemini_engine.py index cdf2d9a..1d4898b 100644 --- a/src/remove_ai_watermarks/gemini_engine.py +++ b/src/remove_ai_watermarks/gemini_engine.py @@ -556,7 +556,9 @@ def detect_sparkle_confidence(image_path: Path) -> float | None: (cv2 returns None for unsupported containers such as HEIC). Kept here so the cv2 dependency stays in this module; callers apply their own threshold. """ - img = cv2.imread(str(image_path)) + from remove_ai_watermarks import image_io + + img = image_io.imread(image_path) if img is None: return None return float(GeminiEngine().detect_watermark(img).confidence) diff --git a/src/remove_ai_watermarks/image_io.py b/src/remove_ai_watermarks/image_io.py new file mode 100644 index 0000000..78cfd47 --- /dev/null +++ b/src/remove_ai_watermarks/image_io.py @@ -0,0 +1,66 @@ +"""Unicode-safe cv2 image IO (issue #17). + +``cv2.imread`` / ``cv2.imwrite`` pass the path to the platform C runtime, which +on Windows uses the narrow (ANSI) code-page API and therefore fails on paths +containing non-ASCII characters (Chinese, Cyrillic, ...). The symptom is a +``can't open/read file`` warning and a ``None`` decode even though the file +exists. + +These wrappers route through numpy buffers instead: ``np.fromfile`` / +``ndarray.tofile`` open the path in Python (full Unicode), and +``cv2.imdecode`` / ``cv2.imencode`` do the codec work. The decoded/encoded +bytes are byte-for-byte identical to ``imread`` / ``imwrite``. On macOS/Linux +cv2 already accepts UTF-8 paths, so the wrappers are behavior-neutral there. + +cv2/numpy are imported lazily inside the functions so importing this module +stays cheap in a bare environment (matching the rest of the package). +""" + +# cv2 ships no type stubs; mirror the pragma used by the other cv2-using modules. +# pyright: reportMissingTypeStubs=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from numpy.typing import NDArray + + +def imread(path: str | Path, flags: int | None = None) -> NDArray[Any] | None: + """Unicode-safe ``cv2.imread``. + + ``flags`` defaults to ``cv2.IMREAD_COLOR`` (same as ``cv2.imread``). Returns + ``None`` when the file is missing or cannot be decoded, matching + ``cv2.imread`` semantics so existing ``if img is None`` checks keep working. + """ + import cv2 + import numpy as np + + if flags is None: + flags = cv2.IMREAD_COLOR + try: + data = np.fromfile(str(path), dtype=np.uint8) + except OSError: + return None + if data.size == 0: + return None + return cv2.imdecode(data, flags) + + +def imwrite(path: str | Path, img: NDArray[Any]) -> bool: + """Unicode-safe ``cv2.imwrite``. + + The output format is taken from the path extension (e.g. ``.png``), exactly + like ``cv2.imwrite``. Returns ``True`` on success, ``False`` if the codec + rejects the image. + """ + import cv2 + + ext = Path(path).suffix or ".png" + ok, buf = cv2.imencode(ext, img) + if not ok: + return False + buf.tofile(str(path)) + return True diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index 8b03689..98225d1 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/src/remove_ai_watermarks/invisible_engine.py @@ -217,7 +217,9 @@ class InvisibleEngine: import cv2 import numpy as np - out_cv = cv2.imread(str(out_path), cv2.IMREAD_COLOR) + from remove_ai_watermarks import image_io + + out_cv = image_io.imread(out_path, cv2.IMREAD_COLOR) if protect_faces and original_faces: if self._progress_callback: @@ -243,20 +245,22 @@ class InvisibleEngine: # Using INTER_LANCZOS4 for high-quality upscaling back to original out_cv = cv2.resize(out_cv, orig_size, interpolation=cv2.INTER_LANCZOS4) - cv2.imwrite(str(out_path), out_cv) + image_io.imwrite(out_path, out_cv) else: # Even if no protect_faces or humanize, we must restore original size if needed import cv2 - out_cv = cv2.imread(str(out_path), cv2.IMREAD_COLOR) + from remove_ai_watermarks import image_io + + out_cv = image_io.imread(out_path, cv2.IMREAD_COLOR) if out_cv is not None and (out_cv.shape[1], out_cv.shape[0]) != orig_size: if self._progress_callback: self._progress_callback( f"Upscaling result back to original resolution {orig_size[0]}x{orig_size[1]}..." ) out_cv = cv2.resize(out_cv, orig_size, interpolation=cv2.INTER_LANCZOS4) - cv2.imwrite(str(out_path), out_cv) + image_io.imwrite(out_path, out_cv) return out_path finally: diff --git a/src/remove_ai_watermarks/invisible_watermark.py b/src/remove_ai_watermarks/invisible_watermark.py index 5b4167e..10d3c62 100644 --- a/src/remove_ai_watermarks/invisible_watermark.py +++ b/src/remove_ai_watermarks/invisible_watermark.py @@ -78,10 +78,11 @@ def detect_invisible_watermark(image_path: Path) -> str | None: """ if not is_available(): return None - import cv2 from imwatermark import WatermarkDecoder - img = cv2.imread(str(image_path)) + from remove_ai_watermarks import image_io + + img = image_io.imread(image_path) if img is None: return None diff --git a/tests/test_cli.py b/tests/test_cli.py index 0f20175..ecea89e 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -509,3 +509,21 @@ class TestBatchCommand: assert result.exit_code == 0 expected_dir = tmp_path / "input_clean" assert expected_dir.exists() + + +class TestGpuHintMarkup: + """The GPU-extra install hint must survive rich markup (the ``[gpu]`` token + is otherwise parsed as a style tag and silently dropped).""" + + def test_invisible_install_hint_keeps_gpu_extra(self, runner, sample_png): + with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False): + result = runner.invoke(main, ["invisible", str(sample_png)]) + assert result.exit_code != 0 + assert "remove-ai-watermarks[gpu]" in result.output + + def test_all_install_hint_keeps_gpu_extra(self, runner, sample_png): + # The `all` pipeline skips the invisible step with a warning that carries + # the same hint; it must keep the [gpu] extra too. + with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False): + result = runner.invoke(main, ["all", str(sample_png)]) + assert "remove-ai-watermarks[gpu]" in result.output diff --git a/tests/test_image_io.py b/tests/test_image_io.py new file mode 100644 index 0000000..9ff65ed --- /dev/null +++ b/tests/test_image_io.py @@ -0,0 +1,74 @@ +"""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 diff --git a/uv.lock b/uv.lock index e019357..124a0a1 100644 --- a/uv.lock +++ b/uv.lock @@ -2865,7 +2865,7 @@ wheels = [ [[package]] name = "remove-ai-watermarks" -version = "0.6.5" +version = "0.6.6" source = { editable = "." } dependencies = [ { name = "click" },