mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-06-05 10:38:00 +02:00
Victor Kuznetsov
9ca2811938
On a dark/textured background (e.g. grass) the captured alpha map over-estimates the real Gemini sparkle's effective opacity (~0.51 captured vs ~0.31 effective), so the fixed-alpha reverse blend over-subtracts (watermarked - alpha*logo goes negative) and drives the footprint to black -- the white sparkle turns into a black diamond (issue #30, reported by @CoolZimo1). remove_watermark now detects this via _reverse_alpha_oversubtracts (fraction of footprint pixels with a negative numerator > 5%) and inpaints the small sparkle footprint from the surrounding pixels (cv2 NS, cropped to a padded box) instead. Behavior-neutral on the working case: a bright background over-subtracts at ~0%, so reverse-alpha is used and the output is byte-identical to before (verified: demo_banana 0.0 frac vs the issue-#30 grass image 0.61 frac; issue-#30 footprint recovers to background grass with no pit, residual sparkle conf 0.25 < 0.35). Guard is scoped to GeminiEngine: doubao/jimeng already NCC-align their alpha to the actual mark per image, which sidesteps the fixed-alpha mismatch. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
287 lines
12 KiB
Python
287 lines
12 KiB
Python
"""Tests for the Gemini visible-watermark engine."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from remove_ai_watermarks.gemini_engine import (
|
|
DetectionResult,
|
|
GeminiEngine,
|
|
WatermarkPosition,
|
|
WatermarkSize,
|
|
_calculate_alpha_map,
|
|
detect_sparkle_confidence,
|
|
get_watermark_config,
|
|
get_watermark_size,
|
|
)
|
|
|
|
# ── WatermarkSize / config helpers ──────────────────────────────────
|
|
|
|
|
|
class TestWatermarkConfig:
|
|
"""Tests for watermark size detection and position calculation."""
|
|
|
|
def test_small_image_gets_small_watermark(self):
|
|
assert get_watermark_size(800, 600) == WatermarkSize.SMALL
|
|
|
|
def test_large_image_gets_large_watermark(self):
|
|
assert get_watermark_size(1920, 1080) == WatermarkSize.LARGE
|
|
|
|
def test_boundary_image_stays_small(self):
|
|
"""Exactly 1024x1024 should be SMALL (rule: > 1024 for LARGE)."""
|
|
assert get_watermark_size(1024, 1024) == WatermarkSize.SMALL
|
|
|
|
def test_one_dimension_small(self):
|
|
"""Only one dimension > 1024 → still SMALL."""
|
|
assert get_watermark_size(2000, 500) == WatermarkSize.SMALL
|
|
|
|
def test_config_small_returns_correct_values(self):
|
|
config = get_watermark_config(800, 600)
|
|
assert config.margin_right == 32
|
|
assert config.margin_bottom == 32
|
|
assert config.logo_size == 48
|
|
|
|
def test_config_large_returns_correct_values(self):
|
|
config = get_watermark_config(1920, 1080)
|
|
assert config.margin_right == 64
|
|
assert config.margin_bottom == 64
|
|
assert config.logo_size == 96
|
|
|
|
def test_position_calculation(self):
|
|
pos = WatermarkPosition(margin_right=32, margin_bottom=32, logo_size=48)
|
|
x, y = pos.get_position(800, 600)
|
|
assert x == 800 - 32 - 48 # 720
|
|
assert y == 600 - 32 - 48 # 520
|
|
|
|
|
|
# ── Alpha map ───────────────────────────────────────────────────────
|
|
|
|
|
|
class TestAlphaMap:
|
|
"""Tests for alpha map calculation."""
|
|
|
|
def test_pure_black_gives_zero_alpha(self):
|
|
black = np.zeros((10, 10, 3), dtype=np.uint8)
|
|
alpha = _calculate_alpha_map(black)
|
|
assert alpha.shape == (10, 10)
|
|
np.testing.assert_array_equal(alpha, 0.0)
|
|
|
|
def test_pure_white_gives_one_alpha(self):
|
|
white = np.full((10, 10, 3), 255, dtype=np.uint8)
|
|
alpha = _calculate_alpha_map(white)
|
|
np.testing.assert_allclose(alpha, 1.0)
|
|
|
|
def test_grayscale_input(self):
|
|
gray = np.full((10, 10), 128, dtype=np.uint8)
|
|
alpha = _calculate_alpha_map(gray)
|
|
np.testing.assert_allclose(alpha, 128 / 255.0)
|
|
|
|
def test_max_channel_used(self):
|
|
"""Alpha should use max(R, G, B)."""
|
|
img = np.zeros((1, 1, 3), dtype=np.uint8)
|
|
img[0, 0] = [50, 200, 100] # BGR
|
|
alpha = _calculate_alpha_map(img)
|
|
assert pytest.approx(alpha[0, 0], rel=1e-3) == 200 / 255.0
|
|
|
|
|
|
# ── GeminiEngine ────────────────────────────────────────────────────
|
|
|
|
|
|
class TestGeminiEngine:
|
|
"""Tests for the GeminiEngine class."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup_engine(self):
|
|
self.engine = GeminiEngine()
|
|
|
|
def test_engine_loads_alpha_maps(self):
|
|
small = self.engine.get_alpha_map(WatermarkSize.SMALL)
|
|
large = self.engine.get_alpha_map(WatermarkSize.LARGE)
|
|
assert small.shape == (48, 48)
|
|
assert large.shape == (96, 96)
|
|
|
|
def test_remove_watermark_returns_same_shape(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.remove_watermark(image)
|
|
assert result.shape == image.shape
|
|
assert result.dtype == np.uint8
|
|
|
|
def test_remove_watermark_does_not_modify_input(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
original = image.copy()
|
|
self.engine.remove_watermark(image)
|
|
np.testing.assert_array_equal(image, original)
|
|
|
|
def test_remove_watermark_large_image(self, tmp_large_image_path):
|
|
image = cv2.imread(str(tmp_large_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.remove_watermark(image)
|
|
assert result.shape == image.shape
|
|
|
|
def test_remove_watermark_custom_region(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.remove_watermark_custom(image, (10, 10, 48, 48))
|
|
assert result.shape == image.shape
|
|
|
|
def test_remove_watermark_custom_large_region(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.remove_watermark_custom(image, (10, 10, 96, 96))
|
|
assert result.shape == image.shape
|
|
|
|
def test_remove_watermark_custom_arbitrary_region(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.remove_watermark_custom(image, (5, 5, 60, 60))
|
|
assert result.shape == image.shape
|
|
|
|
def test_force_size(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.remove_watermark(image, force_size=WatermarkSize.LARGE)
|
|
assert result.shape == image.shape
|
|
|
|
|
|
# ── Detection ───────────────────────────────────────────────────────
|
|
|
|
|
|
class TestDetection:
|
|
"""Tests for watermark detection."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup_engine(self):
|
|
self.engine = GeminiEngine()
|
|
|
|
def test_detect_returns_result_object(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.detect_watermark(image)
|
|
assert isinstance(result, DetectionResult)
|
|
assert 0.0 <= result.confidence <= 1.0
|
|
|
|
def test_detect_empty_image_returns_no_detection(self):
|
|
empty = np.zeros((0, 0, 3), dtype=np.uint8)
|
|
result = self.engine.detect_watermark(empty)
|
|
assert not result.detected
|
|
assert result.confidence == 0.0
|
|
|
|
def test_detect_none_image_returns_no_detection(self):
|
|
result = self.engine.detect_watermark(None)
|
|
assert not result.detected
|
|
|
|
def test_detect_random_image_low_confidence(self, tmp_image_path):
|
|
"""Random noise should not look like a watermark."""
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.detect_watermark(image)
|
|
# Random image may or may not be detected; confidence should be meaningful
|
|
assert isinstance(result.spatial_score, float)
|
|
assert isinstance(result.gradient_score, float)
|
|
|
|
|
|
class TestDetectSparkleConfidence:
|
|
"""File-level entry point used by identify.py."""
|
|
|
|
def test_returns_float_in_range_for_real_image(self, tmp_image_path):
|
|
conf = detect_sparkle_confidence(tmp_image_path)
|
|
assert conf is not None
|
|
assert 0.0 <= conf <= 1.0
|
|
|
|
def test_returns_none_for_unreadable_file(self, tmp_path):
|
|
# cv2.imread returns None for a non-image; the helper must not raise.
|
|
bogus = tmp_path / "not_an_image.png"
|
|
bogus.write_bytes(b"this is not a PNG")
|
|
assert detect_sparkle_confidence(bogus) is None
|
|
|
|
|
|
# ── Inpainting ──────────────────────────────────────────────────────
|
|
|
|
|
|
class TestInpainting:
|
|
"""Tests for residual inpainting."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup_engine(self):
|
|
self.engine = GeminiEngine()
|
|
|
|
def test_inpaint_ns(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.inpaint_residual(image, (150, 150, 48, 48), method="ns")
|
|
assert result.shape == image.shape
|
|
|
|
def test_inpaint_telea(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.inpaint_residual(image, (150, 150, 48, 48), method="telea")
|
|
assert result.shape == image.shape
|
|
|
|
def test_inpaint_gaussian(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.inpaint_residual(image, (150, 150, 48, 48), method="gaussian")
|
|
assert result.shape == image.shape
|
|
|
|
def test_inpaint_zero_strength(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.inpaint_residual(image, (150, 150, 48, 48), strength=0.0)
|
|
np.testing.assert_array_equal(result, image)
|
|
|
|
def test_inpaint_tiny_region_returns_unchanged(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
result = self.engine.inpaint_residual(image, (10, 10, 2, 2))
|
|
np.testing.assert_array_equal(result, image)
|
|
|
|
def test_inpaint_does_not_modify_input(self, tmp_image_path):
|
|
image = cv2.imread(str(tmp_image_path), cv2.IMREAD_COLOR)
|
|
original = image.copy()
|
|
self.engine.inpaint_residual(image, (150, 150, 48, 48))
|
|
np.testing.assert_array_equal(image, original)
|
|
|
|
|
|
class TestOverSubtractionGuard:
|
|
"""Issue #30: reverse-alpha must not turn the sparkle into a black pit.
|
|
|
|
On a dark background the captured alpha over-estimates the real sparkle opacity,
|
|
so the fixed-alpha reverse blend over-subtracts and drives the footprint to black.
|
|
The engine detects this and inpaints the footprint instead.
|
|
"""
|
|
|
|
# Composite the mark at ~60% of the captured opacity: the engine's alpha maxes at
|
|
# ~0.51, real dark-background sparkles sit nearer ~0.31, so 0.6x reproduces the
|
|
# capture-over-estimates-reality mismatch that triggers the bug.
|
|
_REALISTIC_ALPHA_SCALE = 0.6
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup_engine(self):
|
|
self.engine = GeminiEngine()
|
|
|
|
def _composite_sparkle(self, bg_value: int, size: int = 1400, alpha_scale: float = _REALISTIC_ALPHA_SCALE):
|
|
"""Build a flat BGR image of ``bg_value`` with the sparkle composited in.
|
|
|
|
The mark is composited at a LOWER effective opacity than the engine's captured
|
|
alpha map (``alpha_scale`` < 1), reproducing the real-world mismatch behind
|
|
issue #30: the captured alpha (~0.51) over-estimates a real sparkle whose
|
|
effective opacity is lower, so the fixed-alpha reverse blend over-subtracts.
|
|
Placed at the configured large-image position so the detector locates it.
|
|
"""
|
|
img = np.full((size, size, 3), bg_value, dtype=np.float32)
|
|
config = get_watermark_config(size, size)
|
|
x, y = config.get_position(size, size)
|
|
alpha = self.engine.get_alpha_map(WatermarkSize.LARGE)
|
|
ah, aw = alpha.shape[:2]
|
|
a = (alpha * alpha_scale)[:, :, None]
|
|
roi = img[y : y + ah, x : x + aw]
|
|
img[y : y + ah, x : x + aw] = a * 255.0 + (1.0 - a) * roi
|
|
return np.clip(img, 0, 255).astype(np.uint8), (x, y, aw, ah)
|
|
|
|
def test_dark_background_does_not_leave_black_pit(self):
|
|
image, (x, y, w, h) = self._composite_sparkle(bg_value=60)
|
|
out = self.engine.remove_watermark(image)
|
|
footprint = out[y : y + h, x : x + w]
|
|
# The recovered footprint must read like the dark background, not a black hole.
|
|
assert footprint.min() > 25, f"black pit: min={footprint.min()}"
|
|
assert abs(float(footprint.mean()) - 60.0) < 25.0
|
|
|
|
def test_bright_background_keeps_reverse_alpha(self):
|
|
"""A bright background does not over-subtract, so reverse-alpha is used."""
|
|
bright, pos = self._composite_sparkle(bg_value=230)
|
|
alpha = self.engine.get_interpolated_alpha(pos[2])
|
|
assert self.engine._reverse_alpha_oversubtracts(bright, alpha, (pos[0], pos[1])) is False
|
|
dark, dpos = self._composite_sparkle(bg_value=60)
|
|
dalpha = self.engine.get_interpolated_alpha(dpos[2])
|
|
assert self.engine._reverse_alpha_oversubtracts(dark, dalpha, (dpos[0], dpos[1])) is True
|