CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-06-05 10:38:00 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix #30 white box: stop zeroing alpha in the watermark region on save

Browse Source 
On RGBA inputs the CLI forced the watermark bbox alpha to 0 on save, so the
removed-sparkle area became a transparent hole that renders as a solid white
box on any non-transparent viewer. The Gemini app exports opaque RGBA, so
every user hit it. Reverse-alpha already recovers the real pixels there (and
`erase` inpaints them), so there is no artifact to hide -- the hole was the
bug, introduced as an over-correction in d091b9f.

`_write_bgr_with_alpha` now rejoins the input alpha plane unchanged (drops the
`clear_region`/`pad` params); the `visible` / `erase` / `all` / `batch` call
sites drop the cleared-region argument and the orphaned region bookkeeping.
The registry `remove()` still returns the mark bbox (used for inpaint_residual
positioning); the CLI just no longer clears alpha with it.

Inverts the test that locked in the old behavior into a #30 regression guard
(watermark-region alpha stays opaque, no pixel forced transparent). Verified
end-to-end on a real Gemini RGBA export: sparkle gone, zero transparent
pixels, clean over a white background.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-05-30 12:27:37 -07:00
parent 25a1acc53b
commit 89f427852f
4 changed files with 32 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+1 -1
View File
File diff suppressed because one or more lines are too long
src/remove_ai_watermarks/cli.py
+17 -32
View File
@@ -101,15 +101,15 @@ def _write_bgr_with_alpha(
path: Path,
bgr: NDArray[Any],
alpha: NDArray[Any] | None,
clear_region: tuple[int, int, int, int] | None = None,
pad: int = 6,
) -> None:
"""Write BGR (with optional alpha) to ``path``.
When ``alpha`` is provided and the output extension supports it, writes a
4-channel image. If ``clear_region`` is given as ``(x, y, w, h)``, alpha is
forced to 0 inside that bbox (expanded by ``pad`` px) so the watermark area
becomes fully transparent in the saved file.
When ``alpha`` is provided and the output extension supports it, the original
alpha plane is rejoined unchanged. The watermark region is NOT made
transparent: reverse-alpha (and inpaint) recover real pixels there, so
zeroing alpha would punch a transparent hole that renders as a white box on
any non-transparent viewer (issue #30). Preserving the input alpha keeps
genuinely transparent backgrounds intact without inventing new holes.
"""
import numpy as np
@@ -119,17 +119,7 @@ def _write_bgr_with_alpha(
image_io.imwrite(path, bgr)
return
alpha_out = alpha
if clear_region is not None:
alpha_out = alpha.copy()
x, y, w, h = clear_region
height, width = alpha.shape[:2]
x0, y0 = max(0, x - pad), max(0, y - pad)
x1, y1 = min(width, x + w + pad), min(height, y + h + pad)
if x1 > x0 and y1 > y0:
alpha_out[y0:y1, x0:x1] = 0
bgra = np.dstack([bgr, alpha_out])
bgra = np.dstack([bgr, alpha])
image_io.imwrite(path, bgra)
@@ -246,7 +236,7 @@ def cmd_visible(
method: Literal["telea", "ns"] = "ns" if inpaint_method == "ns" else "telea"
t0 = time.monotonic()
with console.status(f"[cyan]Removing {chosen.label}… ({chosen.recovery})[/]"):
result, region = chosen.remove(
result, _ = chosen.remove(
image,
inpaint_method=method,
inpaint=inpaint,
@@ -255,9 +245,9 @@ def cmd_visible(
)
elapsed = time.monotonic() - t0
# Save (preserves transparency by clearing alpha in the watermark region)
# Save (rejoins the original alpha plane unchanged)
output.parent.mkdir(parents=True, exist_ok=True)
_write_bgr_with_alpha(output, result, alpha, clear_region=region)
_write_bgr_with_alpha(output, result, alpha)
# Strip metadata
if strip_metadata:
@@ -349,8 +339,7 @@ def cmd_erase(
elapsed = time.monotonic() - t0
output.parent.mkdir(parents=True, exist_ok=True)
clear = boxes[0] if len(boxes) == 1 else None
_write_bgr_with_alpha(output, result, alpha, clear_region=clear)
_write_bgr_with_alpha(output, result, alpha)
if strip_metadata:
try:
@@ -695,7 +684,6 @@ def cmd_all(
h, w = image.shape[:2]
console.print(f" [dim]Input:[/] {source.name} ({w}x{h})")
region: tuple[int, int, int, int] | None = None
with console.status("[cyan]Removing visible watermark…[/]"):
det = engine.detect_watermark(image)
if det.detected:
@@ -709,7 +697,7 @@ def cmd_all(
console.print(" [dim]Skipped (no visible watermark detected)[/]")
# Save to temp file for invisible engine input (preserve alpha if present)
_write_bgr_with_alpha(tmp_path, result, alpha, clear_region=region)
_write_bgr_with_alpha(tmp_path, result, alpha)
# ── Step 2: Invisible watermark ──────────────────────────────
console.print("\n [bold cyan]② Invisible watermark removal[/]")
@@ -761,14 +749,14 @@ def cmd_all(
# ── Write final result ────────────────────────────────────────
# The invisible step (and downstream cv2.IMREAD_COLOR paths) drops alpha,
# so re-attach the original alpha (with the watermark region cleared)
# when writing the final output for transparent formats.
# so re-attach the original alpha plane unchanged when writing the final
# output for transparent formats.
output.parent.mkdir(parents=True, exist_ok=True)
final_bgr, _ = _read_bgr_and_alpha(tmp_path)
if final_bgr is None:
console.print(f"[red]Error:[/] Failed to read intermediate file: {tmp_path}")
raise SystemExit(1)
_write_bgr_with_alpha(output, final_bgr, alpha, clear_region=region)
_write_bgr_with_alpha(output, final_bgr, alpha)
finally:
# Clean up temp file if it still exists
@@ -808,7 +796,6 @@ def _process_batch_image(
ValueError: If the image cannot be opened.
"""
saved_alpha: NDArray[Any] | None = None
saved_region: tuple[int, int, int, int] | None = None
if mode in ("visible", "all"):
from remove_ai_watermarks.gemini_engine import GeminiEngine
@@ -823,7 +810,6 @@ def _process_batch_image(
if image is None:
raise ValueError("Failed to read image")
region: tuple[int, int, int, int] | None = None
det = engine.detect_watermark(image)
if det.detected:
result = engine.remove_watermark(image)
@@ -834,9 +820,8 @@ def _process_batch_image(
else:
result = image.copy()
_write_bgr_with_alpha(out_path, result, alpha, clear_region=region)
_write_bgr_with_alpha(out_path, result, alpha)
saved_alpha = alpha
saved_region = region
if mode in ("invisible", "all"):
from remove_ai_watermarks.invisible_engine import (
@@ -873,7 +858,7 @@ def _process_batch_image(
if mode == "all" and saved_alpha is not None:
final_bgr, _ = _read_bgr_and_alpha(out_path)
if final_bgr is not None:
_write_bgr_with_alpha(out_path, final_bgr, saved_alpha, clear_region=saved_region)
_write_bgr_with_alpha(out_path, final_bgr, saved_alpha)
@main.command("batch")
src/remove_ai_watermarks/watermark_registry.py
+4 -2
View File
@@ -71,8 +71,10 @@ class KnownMark:
inpaint_strength: float = 0.85,
force: bool = False,
) -> tuple[NDArray[Any], Region | None]:
"""Remove this mark by reverse-alpha; returns ``(result, cleared_region)``
(region for clearing alpha on save, or None if nothing was removed).
"""Remove this mark by reverse-alpha; returns ``(result, region)`` where
``region`` is the removed mark's bbox (for residual-inpaint positioning),
or None if nothing was removed. NB: the CLI does NOT use ``region`` to
clear alpha on save -- that zeroing caused the issue-#30 white box.
``inpaint`` / ``inpaint_strength`` / ``inpaint_method`` tune the Gemini
reverse-alpha edge-residual cleanup only. ``force`` removes at the mark's
tests/test_cli.py
+10 -6
View File
@@ -198,15 +198,17 @@ class TestVisibleCommand:
# which doesn't overlap the centre square at 200x200).
assert out[100, 100, 3] == 255
def test_visible_clears_alpha_in_watermark_region(self, runner, tmp_path):
"""When inpainting an RGBA image, the watermark region must be cleared
in the alpha channel so the sparkle area becomes transparent, not opaque-black.
def test_visible_keeps_alpha_opaque_in_watermark_region(self, runner, tmp_path):
"""Regression for issue #30 (white box): on an opaque RGBA image, the
watermark region must stay OPAQUE. Reverse-alpha recovers real pixels
there, so zeroing alpha would punch a transparent hole that renders as a
solid white box on any non-transparent viewer.
"""
rgba = np.full((200, 200, 4), 255, dtype=np.uint8) # fully opaque white
src = tmp_path / "rgba_full.png"
cv2.imwrite(str(src), rgba)
output = tmp_path / "rgba_cleared.png"
output = tmp_path / "rgba_kept.png"
result = runner.invoke(
main,
["visible", str(src), "-o", str(output), "--no-detect"],
@@ -215,13 +217,15 @@ class TestVisibleCommand:
assert result.exit_code == 0, result.output
out = cv2.imread(str(output), cv2.IMREAD_UNCHANGED)
assert out.shape[2] == 4
# Default sparkle position is in the bottom-right; alpha there must be 0.
# Default sparkle position is in the bottom-right; alpha there must stay 255.
from remove_ai_watermarks.gemini_engine import get_watermark_config
cfg = get_watermark_config(200, 200)
px, py = cfg.get_position(200, 200)
size = cfg.logo_size
assert out[py + size // 2, px + size // 2, 3] == 0, "alpha in the watermark region was not cleared"
assert out[py + size // 2, px + size // 2, 3] == 255, "watermark region alpha was zeroed (white-box regression)"
# No pixel anywhere should have been forced transparent.
assert int((out[:, :, 3] == 0).sum()) == 0, "spurious transparent pixels introduced"
def test_visible_rgb_input_stays_rgb(self, runner, sample_png, tmp_path):
"""Regression: a plain RGB PNG must NOT gain a spurious alpha channel."""