mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-06-10 20:57:47 +02:00
Victor Kuznetsov
6d11c11b52
Three content-quality features for the invisible/all/batch pipeline.
DBNet text detector (auto_config): replace the MSER text heuristic with
PP-OCRv3 differentiable-binarization via cv2.dnn.TextDetectionModel_DB,
using a bundled 2.4 MB Apache-2.0 model (en/cn detection nets are
byte-identical, so it ships language-neutral). cv2.dnn is core OpenCV, so
no new pip dep. MSER stays as the fallback when the model can't load.
Validated on real images: matches MSER everywhere and additionally catches
the Doubao CJK mark MSER missed; routing decisions unchanged otherwise.
Real-ESRGAN upscaler (new upscaler.py, esrgan extra): optional
pre-diffusion super-resolution for the min-resolution floor upscale, loaded
via spandrel (MIT, no basicsr) with BSD-3-Clause weights downloaded on
first use. New --upscaler {lanczos,esrgan} on invisible/all/batch; default
stays lanczos and the engine falls back to lanczos when the extra is absent
or the model errors (never breaks removal). It is a manual opt-in knob (the
auto plan never selects it) -- as a generic GAN it sharpens photo/texture
content strongly but can degrade faces (the diffusion pass regenerates
them) and thin text, documented accordingly.
batch --auto: wire the content-adaptive --auto (+ --adaptive-polish) into
cmd_batch. The plan is recomputed per image and the invisible engine is
cached per resolved pipeline (default/controlnet), so a mixed directory
builds at most one engine of each kind. Verified end-to-end: 3 mixed
images routed correctly with only 2 pipeline loads (controlnet reused).
ruff + strict pyright(src/) clean; 558 tests pass.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
33 lines
1.0 KiB
Python
33 lines
1.0 KiB
Python
"""Tests for the optional Real-ESRGAN upscaler (no model download).
|
|
|
|
The model-running path is exercised manually (it downloads ~67 MB of BSD-3-Clause
|
|
weights on first use); these tests cover the availability guard and the no-model
|
|
control flow, mirroring the repo convention for ML-adjacent modules.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from remove_ai_watermarks import upscaler
|
|
|
|
|
|
class TestIsAvailable:
|
|
def test_returns_bool(self):
|
|
assert isinstance(upscaler.is_available(), bool)
|
|
|
|
|
|
class TestUpscaleGuard:
|
|
def test_raises_without_extra(self, monkeypatch):
|
|
monkeypatch.setattr(upscaler, "is_available", lambda: False)
|
|
with pytest.raises(RuntimeError, match="esrgan"):
|
|
upscaler.upscale(np.full((32, 32, 3), 128, dtype=np.uint8))
|
|
|
|
|
|
class TestModelCachePath:
|
|
def test_cache_path_uses_model_filename(self):
|
|
if not upscaler.is_available():
|
|
pytest.skip("esrgan extra (torch) not installed")
|
|
assert upscaler._model_cache_path().name == upscaler._MODEL_FILENAME
|