diff --git a/docs/cli.md b/docs/cli.md index 5241a33..90b8c3b 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -469,6 +469,11 @@ detected. Use `--force` when you know the image should be processed: remove-ai-watermarks invisible image.png -o clean.png --force ``` +For transparent inputs, use PNG or WebP output to retain transparency. The engine +preserves the original alpha channel, including partially transparent edges, +while regenerating RGB pixels. JPEG output cannot preserve transparency. This +also applies to `batch --mode invisible`. + ### Choose a strength cohort `--vendor` selects the cohort the default strength resolves from. `auto` diff --git a/docs/module-internals.md b/docs/module-internals.md index b83d915..55e9ebe 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -2160,6 +2160,13 @@ image sizing, postprocessing, and the public engine interface. It delegates model execution to [`_internal/watermark_remover.py`](../src/remove_ai_watermarks/_internal/watermark_remover.py). +The engine preserves source transparency for alpha-capable outputs. It captures +alpha after EXIF orientation is applied and before resizing, then rejoins the +unchanged alpha plane after restoring the RGB result to the original dimensions +and applying postprocessing. This also handles palette and colour-key PNG +transparency; diffusion itself still receives RGB. Regression coverage lives in +`test_invisible_engine.py::TestTransparency`. + `get_device` in that module answers only `cuda` or `cpu`. An `mps` or `xpu` answer would travel one frame to the same CUDA-only refusal while costing a device probe, and reporting it implied an Apple-silicon or Intel-GPU path that does not exist. diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index 559adac..5e57a8c 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/src/remove_ai_watermarks/invisible_engine.py @@ -311,6 +311,15 @@ class InvisibleEngine: with Image.open(image_path) as opened: image = ImageOps.exif_transpose(opened) orig_size = image.size # (width, height) + # Diffusion only processes RGB. Keep alpha at the original display geometry + # (after EXIF, before any resize) and rejoin it after all RGB postprocessing. + # Palette and colour-key PNGs can carry transparency without an A band. + alpha = ( + image.convert("RGBA").getchannel("A") + if output_path.suffix.lower() in image_io.ALPHA_FORMATS + and ("A" in image.getbands() or "transparency" in image.info) + else None + ) # Full-res original, kept for the adaptive-polish detail target (image is # reassigned to the resized copy below; PIL resize returns a new object). reference_pil = image @@ -369,7 +378,7 @@ class InvisibleEngine: # Route those outputs through the same final resize so --no-adaptive-polish does # not silently change e.g. 1448x1086 into 1448x1080. needs_restore = target is not None or any(dimension % 8 for dimension in orig_size) - if humanize > 0.0 or unsharp > 0.0 or adaptive_polish or needs_restore: + if humanize > 0.0 or unsharp > 0.0 or adaptive_polish or needs_restore or alpha is not None: import cv2 import numpy as np @@ -390,7 +399,13 @@ class InvisibleEngine: progress=self._progress_callback, ) - if not image_io.imwrite(out_path, out_cv, display_tags_from=image_path, orientation_applied=True): + if not image_io.write_bgr_with_alpha( + out_path, + out_cv, + np.asarray(alpha) if alpha is not None else None, + display_tags_from=image_path, + orientation_applied=True, + ): raise OSError(f"failed to write postprocessed output: {out_path}") return out_path diff --git a/tests/test_invisible_engine.py b/tests/test_invisible_engine.py index 1e6b71d..530a8b8 100644 --- a/tests/test_invisible_engine.py +++ b/tests/test_invisible_engine.py @@ -4,6 +4,8 @@ from __future__ import annotations from types import SimpleNamespace +import numpy as np +import pytest from PIL import Image from remove_ai_watermarks.invisible_engine import ( @@ -164,6 +166,126 @@ class TestNativeOutputSize: assert Image.open(out).size == (24, 18) +class TestTransparency: + """Exercise real image IO and postprocessing around an RGB-only model stand-in.""" + + @staticmethod + def _engine(): + engine = object.__new__(InvisibleEngine) + engine._progress_callback = None + + def remove_watermark(image_path, output_path, **_kwargs): + with Image.open(image_path) as image: + assert image.mode == "RGB" + width, height = image.size + # Model-side grid rounding and changed pixels, not a copy of the input. + Image.new("RGB", (width // 8 * 8, height // 8 * 8), (40, 120, 200)).save(output_path) + return output_path + + engine._remover = SimpleNamespace(configured_profile="qwen-zimage", remove_watermark=remove_watermark) + return engine + + @staticmethod + def _source(mode, size): + image = Image.new("RGB", size, (180, 90, 30)) + alpha = np.resize(np.array([0, 64, 128, 255], dtype=np.uint8), (size[1], size[0])) + if mode == "P": + image = Image.fromarray(alpha).convert("P") + image.info["transparency"] = bytes(range(256)) + elif mode == "RGB-key": + image.paste((0, 0, 0), (0, 0, size[0] // 2, size[1])) + image.info["transparency"] = (0, 0, 0) + else: + image = image.convert(mode) + image.putalpha(Image.fromarray(alpha)) + return image + + @pytest.mark.parametrize("mode", ["RGBA", "LA", "P", "RGB-key"]) + @pytest.mark.parametrize(("size", "max_resolution"), [((32, 24), 0), ((35, 27), 0), ((32, 24), 16)]) + def test_preserves_exact_alpha_after_model_and_size_restoration(self, tmp_path, mode, size, max_resolution): + source, output = tmp_path / "source.png", tmp_path / "output.png" + self._source(mode, size).save(source) + with Image.open(source) as image: + expected = np.asarray(image.convert("RGBA"))[:, :, 3].copy() + + self._engine().remove_watermark(source, output, max_resolution=max_resolution, adaptive_polish=False) + + with Image.open(output) as image: + assert image.mode == "RGBA" + assert image.size == size + actual = np.asarray(image) + np.testing.assert_array_equal(actual[:, :, 3], expected) + # Retain the model's RGB result rather than restoring the source image. + np.testing.assert_array_equal(actual[:, :, :3], np.broadcast_to([40, 120, 200], actual[:, :, :3].shape)) + + @pytest.mark.parametrize("orientation", [2, 6, 7]) + def test_in_place_output_keeps_alpha_in_display_orientation(self, tmp_path, orientation): + from PIL import ImageOps + + source = tmp_path / "source.png" + image = self._source("RGBA", (32, 24)) + image.getexif()[274] = orientation + image.save(source) + with Image.open(source) as opened: + expected = np.asarray(ImageOps.exif_transpose(opened).getchannel("A")).copy() + + self._engine().remove_watermark(source, adaptive_polish=False) + + with Image.open(source) as result: + assert result.mode == "RGBA" + assert result.getexif().get(274, 1) == 1 + np.testing.assert_array_equal(np.asarray(result.getchannel("A")), expected) + + def test_webp_keeps_alpha_after_postprocessing(self, tmp_path): + source, output = tmp_path / "source.png", tmp_path / "output.webp" + image = self._source("RGBA", (32, 24)) + image.save(source) + + self._engine().remove_watermark(source, output, unsharp=0.5, adaptive_polish=True, humanize=1.0) + + with Image.open(output) as result: + assert result.mode == "RGBA" + np.testing.assert_array_equal(np.asarray(result.getchannel("A")), np.asarray(image.getchannel("A"))) + + @pytest.mark.parametrize(("source_mode", "output_suffix"), [("RGB", ".png"), ("RGBA", ".jpg")]) + def test_rgb_and_non_alpha_output_remain_rgb(self, tmp_path, source_mode, output_suffix): + source, output = tmp_path / "source.png", tmp_path / ("output" + output_suffix) + Image.new(source_mode, (32, 24)).save(source) + + self._engine().remove_watermark(source, output, adaptive_polish=False) + + with Image.open(output) as result: + assert result.mode == "RGB" + + @pytest.mark.parametrize("command", ["invisible", "batch"]) + def test_cli_preserves_alpha_through_the_real_engine(self, tmp_path, monkeypatch, command): + from click.testing import CliRunner + + from remove_ai_watermarks import invisible_engine + from remove_ai_watermarks.cli import main + + source = tmp_path / "source.png" + image = self._source("RGBA", (32, 24)) + image.save(source) + engine = self._engine() + monkeypatch.setattr(invisible_engine, "InvisibleEngine", lambda **_kwargs: engine) + monkeypatch.setattr(invisible_engine, "is_available", lambda: True) + if command == "batch": + output_dir = tmp_path / "results" + args = ["batch", str(tmp_path), "--mode", "invisible", "-o", str(output_dir), "--force"] + output = output_dir / "source.png" + else: + output = tmp_path / "output.png" + args = ["invisible", str(source), "-o", str(output), "--force"] + + result = CliRunner().invoke(main, args) + + assert result.exit_code == 0, result.output + with Image.open(output) as processed: + assert processed.mode == "RGBA" + np.testing.assert_array_equal(np.asarray(processed.getchannel("A")), np.asarray(image.getchannel("A"))) + + class TestTargetSize: """Regression guard for the native-resolution decision (issues #10 / #15).