mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-07-25 17:00:57 +02:00
feat(identify): detect visible Doubao/Jimeng marks; keep identify import torch-free
identify previously ran only the Gemini sparkle as a visible detector, so a Doubao/Jimeng image with stripped TC260 metadata had no visible fallback. Add `_visible_text_marks` (registry-backed) so the ByteDance Doubao 豆包AI生成 and Jimeng 即梦AI marks are detected too, each gated by its own engine NCC threshold via MarkDetection.detected. New signals `visible_doubao` / `visible_jimeng` (medium), same stripped-metadata fallback role as the sparkle; excluded from integrity-clash vendor claims; set platform only when no harder signal did. Also make `noai/__init__` lazy (PEP 562 __getattr__): importing the light `noai.c2pa` / `noai.constants` submodules (which identify needs) no longer eagerly pulls `watermark_remover`, which imports torch + diffusers at module top. `import remove_ai_watermarks.identify` drops from ~420 MB to ~21 MB in a full gpu/detect install (torch not loaded), so it fits a 512 MB host; the removal API resolves lazily on first access. Guarded by TestIdentifyImportIsLight. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -7,6 +7,8 @@ against the real committed C2PA / IPTC fixtures in data/samples/.
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
@@ -371,6 +373,68 @@ class TestIdentifyVisibleSparkle:
|
||||
assert r.confidence == "high"
|
||||
|
||||
|
||||
class TestIdentifyImportIsLight:
|
||||
"""`import identify` must stay torch-free (lazy noai/__init__): the package
|
||||
is deployed on a 512 MB host where eagerly pulling torch/diffusers OOMs."""
|
||||
|
||||
def test_import_identify_does_not_pull_torch(self):
|
||||
# Only meaningful where torch is installed (the gpu/detect extra); on a
|
||||
# core-only CI runner torch can't be in sys.modules anyway.
|
||||
pytest.importorskip("torch")
|
||||
code = "import sys, remove_ai_watermarks.identify; sys.exit(1 if 'torch' in sys.modules else 0)"
|
||||
result = subprocess.run([sys.executable, "-c", code], capture_output=True, check=False) # noqa: S603
|
||||
assert result.returncode == 0, f"import identify pulled torch: {result.stderr.decode()[-500:]}"
|
||||
|
||||
|
||||
# Where the registry-backed Doubao/Jimeng visible detector resolves.
|
||||
_TEXT_MARKS_TARGET = "remove_ai_watermarks.identify._visible_text_marks"
|
||||
|
||||
|
||||
class TestIdentifyVisibleTextMarks:
|
||||
"""The visible Doubao/Jimeng marks are a stripped-metadata visual fallback,
|
||||
parallel to the Gemini sparkle: each lifts an Unknown verdict to medium."""
|
||||
|
||||
@staticmethod
|
||||
def _detection(key: str, label: str, conf: float):
|
||||
from remove_ai_watermarks.watermark_registry import MarkDetection
|
||||
|
||||
return MarkDetection(key, label, "bottom-right", True, conf, (0, 0, 10, 10))
|
||||
|
||||
def test_doubao_promotes_to_medium(self, tmp_clean_png: Path):
|
||||
det = self._detection("doubao", "Doubao 豆包AI生成 text", 0.8)
|
||||
with patch(_SPARKLE_TARGET, return_value=None), patch(_TEXT_MARKS_TARGET, return_value=[det]):
|
||||
r = identify(tmp_clean_png, check_visible=True)
|
||||
assert r.is_ai_generated is True
|
||||
assert r.confidence == "medium"
|
||||
assert r.platform is not None
|
||||
assert "Doubao" in r.platform
|
||||
signal = next(s for s in r.signals if s.name == "visible_doubao")
|
||||
assert signal.confidence == "medium"
|
||||
|
||||
def test_jimeng_promotes_to_medium(self, tmp_clean_png: Path):
|
||||
det = self._detection("jimeng", "Jimeng 即梦AI wordmark", 0.9)
|
||||
with patch(_SPARKLE_TARGET, return_value=None), patch(_TEXT_MARKS_TARGET, return_value=[det]):
|
||||
r = identify(tmp_clean_png, check_visible=True)
|
||||
assert r.is_ai_generated is True
|
||||
assert r.confidence == "medium"
|
||||
assert r.platform is not None
|
||||
assert "Jimeng" in r.platform
|
||||
assert any(s.name == "visible_jimeng" for s in r.signals)
|
||||
|
||||
def test_check_visible_false_skips_text_marks(self, tmp_clean_png: Path):
|
||||
det = self._detection("doubao", "Doubao 豆包AI生成 text", 0.99)
|
||||
with patch(_SPARKLE_TARGET, return_value=None), patch(_TEXT_MARKS_TARGET, return_value=[det]) as mock:
|
||||
r = identify(tmp_clean_png, check_visible=False)
|
||||
mock.assert_not_called()
|
||||
assert not any(s.name == "visible_doubao" for s in r.signals)
|
||||
|
||||
def test_metadata_keeps_high_even_with_text_mark(self, tmp_png_with_ai_metadata: Path):
|
||||
det = self._detection("doubao", "Doubao 豆包AI生成 text", 0.8)
|
||||
with patch(_SPARKLE_TARGET, return_value=None), patch(_TEXT_MARKS_TARGET, return_value=[det]):
|
||||
r = identify(tmp_png_with_ai_metadata, check_visible=True)
|
||||
assert r.confidence == "high"
|
||||
|
||||
|
||||
# ── Caveats and serialization ───────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user