Merge pull request #121 from dukagjina/fix/preserve-invisible-transparency

Preserve alpha in invisible image processing

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`

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.

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

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).