fix: address whole-project code review (visible all/batch, engine consolidation, I/O)

Nine findings from a high-effort project-wide review, fixed and verified
(571 passed, ruff/pyright clean):

Correctness:
- all/batch now remove Doubao/Jimeng/Samsung visible text marks: the visible
  step routes through the registry (new cli._remove_visible_auto) instead of a
  hardcoded GeminiEngine, so they no longer leave the wordmark intact.
- batch always reads the original source (dropped the out_path-reuse that
  re-processed already-cleaned outputs on a re-run).
- img2img_runner only retries the diffusion call on the deprecated-callback
  TypeError; any other TypeError now propagates instead of double-running.
- gemini detect/remove and the reverse-alpha engines normalize channels via a
  new image_io.to_bgr, fixing a grayscale/BGRA crash in the FP-gate path.
- _png_late_metadata advances its cursor by the clamped length, so a malformed
  chunk length no longer aborts the late AI-label scan.

Cleanup / efficiency:
- Consolidate the ~90%-identical Doubao/Jimeng/Samsung engines into a shared
  config-driven _text_mark_engine.TextMarkEngine base; each engine is now a thin
  subclass (TextMarkConfig + test shims). Behavior is byte-exact (the three
  engine test suites pass unchanged). Registry adapters collapse to one
  _text_mark(...) row each. Gemini stays a separate engine.
- scan_head is memoized per (path, size, mtime), so identify() reads the file
  head once instead of ~8 times.
- invisible_engine post-processing decodes/encodes the output once (chained in
  memory) instead of 2-4 times across stages.
- Remove the orphaned get_model_id_for_profile (+ CONTROLNET_PROFILE); derive
  the --strength help from the strength constants (strength_default_help) so it
  cannot drift; share the --pipeline/--strength click options; simplify the
  retired --auto resolver.

Net -835 lines. Tests added for the registry-routed visible pass, to_bgr,
the polish/model/guidance wiring, and strength_default_help. CLAUDE.md updated
for the new base module, the engine/registry changes, image_io.to_bgr, and the
scan_head cache.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-06-09 13:21:13 -07:00
parent b1189549b8
commit 2fcd00ced0
17 changed files with 777 additions and 1264 deletions
+34 -38
View File
@@ -29,7 +29,6 @@ if TYPE_CHECKING:
from numpy.typing import NDArray
from remove_ai_watermarks.gemini_engine import DetectionResult
# --- plain-text output layer (replaces rich: no colors, no markup, no boxes) ---
@@ -291,15 +290,32 @@ def _warn_if_esrgan_unavailable(upscaler: str) -> None:
console.print(" Note: --upscaler esrgan needs the 'esrgan' extra; falling back to Lanczos.")
def _watermark_region(det: DetectionResult, width: int, height: int) -> tuple[int, int, int, int]:
"""Pick a watermark bbox: detector's region if confident, else the default config slot."""
if det.confidence > 0.15:
return det.region
from remove_ai_watermarks.gemini_engine import get_watermark_config
def _remove_visible_auto(
image: NDArray[Any],
*,
inpaint: bool = True,
inpaint_method: str = "ns",
inpaint_strength: float = 0.85,
) -> tuple[NDArray[Any], str | None]:
"""Remove the strongest auto-detected visible mark via the registry.
config = get_watermark_config(width, height)
px, py = config.get_position(width, height)
return (px, py, config.logo_size, config.logo_size)
Routes the ``all``/``batch`` visible step through the same registry path the
standalone ``visible`` command uses, so EVERY registered mark is handled (the
Gemini sparkle AND the Doubao/Jimeng/Samsung text marks), not just the sparkle.
Returns ``(result, label-or-None)``; when no ``in_auto`` mark fires the image is
returned unchanged with ``None``. ``inpaint*`` tune the Gemini edge-residual
cleanup only (the text engines ignore them).
"""
from remove_ai_watermarks import watermark_registry
best = watermark_registry.best_auto_mark(image)
if best is None:
return image, None
method: Literal["telea", "ns"] = "ns" if inpaint_method == "ns" else "telea"
result, _ = watermark_registry.get_mark(best.key).remove(
image, inpaint_method=method, inpaint=inpaint, inpaint_strength=inpaint_strength, force=False
)
return result, best.label
def _read_bgr_and_alpha(path: Path) -> tuple[NDArray[Any] | None, NDArray[Any] | None]:
@@ -893,8 +909,6 @@ def cmd_all(
If invisible watermark deps are not installed, skips step 2 with a warning.
"""
from remove_ai_watermarks.gemini_engine import GeminiEngine
_banner()
source = _validate_image(source)
_warn_if_esrgan_unavailable(upscaler)
@@ -918,7 +932,6 @@ def cmd_all(
# -- Step 1: Visible watermark --------------------------------
console.print("\n 1) Visible watermark removal")
engine = GeminiEngine()
image, alpha = _read_bgr_and_alpha(source)
if image is None:
console.print(f"Error: Failed to read image: {source}")
@@ -928,15 +941,10 @@ def cmd_all(
console.print(f" Input: {source.name} ({w}x{h})")
with console.status("Removing visible watermark..."):
det = engine.detect_watermark(image)
if det.detected:
result = engine.remove_watermark(image)
if inpaint:
region = _watermark_region(det, w, h)
result = engine.inpaint_residual(result, region, method=inpaint_method)
console.print(" Visible watermark removed")
result, removed_label = _remove_visible_auto(image, inpaint=inpaint, inpaint_method=inpaint_method)
if removed_label is not None:
console.print(f" Visible watermark removed ({removed_label})")
else:
result = image.copy()
console.print(" Skipped (no visible watermark detected)")
# Save to temp file for invisible engine input (preserve alpha if present)
@@ -1058,27 +1066,15 @@ def _process_batch_image(
saved_alpha: NDArray[Any] | None = None
if mode in ("visible", "all"):
from remove_ai_watermarks.gemini_engine import GeminiEngine
if "_vis_engine" not in ctx.obj:
ctx.obj["_vis_engine"] = GeminiEngine()
engine = ctx.obj["_vis_engine"]
read_path = img_path
if mode == "all" and out_path.exists():
read_path = out_path
image, alpha = _read_bgr_and_alpha(read_path)
# Always read the ORIGINAL source: the visible pass is the first step, so a
# stale out_path from a previous run must not be re-processed as if it were
# the input. (The invisible step below reads out_path for `all` -- that chain
# is within a single run.)
image, alpha = _read_bgr_and_alpha(img_path)
if image is None:
raise ValueError("Failed to read image")
det = engine.detect_watermark(image)
if det.detected:
result = engine.remove_watermark(image)
if inpaint:
h, w = image.shape[:2]
region = _watermark_region(det, w, h)
result = engine.inpaint_residual(result, region)
else:
result = image.copy()
result, _ = _remove_visible_auto(image, inpaint=inpaint)
_write_bgr_with_alpha(out_path, result, alpha)
saved_alpha = alpha