mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-20 12:37:12 +02:00
Add project files, tests, and documentation for GitHub release
- CLI with visible, invisible, all, metadata, and batch commands - Gemini watermark removal via reverse alpha blending - Invisible watermark removal via diffusion regeneration (SynthID, TreeRing) - AI metadata stripping (EXIF, PNG text, C2PA) - Face protection (YOLO/Haar) and analog humanizer - 137 tests covering all CLI modes and core engines - Ruff and Pyright clean
This commit is contained in:
test-user
@@ -0,0 +1 @@
|
||||
"""Tests for remove-ai-watermarks."""
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Shared fixtures for remove-ai-watermarks test suite."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def tmp_image_path(tmp_path: Path) -> Path:
|
||||
"""Create a minimal 200×200 test PNG image and return its path."""
|
||||
img = np.random.randint(0, 255, (200, 200, 3), dtype=np.uint8)
|
||||
path = tmp_path / "test_image.png"
|
||||
cv2.imwrite(str(path), img)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def tmp_large_image_path(tmp_path: Path) -> Path:
|
||||
"""Create a 1200×1200 test PNG image (triggers large watermark branch)."""
|
||||
img = np.random.randint(0, 255, (1200, 1200, 3), dtype=np.uint8)
|
||||
path = tmp_path / "test_large.png"
|
||||
cv2.imwrite(str(path), img)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def tmp_jpeg_path(tmp_path: Path) -> Path:
|
||||
"""Create a minimal JPEG test image."""
|
||||
img = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
|
||||
path = tmp_path / "test_image.jpg"
|
||||
cv2.imwrite(str(path), img)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def tmp_png_with_ai_metadata(tmp_path: Path) -> Path:
|
||||
"""Create a PNG with AI-related metadata keys."""
|
||||
img = Image.new("RGB", (64, 64), color=(128, 128, 128))
|
||||
pnginfo = PngInfo()
|
||||
pnginfo.add_text("parameters", "Steps: 20, Sampler: Euler, CFG scale: 7")
|
||||
pnginfo.add_text("prompt", "a beautiful landscape")
|
||||
pnginfo.add_text("Author", "Test Author")
|
||||
path = tmp_path / "ai_metadata.png"
|
||||
img.save(path, pnginfo=pnginfo)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def tmp_clean_png(tmp_path: Path) -> Path:
|
||||
"""Create a PNG with no AI metadata."""
|
||||
img = Image.new("RGB", (64, 64), color=(200, 100, 50))
|
||||
pnginfo = PngInfo()
|
||||
pnginfo.add_text("Author", "Human Artist")
|
||||
pnginfo.add_text("Title", "Test Artwork")
|
||||
path = tmp_path / "clean.png"
|
||||
img.save(path, pnginfo=pnginfo)
|
||||
return path
|
||||
@@ -0,0 +1,322 @@
|
||||
"""Tests for the CLI entry point."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
from remove_ai_watermarks.cli import main
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def runner():
|
||||
return CliRunner()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def sample_png(tmp_path: Path) -> Path:
|
||||
"""Create a sample PNG for CLI testing."""
|
||||
img = np.random.randint(0, 255, (200, 200, 3), dtype=np.uint8)
|
||||
path = tmp_path / "input.png"
|
||||
cv2.imwrite(str(path), img)
|
||||
return path
|
||||
|
||||
|
||||
def _make_batch_dir(tmp_path: Path, count: int = 3) -> Path:
|
||||
"""Create a directory with test images for batch testing."""
|
||||
input_dir = tmp_path / "input"
|
||||
input_dir.mkdir()
|
||||
for i in range(count):
|
||||
img = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
|
||||
cv2.imwrite(str(input_dir / f"img_{i}.png"), img)
|
||||
return input_dir
|
||||
|
||||
|
||||
def _make_batch_dir_with_metadata(tmp_path: Path, count: int = 3) -> Path:
|
||||
"""Create a directory with PNG images containing AI metadata."""
|
||||
input_dir = tmp_path / "input"
|
||||
input_dir.mkdir()
|
||||
for i in range(count):
|
||||
img = Image.new("RGB", (64, 64), color=(100 + i, 150, 200))
|
||||
pnginfo = PngInfo()
|
||||
pnginfo.add_text("parameters", f"Steps: 20, Sampler: Euler, img_{i}")
|
||||
pnginfo.add_text("prompt", "a test landscape")
|
||||
img.save(input_dir / f"img_{i}.png", pnginfo=pnginfo)
|
||||
return input_dir
|
||||
|
||||
|
||||
def _mock_invisible_engine():
|
||||
"""Create a mock InvisibleEngine that writes a copy of the input image."""
|
||||
|
||||
def _mock_remove_watermark(image_path, output_path=None, **kwargs):
|
||||
out = output_path or image_path.with_stem(image_path.stem + "_clean")
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
img = Image.open(image_path)
|
||||
img.save(out)
|
||||
return out
|
||||
|
||||
mock_engine = MagicMock()
|
||||
mock_engine.remove_watermark.side_effect = _mock_remove_watermark
|
||||
mock_cls = MagicMock(return_value=mock_engine)
|
||||
return mock_cls, mock_engine
|
||||
|
||||
|
||||
class TestMainGroup:
|
||||
"""Tests for the top-level CLI group."""
|
||||
|
||||
def test_help(self, runner):
|
||||
result = runner.invoke(main, ["--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "Remove visible and invisible" in result.output
|
||||
|
||||
def test_version(self, runner):
|
||||
result = runner.invoke(main, ["--version"])
|
||||
assert result.exit_code == 0
|
||||
assert "0.1.0" in result.output
|
||||
|
||||
def test_no_command_shows_banner(self, runner):
|
||||
result = runner.invoke(main, [])
|
||||
assert result.exit_code == 0
|
||||
assert "Remove-AI-Watermarks" in result.output
|
||||
|
||||
|
||||
class TestVisibleCommand:
|
||||
"""Tests for the 'visible' subcommand."""
|
||||
|
||||
def test_visible_help(self, runner):
|
||||
result = runner.invoke(main, ["visible", "--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "Gemini watermark" in result.output
|
||||
|
||||
def test_visible_basic(self, runner, sample_png, tmp_path):
|
||||
output = tmp_path / "clean.png"
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["visible", str(sample_png), "-o", str(output), "--no-detect"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
assert output.exists()
|
||||
assert "Saved" in result.output
|
||||
|
||||
def test_visible_default_output_name(self, runner, sample_png):
|
||||
result = runner.invoke(main, ["visible", str(sample_png), "--no-detect"])
|
||||
assert result.exit_code == 0
|
||||
expected = sample_png.with_stem(sample_png.stem + "_clean")
|
||||
assert expected.exists()
|
||||
|
||||
def test_visible_no_inpaint(self, runner, sample_png, tmp_path):
|
||||
output = tmp_path / "clean.png"
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"visible",
|
||||
str(sample_png),
|
||||
"-o",
|
||||
str(output),
|
||||
"--no-inpaint",
|
||||
"--no-detect",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
assert output.exists()
|
||||
|
||||
def test_visible_no_detect(self, runner, sample_png, tmp_path):
|
||||
output = tmp_path / "clean.png"
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["visible", str(sample_png), "-o", str(output), "--no-detect"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
|
||||
def test_visible_nonexistent_file(self, runner):
|
||||
result = runner.invoke(main, ["visible", "/nonexistent/file.png"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
|
||||
class TestInvisibleCommand:
|
||||
"""Tests for the 'invisible' subcommand."""
|
||||
|
||||
def test_invisible_help(self, runner):
|
||||
result = runner.invoke(main, ["invisible", "--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "invisible" in result.output.lower()
|
||||
|
||||
def test_invisible_basic(self, runner, sample_png, tmp_path):
|
||||
mock_cls, mock_engine = _mock_invisible_engine()
|
||||
output = tmp_path / "clean.png"
|
||||
with patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch(
|
||||
"remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls
|
||||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["invisible", str(sample_png), "-o", str(output)],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert output.exists()
|
||||
mock_engine.remove_watermark.assert_called_once()
|
||||
|
||||
def test_invisible_default_output(self, runner, sample_png):
|
||||
mock_cls, mock_engine = _mock_invisible_engine()
|
||||
with patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch(
|
||||
"remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls
|
||||
):
|
||||
result = runner.invoke(main, ["invisible", str(sample_png)])
|
||||
assert result.exit_code == 0, result.output
|
||||
expected = sample_png.with_stem(sample_png.stem + "_clean")
|
||||
assert expected.exists()
|
||||
|
||||
def test_invisible_nonexistent_file(self, runner):
|
||||
result = runner.invoke(main, ["invisible", "/nonexistent/file.png"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
|
||||
class TestAllCommand:
|
||||
"""Tests for the 'all' subcommand (full pipeline)."""
|
||||
|
||||
def test_all_help(self, runner):
|
||||
result = runner.invoke(main, ["all", "--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "visible" in result.output.lower()
|
||||
|
||||
def test_all_basic(self, runner, sample_png, tmp_path):
|
||||
mock_cls, mock_engine = _mock_invisible_engine()
|
||||
output = tmp_path / "clean.png"
|
||||
with patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch(
|
||||
"remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls
|
||||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["all", str(sample_png), "-o", str(output)],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert output.exists()
|
||||
|
||||
def test_all_nonexistent_file(self, runner):
|
||||
result = runner.invoke(main, ["all", "/nonexistent/file.png"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
|
||||
class TestMetadataCommand:
|
||||
"""Tests for the 'metadata' subcommand."""
|
||||
|
||||
def test_metadata_help(self, runner):
|
||||
result = runner.invoke(main, ["metadata", "--help"])
|
||||
assert result.exit_code == 0
|
||||
|
||||
def test_metadata_check_clean(self, runner, tmp_clean_png):
|
||||
result = runner.invoke(main, ["metadata", str(tmp_clean_png), "--check"])
|
||||
assert result.exit_code == 0
|
||||
assert "No AI metadata" in result.output
|
||||
|
||||
def test_metadata_check_ai(self, runner, tmp_png_with_ai_metadata):
|
||||
result = runner.invoke(main, ["metadata", str(tmp_png_with_ai_metadata), "--check"])
|
||||
assert result.exit_code == 0
|
||||
assert "AI metadata detected" in result.output
|
||||
|
||||
def test_metadata_remove(self, runner, tmp_png_with_ai_metadata, tmp_path):
|
||||
output = tmp_path / "stripped.png"
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"metadata",
|
||||
str(tmp_png_with_ai_metadata),
|
||||
"--remove",
|
||||
"-o",
|
||||
str(output),
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
assert "stripped" in result.output
|
||||
|
||||
|
||||
class TestBatchCommand:
|
||||
"""Tests for the 'batch' subcommand."""
|
||||
|
||||
def test_batch_help(self, runner):
|
||||
result = runner.invoke(main, ["batch", "--help"])
|
||||
assert result.exit_code == 0
|
||||
|
||||
def test_batch_empty_dir(self, runner, tmp_path):
|
||||
empty_dir = tmp_path / "empty"
|
||||
empty_dir.mkdir()
|
||||
result = runner.invoke(main, ["batch", str(empty_dir)])
|
||||
assert result.exit_code == 0
|
||||
assert "No supported images" in result.output
|
||||
|
||||
def test_batch_visible_mode(self, runner, tmp_path):
|
||||
input_dir = _make_batch_dir(tmp_path)
|
||||
output_dir = tmp_path / "output"
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["batch", str(input_dir), "-o", str(output_dir), "--mode", "visible"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
assert "3 processed" in result.output
|
||||
assert output_dir.exists()
|
||||
assert len(list(output_dir.glob("*.png"))) == 3
|
||||
|
||||
def test_batch_metadata_mode(self, runner, tmp_path):
|
||||
input_dir = _make_batch_dir_with_metadata(tmp_path)
|
||||
output_dir = tmp_path / "output"
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["batch", str(input_dir), "-o", str(output_dir), "--mode", "metadata"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
assert "3 processed" in result.output
|
||||
assert output_dir.exists()
|
||||
assert len(list(output_dir.glob("*.png"))) == 3
|
||||
# Verify AI metadata was stripped
|
||||
for out_img in output_dir.glob("*.png"):
|
||||
with Image.open(out_img) as img:
|
||||
assert "parameters" not in img.info
|
||||
|
||||
def test_batch_invisible_mode(self, runner, tmp_path):
|
||||
input_dir = _make_batch_dir(tmp_path)
|
||||
output_dir = tmp_path / "output"
|
||||
mock_cls, mock_engine = _mock_invisible_engine()
|
||||
with patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch(
|
||||
"remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls
|
||||
), patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True), patch(
|
||||
"remove_ai_watermarks.invisible_engine.is_available", return_value=True
|
||||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["batch", str(input_dir), "-o", str(output_dir), "--mode", "invisible"],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "3 processed" in result.output
|
||||
|
||||
def test_batch_all_mode(self, runner, tmp_path):
|
||||
input_dir = _make_batch_dir(tmp_path)
|
||||
output_dir = tmp_path / "output"
|
||||
mock_cls, mock_engine = _mock_invisible_engine()
|
||||
with patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch(
|
||||
"remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls
|
||||
), patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True), patch(
|
||||
"remove_ai_watermarks.invisible_engine.is_available", return_value=True
|
||||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["batch", str(input_dir), "-o", str(output_dir), "--mode", "all"],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "3 processed" in result.output
|
||||
|
||||
def test_batch_default_output_dir(self, runner, tmp_path):
|
||||
input_dir = _make_batch_dir(tmp_path)
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["batch", str(input_dir), "--mode", "visible"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
expected_dir = tmp_path / "input_clean"
|
||||
assert expected_dir.exists()
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
import numpy as np
|
||||
|
||||
from remove_ai_watermarks.face_protector import FaceProtector
|
||||
|
||||
|
||||
def test_face_protector_initialization():
|
||||
# Will fallback to Haar cascade if ultralytics is missing
|
||||
fp = FaceProtector(use_yolo=False)
|
||||
assert fp.use_yolo is False
|
||||
assert fp.haar_cascade is not None
|
||||
|
||||
|
||||
def test_face_protector_lifecycle():
|
||||
fp = FaceProtector(use_yolo=False)
|
||||
|
||||
# Create dummy black image
|
||||
img = np.zeros((200, 200, 3), dtype=np.uint8)
|
||||
|
||||
# Since it's a black image, haar cascade should find 0 faces
|
||||
faces = fp.extract_faces(img)
|
||||
assert isinstance(faces, list)
|
||||
assert len(faces) == 0
|
||||
|
||||
# Restoring 0 faces should result in strictly equal image
|
||||
restored = fp.restore_faces(img, faces)
|
||||
assert np.array_equal(img, restored)
|
||||
|
||||
|
||||
def test_face_protector_restore_bypass_on_size_mismatch():
|
||||
fp = FaceProtector(use_yolo=False)
|
||||
img_small = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
# Manually mock a face that is OUT OF BOUNDS for img_small
|
||||
mock_bbox = (80, 80, 130, 130)
|
||||
mock_crop = np.ones((50, 50, 3), dtype=np.uint8) * 255
|
||||
mock_faces = [(mock_bbox, mock_crop)]
|
||||
|
||||
# Attempt to restore onto an image too small for this box
|
||||
restored = fp.restore_faces(img_small, mock_faces)
|
||||
|
||||
# Should safely skip restoring and not crash
|
||||
assert np.array_equal(restored, img_small)
|
||||
|
||||
|
||||
def test_face_protector_restore_blending():
|
||||
fp = FaceProtector(use_yolo=False)
|
||||
# Background is black
|
||||
img_target = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
# Face crop is white
|
||||
mock_bbox = (25, 25, 75, 75)
|
||||
mock_crop = np.ones((50, 50, 3), dtype=np.uint8) * 255
|
||||
mock_faces = [(mock_bbox, mock_crop)]
|
||||
|
||||
restored = fp.restore_faces(img_target, mock_faces)
|
||||
|
||||
# The center of the face should be perfectly white (255)
|
||||
assert restored[50, 50, 0] >= 254
|
||||
# The corner of the target should remain perfectly black (0)
|
||||
assert restored[0, 0, 0] == 0
|
||||
# We should have a blending gradient between them due to the gaussian blur mask
|
||||
# For example, around (30, 30) or similar
|
||||
assert 0 <= restored[28, 28, 0] <= 255
|
||||
@@ -0,0 +1,216 @@
|
||||
"""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,
|
||||
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 1024×1024 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)
|
||||
|
||||
|
||||
# ── 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)
|
||||
@@ -0,0 +1,52 @@
|
||||
import numpy as np
|
||||
|
||||
from remove_ai_watermarks.humanizer import apply_analog_humanizer
|
||||
|
||||
|
||||
def test_humanizer_does_not_modify_original_if_disabled():
|
||||
img = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
img[50, 50] = [100, 150, 200]
|
||||
org_img = img.copy()
|
||||
|
||||
# grain=0, shift=0 means disabled essentially. But wait, apply_analog_humanizer currently applies chromatic shift even if grain=0.
|
||||
result = apply_analog_humanizer(img, grain_intensity=0.0, chromatic_shift=0)
|
||||
assert np.array_equal(result, org_img)
|
||||
|
||||
|
||||
def test_chromatic_shift():
|
||||
# Only green channel is centered, red/blue should shift.
|
||||
img = np.zeros((5, 5, 3), dtype=np.uint8)
|
||||
img[2, 2] = [255, 255, 255] # B, G, R
|
||||
|
||||
# shift=1
|
||||
result = apply_analog_humanizer(img, grain_intensity=0.0, chromatic_shift=1)
|
||||
|
||||
# G (index 1) stays at [2,2]
|
||||
assert result[2, 2, 1] == 255
|
||||
# B (index 0) shifted right (+1 axis 1) -> [2, 3]
|
||||
assert result[2, 3, 0] == 255
|
||||
# R (index 2) shifted left (-1 axis 1) -> [2, 1]
|
||||
assert result[2, 1, 2] == 255
|
||||
|
||||
|
||||
def test_grain_intensity():
|
||||
# Gray image
|
||||
img = np.full((100, 100, 3), 128, dtype=np.uint8)
|
||||
|
||||
# Add strong noise
|
||||
result = apply_analog_humanizer(img, grain_intensity=10.0, chromatic_shift=0)
|
||||
|
||||
# Image should no longer be purely 128
|
||||
unique_vals = np.unique(result)
|
||||
assert len(unique_vals) > 5
|
||||
|
||||
# Mean should roughly be 128
|
||||
assert 126 < np.mean(result) < 130
|
||||
|
||||
|
||||
def test_invalid_shape():
|
||||
# Missing color channel
|
||||
img = np.zeros((100, 100), dtype=np.uint8)
|
||||
img[0, 0] = 50
|
||||
result = apply_analog_humanizer(img)
|
||||
assert np.array_equal(img, result)
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Tests for the invisible watermark engine (unit tests, no GPU required)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from remove_ai_watermarks.invisible_engine import InvisibleEngine, is_available
|
||||
|
||||
|
||||
class TestIsAvailable:
|
||||
"""Tests for dependency checking."""
|
||||
|
||||
def test_returns_bool(self):
|
||||
result = is_available()
|
||||
assert isinstance(result, bool)
|
||||
|
||||
def test_available_when_torch_installed(self):
|
||||
"""torch + diffusers should be installed in dev env."""
|
||||
assert is_available() is True
|
||||
|
||||
|
||||
class TestInvisibleEngineInit:
|
||||
"""Tests for InvisibleEngine construction (no GPU required)."""
|
||||
|
||||
def test_default_model_id(self):
|
||||
assert InvisibleEngine.DEFAULT_MODEL_ID == "Lykon/dreamshaper-8"
|
||||
|
||||
def test_ctrlregen_model_id(self):
|
||||
assert InvisibleEngine.CTRLREGEN_MODEL_ID == "yepengliu/ctrlregen"
|
||||
@@ -0,0 +1,150 @@
|
||||
"""Tests for AI metadata detection and removal."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
from remove_ai_watermarks.metadata import (
|
||||
_is_ai_key,
|
||||
get_ai_metadata,
|
||||
has_ai_metadata,
|
||||
remove_ai_metadata,
|
||||
)
|
||||
|
||||
# ── Key detection ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestIsAiKey:
|
||||
"""Tests for _is_ai_key helper."""
|
||||
|
||||
def test_exact_match_lowercase(self):
|
||||
assert _is_ai_key("parameters")
|
||||
|
||||
def test_exact_match_mixed_case(self):
|
||||
assert _is_ai_key("Parameters")
|
||||
|
||||
def test_keyword_substring(self):
|
||||
assert _is_ai_key("stable_diffusion_model_v2")
|
||||
|
||||
def test_c2pa_detected(self):
|
||||
assert _is_ai_key("c2pa_chunk")
|
||||
|
||||
def test_standard_key_not_flagged(self):
|
||||
assert not _is_ai_key("Author")
|
||||
|
||||
def test_innocuous_key_not_flagged(self):
|
||||
assert not _is_ai_key("Title")
|
||||
|
||||
def test_dpi_not_flagged(self):
|
||||
assert not _is_ai_key("dpi")
|
||||
|
||||
|
||||
# ── has_ai_metadata / get_ai_metadata ───────────────────────────────
|
||||
|
||||
|
||||
class TestHasAiMetadata:
|
||||
"""Tests for detecting AI metadata in images."""
|
||||
|
||||
def test_detects_ai_metadata(self, tmp_png_with_ai_metadata):
|
||||
assert has_ai_metadata(tmp_png_with_ai_metadata)
|
||||
|
||||
def test_clean_image_no_ai(self, tmp_clean_png):
|
||||
assert not has_ai_metadata(tmp_clean_png)
|
||||
|
||||
|
||||
class TestGetAiMetadata:
|
||||
"""Tests for extracting AI metadata."""
|
||||
|
||||
def test_extracts_parameters_key(self, tmp_png_with_ai_metadata):
|
||||
meta = get_ai_metadata(tmp_png_with_ai_metadata)
|
||||
assert "parameters" in meta
|
||||
assert "Euler" in meta["parameters"]
|
||||
|
||||
def test_extracts_prompt_key(self, tmp_png_with_ai_metadata):
|
||||
meta = get_ai_metadata(tmp_png_with_ai_metadata)
|
||||
assert "prompt" in meta
|
||||
|
||||
def test_does_not_extract_author(self, tmp_png_with_ai_metadata):
|
||||
meta = get_ai_metadata(tmp_png_with_ai_metadata)
|
||||
assert "Author" not in meta
|
||||
|
||||
def test_clean_image_empty_dict(self, tmp_clean_png):
|
||||
meta = get_ai_metadata(tmp_clean_png)
|
||||
assert meta == {}
|
||||
|
||||
|
||||
# ── remove_ai_metadata ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestRemoveAiMetadata:
|
||||
"""Tests for stripping AI metadata."""
|
||||
|
||||
def test_removes_ai_keys(self, tmp_png_with_ai_metadata):
|
||||
output = tmp_png_with_ai_metadata.parent / "cleaned.png"
|
||||
remove_ai_metadata(tmp_png_with_ai_metadata, output)
|
||||
|
||||
with Image.open(output) as img:
|
||||
assert "parameters" not in img.info
|
||||
assert "prompt" not in img.info
|
||||
|
||||
def test_keeps_standard_metadata(self, tmp_png_with_ai_metadata):
|
||||
output = tmp_png_with_ai_metadata.parent / "cleaned.png"
|
||||
remove_ai_metadata(tmp_png_with_ai_metadata, output, keep_standard=True)
|
||||
|
||||
with Image.open(output) as img:
|
||||
assert "Author" in img.info
|
||||
assert img.info["Author"] == "Test Author"
|
||||
|
||||
def test_remove_all_metadata(self, tmp_png_with_ai_metadata):
|
||||
output = tmp_png_with_ai_metadata.parent / "cleaned.png"
|
||||
remove_ai_metadata(tmp_png_with_ai_metadata, output, keep_standard=False)
|
||||
with Image.open(output) as img:
|
||||
assert "Author" not in img.info
|
||||
assert "parameters" not in img.info
|
||||
|
||||
def test_overwrite_in_place(self, tmp_path):
|
||||
"""When output_path is None, should overwrite source."""
|
||||
img = Image.new("RGB", (32, 32))
|
||||
pnginfo = PngInfo()
|
||||
pnginfo.add_text("parameters", "test data")
|
||||
path = tmp_path / "inplace.png"
|
||||
img.save(path, pnginfo=pnginfo)
|
||||
|
||||
result = remove_ai_metadata(path)
|
||||
assert result == path
|
||||
|
||||
with Image.open(path) as cleaned:
|
||||
assert "parameters" not in cleaned.info
|
||||
|
||||
def test_jpeg_output(self, tmp_path):
|
||||
"""Test metadata removal for JPEG format."""
|
||||
img = Image.new("RGB", (64, 64), color=(100, 150, 200))
|
||||
pnginfo = PngInfo()
|
||||
pnginfo.add_text("parameters", "test")
|
||||
png_path = tmp_path / "source.png"
|
||||
img.save(png_path, pnginfo=pnginfo)
|
||||
|
||||
jpg_path = tmp_path / "output.jpg"
|
||||
result = remove_ai_metadata(png_path, jpg_path)
|
||||
assert result == jpg_path
|
||||
assert jpg_path.exists()
|
||||
|
||||
def test_creates_parent_directories(self, tmp_path):
|
||||
img = Image.new("RGB", (32, 32))
|
||||
pnginfo = PngInfo()
|
||||
pnginfo.add_text("prompt", "test")
|
||||
path = tmp_path / "source.png"
|
||||
img.save(path, pnginfo=pnginfo)
|
||||
|
||||
output = tmp_path / "sub" / "dir" / "cleaned.png"
|
||||
remove_ai_metadata(path, output)
|
||||
assert output.exists()
|
||||
|
||||
def test_returns_path(self, tmp_clean_png):
|
||||
output = tmp_clean_png.parent / "out.png"
|
||||
result = remove_ai_metadata(tmp_clean_png, output)
|
||||
assert isinstance(result, Path)
|
||||
assert result == output
|
||||
@@ -0,0 +1,130 @@
|
||||
"""Tests for vendored noai submodules: constants, extractor, cleaner, c2pa."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from remove_ai_watermarks.noai.c2pa import (
|
||||
extract_c2pa_chunk,
|
||||
extract_c2pa_info,
|
||||
has_c2pa_metadata,
|
||||
)
|
||||
from remove_ai_watermarks.noai.cleaner import (
|
||||
has_ai_content,
|
||||
)
|
||||
from remove_ai_watermarks.noai.cleaner import (
|
||||
remove_ai_metadata as noai_remove_ai_metadata,
|
||||
)
|
||||
from remove_ai_watermarks.noai.constants import (
|
||||
AI_KEYWORDS,
|
||||
AI_METADATA_KEYS,
|
||||
C2PA_CHUNK_TYPE,
|
||||
PNG_SIGNATURE,
|
||||
SUPPORTED_FORMATS,
|
||||
)
|
||||
from remove_ai_watermarks.noai.extractor import (
|
||||
extract_ai_metadata,
|
||||
extract_metadata,
|
||||
get_ai_metadata_summary,
|
||||
has_ai_metadata,
|
||||
)
|
||||
|
||||
# ── Constants ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestConstants:
|
||||
"""Verify constant integrity."""
|
||||
|
||||
def test_supported_formats_include_png(self):
|
||||
assert ".png" in SUPPORTED_FORMATS
|
||||
|
||||
def test_supported_formats_include_jpg(self):
|
||||
assert ".jpg" in SUPPORTED_FORMATS
|
||||
|
||||
def test_ai_metadata_keys_not_empty(self):
|
||||
assert len(AI_METADATA_KEYS) > 0
|
||||
|
||||
def test_ai_keywords_not_empty(self):
|
||||
assert len(AI_KEYWORDS) > 0
|
||||
|
||||
def test_png_signature_bytes(self):
|
||||
assert PNG_SIGNATURE == b"\x89PNG\r\n\x1a\n"
|
||||
|
||||
def test_c2pa_chunk_type(self):
|
||||
assert C2PA_CHUNK_TYPE == b"caBX"
|
||||
|
||||
|
||||
# ── Extractor ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestExtractor:
|
||||
"""Tests for noai.extractor functions."""
|
||||
|
||||
def test_extract_metadata_returns_dict(self, tmp_clean_png):
|
||||
meta = extract_metadata(tmp_clean_png)
|
||||
assert isinstance(meta, dict)
|
||||
|
||||
def test_extract_metadata_gets_standard_keys(self, tmp_clean_png):
|
||||
meta = extract_metadata(tmp_clean_png)
|
||||
assert "Author" in meta
|
||||
|
||||
def test_extract_ai_metadata_from_ai_image(self, tmp_png_with_ai_metadata):
|
||||
meta = extract_ai_metadata(tmp_png_with_ai_metadata)
|
||||
assert "parameters" in meta
|
||||
|
||||
def test_extract_ai_metadata_from_clean_image(self, tmp_clean_png):
|
||||
meta = extract_ai_metadata(tmp_clean_png)
|
||||
assert len(meta) == 0
|
||||
|
||||
def test_has_ai_metadata_detects(self, tmp_png_with_ai_metadata):
|
||||
assert has_ai_metadata(tmp_png_with_ai_metadata)
|
||||
|
||||
def test_has_ai_metadata_clean(self, tmp_clean_png):
|
||||
assert not has_ai_metadata(tmp_clean_png)
|
||||
|
||||
def test_summary_with_ai(self, tmp_png_with_ai_metadata):
|
||||
summary = get_ai_metadata_summary(tmp_png_with_ai_metadata)
|
||||
assert "AI Image Metadata" in summary
|
||||
|
||||
def test_summary_clean(self, tmp_clean_png):
|
||||
summary = get_ai_metadata_summary(tmp_clean_png)
|
||||
assert "No AI metadata" in summary
|
||||
|
||||
|
||||
# ── Cleaner ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestCleaner:
|
||||
"""Tests for noai.cleaner functions."""
|
||||
|
||||
def test_remove_ai_metadata(self, tmp_png_with_ai_metadata, tmp_path):
|
||||
output = tmp_path / "cleaned.png"
|
||||
noai_remove_ai_metadata(tmp_png_with_ai_metadata, output)
|
||||
assert output.exists()
|
||||
# Verify AI metadata removed
|
||||
meta = extract_ai_metadata(output)
|
||||
assert "parameters" not in meta
|
||||
|
||||
def test_has_ai_content(self, tmp_png_with_ai_metadata):
|
||||
assert has_ai_content(tmp_png_with_ai_metadata)
|
||||
|
||||
|
||||
# ── C2PA ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestC2PA:
|
||||
"""Tests for C2PA detection on regular (non-C2PA) images."""
|
||||
|
||||
def test_no_c2pa_on_regular_png(self, tmp_clean_png):
|
||||
assert not has_c2pa_metadata(tmp_clean_png)
|
||||
|
||||
def test_no_c2pa_on_jpeg(self, tmp_jpeg_path):
|
||||
assert not has_c2pa_metadata(tmp_jpeg_path)
|
||||
|
||||
def test_extract_c2pa_none_on_regular(self, tmp_clean_png):
|
||||
assert extract_c2pa_chunk(tmp_clean_png) is None
|
||||
|
||||
def test_extract_c2pa_info_empty(self, tmp_clean_png):
|
||||
info = extract_c2pa_info(tmp_clean_png)
|
||||
assert info == {}
|
||||
|
||||
def test_c2pa_returns_false_for_non_png(self, tmp_jpeg_path):
|
||||
assert not has_c2pa_metadata(tmp_jpeg_path)
|
||||
@@ -0,0 +1,165 @@
|
||||
"""Tests for cross-platform and cross-device compatibility.
|
||||
|
||||
Verifies that device detection, MPS fallback, and platform-specific
|
||||
code paths work correctly on CPU, MPS (macOS), and CUDA (Linux/Windows).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from remove_ai_watermarks.noai.progress import is_mps_error
|
||||
from remove_ai_watermarks.noai.utils import get_image_format, is_supported_format
|
||||
from remove_ai_watermarks.noai.watermark_profiles import (
|
||||
detect_model_profile,
|
||||
get_model_id_for_profile,
|
||||
get_recommended_strength,
|
||||
)
|
||||
from remove_ai_watermarks.noai.watermark_remover import get_device, is_watermark_removal_available
|
||||
|
||||
# ── Device detection ────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestDeviceDetection:
|
||||
"""Tests for get_device() across platforms."""
|
||||
|
||||
def test_returns_valid_device(self):
|
||||
device = get_device()
|
||||
assert device in ("cpu", "mps", "cuda")
|
||||
|
||||
def test_cpu_fallback_when_no_gpu(self):
|
||||
"""On CI / machines without GPU, should fall back to cpu or mps."""
|
||||
device = get_device()
|
||||
# Just verify it doesn't crash and returns a valid string
|
||||
assert isinstance(device, str)
|
||||
|
||||
@patch("remove_ai_watermarks.noai.watermark_remover._HAS_TORCH", False)
|
||||
def test_no_torch_returns_cpu(self):
|
||||
assert get_device() == "cpu"
|
||||
|
||||
|
||||
class TestMpsErrorDetection:
|
||||
"""Tests for MPS error detection helper."""
|
||||
|
||||
def test_detects_mps_error(self):
|
||||
err = RuntimeError("MPS backend out of memory")
|
||||
assert is_mps_error(err) is True
|
||||
|
||||
def test_non_mps_error(self):
|
||||
err = RuntimeError("CUDA out of memory")
|
||||
assert is_mps_error(err) is False
|
||||
|
||||
def test_generic_error(self):
|
||||
err = RuntimeError("something went wrong")
|
||||
assert is_mps_error(err) is False
|
||||
|
||||
|
||||
# ── Model profiles ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestModelProfiles:
|
||||
"""Tests for watermark_profiles.py."""
|
||||
|
||||
def test_default_profile(self):
|
||||
assert get_model_id_for_profile("default") == "Lykon/dreamshaper-8"
|
||||
|
||||
def test_ctrlregen_profile(self):
|
||||
assert get_model_id_for_profile("ctrlregen") == "yepengliu/ctrlregen"
|
||||
|
||||
def test_unknown_profile_raises(self):
|
||||
with pytest.raises(ValueError, match="Unknown model profile"):
|
||||
get_model_id_for_profile("nonexistent")
|
||||
|
||||
def test_detect_default(self):
|
||||
assert detect_model_profile("Lykon/dreamshaper-8") == "default"
|
||||
|
||||
def test_detect_ctrlregen(self):
|
||||
assert detect_model_profile("yepengliu/ctrlregen") == "ctrlregen"
|
||||
|
||||
def test_recommended_strength_high(self):
|
||||
assert get_recommended_strength("treering") == 0.7
|
||||
|
||||
def test_recommended_strength_low(self):
|
||||
assert get_recommended_strength("stablesignature") == 0.04
|
||||
|
||||
def test_recommended_strength_medium(self):
|
||||
assert get_recommended_strength("unknown_type") == 0.35
|
||||
|
||||
|
||||
# ── Format utilities ────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestFormatUtils:
|
||||
"""Tests for utils.py format helpers."""
|
||||
|
||||
def test_supported_png(self, tmp_path):
|
||||
assert is_supported_format(tmp_path / "test.png")
|
||||
|
||||
def test_supported_jpg(self, tmp_path):
|
||||
assert is_supported_format(tmp_path / "test.jpg")
|
||||
|
||||
def test_supported_jpeg(self, tmp_path):
|
||||
assert is_supported_format(tmp_path / "test.jpeg")
|
||||
|
||||
def test_supported_webp(self, tmp_path):
|
||||
assert is_supported_format(tmp_path / "test.webp")
|
||||
|
||||
def test_unsupported_bmp(self, tmp_path):
|
||||
assert not is_supported_format(tmp_path / "test.bmp")
|
||||
|
||||
def test_unsupported_gif(self, tmp_path):
|
||||
assert not is_supported_format(tmp_path / "test.gif")
|
||||
|
||||
def test_get_format_png(self, tmp_path):
|
||||
assert get_image_format(tmp_path / "x.png") == "PNG"
|
||||
|
||||
def test_get_format_jpg(self, tmp_path):
|
||||
assert get_image_format(tmp_path / "x.jpg") == "JPEG"
|
||||
|
||||
def test_get_format_jpeg(self, tmp_path):
|
||||
assert get_image_format(tmp_path / "x.jpeg") == "JPEG"
|
||||
|
||||
def test_get_format_webp_defaults_png(self, tmp_path):
|
||||
# .webp falls through to PNG in current implementation
|
||||
assert get_image_format(tmp_path / "x.webp") == "PNG"
|
||||
|
||||
|
||||
# ── Availability checks ────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestAvailability:
|
||||
"""Tests for dependency availability checks."""
|
||||
|
||||
def test_watermark_removal_available(self):
|
||||
# In dev env with torch+diffusers installed
|
||||
assert is_watermark_removal_available() is True
|
||||
|
||||
def test_invisible_is_available(self):
|
||||
from remove_ai_watermarks.invisible_engine import is_available
|
||||
|
||||
assert is_available() is True
|
||||
|
||||
|
||||
# ── Platform-specific path handling ─────────────────────────────────
|
||||
|
||||
|
||||
class TestPlatformPaths:
|
||||
"""Verify path handling works on current platform."""
|
||||
|
||||
def test_pathlib_works_for_assets(self):
|
||||
from pathlib import Path
|
||||
|
||||
asset_dir = Path(__file__).parent.parent / "src" / "remove_ai_watermarks" / "assets"
|
||||
assert (asset_dir / "gemini_bg_48.png").exists()
|
||||
assert (asset_dir / "gemini_bg_96.png").exists()
|
||||
|
||||
def test_asset_loading_works(self):
|
||||
"""Verify embedded assets load correctly (critical for packaging)."""
|
||||
from remove_ai_watermarks.gemini_engine import GeminiEngine
|
||||
|
||||
engine = GeminiEngine()
|
||||
# If we get here without error, asset loading works
|
||||
assert engine._alpha_small.shape == (48, 48)
|
||||
assert engine._alpha_large.shape == (96, 96)
|
||||
Reference in New Issue
Block a user