mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-07-07 00:47:50 +02:00
Victor Kuznetsov
0e5a4cbc54
- Add the Jimeng-basic top-left "AI生成" pill as a CAPTURE-LESS mark (pill_engine.py): synthetic-silhouette edge-NCC detect + inpaint-only removal. Gated in remove_auto_marks: kept only when Jimeng is confirmed (TC260 metadata OR the bottom-right "★ 即梦AI" wordmark fired -- the wordmark keeps recall on metadata-STRIPPED uploads) AND Doubao did not fire. - Add an inpaint-fallback removal path + MI-GAN ONNX backend (migan extra, MIT, ~28 MB / ~1 GB peak -- droplet-friendly) alongside big-LaMa. New --method auto|reverse-alpha|inpaint (shared across visible/all/batch) and erase --backend migan; footprint_mask on each engine. - auto is deterministic: reverse-alpha for capture marks (recovers exact pixels, lighter -- measured cleaner than MI-GAN on structured backgrounds) and inpaint only for the capture-less pill. - --mark auto now removes EVERY detected mark in one pass (remove_auto_marks), so a Jimeng-basic image's top-left pill AND bottom-right wordmark both clear. - Bump 0.12.1 -> 0.13.0. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
"""Tests for the known-visible-watermark registry (reverse-alpha only)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from remove_ai_watermarks import watermark_registry as reg
|
|
|
|
DOUBAO_SAMPLE = Path(__file__).resolve().parents[1] / "data" / "samples" / "doubao-1.png"
|
|
|
|
|
|
class TestCatalog:
|
|
def test_keys(self):
|
|
assert reg.mark_keys() == ["gemini", "doubao", "jimeng", "samsung", "jimeng_pill"]
|
|
|
|
def test_all_in_auto(self):
|
|
assert all(m.in_auto for m in reg.known_marks())
|
|
|
|
def test_recovery(self):
|
|
# Capture marks recover by reverse-alpha; the capture-less pill is inpaint-only.
|
|
by_key = {m.key: m for m in reg.known_marks()}
|
|
assert all(by_key[k].recovery == "reverse-alpha" for k in ("gemini", "doubao", "jimeng", "samsung"))
|
|
assert by_key["jimeng_pill"].recovery == "inpaint"
|
|
assert by_key["jimeng_pill"].has_capture is False
|
|
|
|
def test_locations(self):
|
|
by_key = {m.key: m for m in reg.known_marks()}
|
|
assert by_key["gemini"].location == "bottom-right"
|
|
assert by_key["doubao"].location == "bottom-right"
|
|
assert by_key["jimeng"].location == "bottom-right"
|
|
assert by_key["samsung"].location == "bottom-left"
|
|
assert by_key["jimeng_pill"].location == "top-left"
|
|
|
|
def test_get_mark_unknown_raises(self):
|
|
with pytest.raises(KeyError):
|
|
reg.get_mark("nope")
|
|
|
|
|
|
class TestScan:
|
|
def test_detect_marks_scans_all(self):
|
|
img = np.zeros((256, 256, 3), np.uint8)
|
|
keys = {d.key for d in reg.detect_marks(img)}
|
|
assert keys == {"gemini", "doubao", "jimeng", "samsung", "jimeng_pill"}
|
|
|
|
def test_blank_image_no_auto_mark(self):
|
|
assert reg.best_auto_mark(np.zeros((256, 256, 3), np.uint8)) is None
|
|
|
|
|
|
@pytest.mark.skipif(not DOUBAO_SAMPLE.exists(), reason="doubao sample not present")
|
|
class TestRealSample:
|
|
def test_doubao_sample_wins_auto(self):
|
|
from remove_ai_watermarks.image_io import imread
|
|
|
|
best = reg.best_auto_mark(imread(DOUBAO_SAMPLE))
|
|
assert best is not None
|
|
assert best.key == "doubao"
|
|
|
|
def test_doubao_remove_returns_region(self):
|
|
from remove_ai_watermarks.image_io import imread
|
|
|
|
img = imread(DOUBAO_SAMPLE) # 2048 wide -> reverse-alpha applies
|
|
result, region = reg.get_mark("doubao").remove(img)
|
|
assert region is not None
|
|
assert result.shape == img.shape
|
|
|
|
|
|
class TestReverseAlphaOnly:
|
|
def test_doubao_off_resolution_is_skipped(self):
|
|
# No alpha capture for this width -> no inpaint fallback, image untouched.
|
|
img = np.zeros((512, 512, 3), np.uint8)
|
|
result, region = reg.get_mark("doubao").remove(img)
|
|
assert region is None
|
|
assert np.array_equal(result, img)
|