From 646f72534b255e8acdccb5f307ef1e84fe8c6c4b Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Sun, 6 Sep 2026 12:54:13 -0700 Subject: [PATCH] Carry the ICC profile and EXIF orientation across re-encodes Issue #98. cv2's encoders write no container metadata, so every pixel-path output (all, visible, invisible, erase) came back without its colour profile and orientation tag, even with --keep-metadata. imwrite now takes display_tags_from, the file whose decode produced the pixels, and restores the two display tags into the re-encoded output: PNG and JPEG by splicing iCCP/eXIf chunks and APP1/APP2 segments into the encoded bytes, WebP by a lossless Pillow re-save (libwebp hides hand-spliced chunks from readers without a VP8X header), HEIF through the save call itself. Orientation is re-tagged only when the raster still has the source's stored dimensions: cv2 rotates JPEG and (build-dependent) PNG decodes under IMREAD_COLOR but never IMREAD_UNCHANGED, which is what the pixel paths read with, so the decision is made against the raster's own shape and cannot double-rotate. The default metadata strip now preserves what the restore put back: the PIL path forwards the profile and a scrubbed EXIF (orientation included) for PNG and WebP outputs, the 16-bit chunk walker re-emits eXIf scrubbed instead of dropping it wholesale, and --remove-all actually drops the ICC profile on every format (Pillow's saver fallback kept it alive, contradicting the documented "removed by --remove-all"). Mirror and 180-degree orientations (2, 3, 4) stay untagged: no cheap signal distinguishes a rotated decode from a stored one for them. --- docs/cli.md | 12 +- docs/module-internals.md | 23 +- .../_internal/watermark_remover.py | 6 +- src/remove_ai_watermarks/api.py | 2 +- src/remove_ai_watermarks/cli.py | 17 +- src/remove_ai_watermarks/image_io.py | 280 +++++++++++++++++- src/remove_ai_watermarks/invisible_engine.py | 2 +- src/remove_ai_watermarks/metadata.py | 65 +++- tests/test_cli.py | 24 ++ tests/test_image_io.py | 197 ++++++++++++ tests/test_metadata.py | 106 +++++++ typings/piexif/__init__.pyi | 1 + 12 files changed, 704 insertions(+), 31 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 265e194..a077e5b 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -156,6 +156,9 @@ The default behavior: - checks the filled region once without changing it or retrying the fill; - strips AI metadata from the output. +The re-encode never costs display fidelity: the ICC colour profile and the EXIF +orientation tag are carried over from the source (see [strip AI metadata](#strip-ai-metadata)). + The automatic command reports `Removed and validated` when the detector no longer accepts an overlapping mark, `Partial` when an overlapping residual is still detected, and `Post-removal validation unavailable` when the check fails. @@ -236,7 +239,8 @@ pixels before inpainting, which helps when a mark has a soft edge or a drop shadow just outside the box you measured; it applies to every backend because it shapes the mask. `--inpaint-method telea|ns` selects the classical algorithm and only affects the `cv2` backend. Like `visible`, `erase` strips AI metadata from -the output by default; pass `--keep-metadata` to retain it. +the output by default; pass `--keep-metadata` to retain it. The ICC colour profile +and the EXIF orientation tag survive either way. ## Strip AI metadata @@ -253,7 +257,11 @@ remove-ai-watermarks metadata image.png --remove -o clean.png ``` When `-o` is omitted, removal overwrites the source. Standard metadata is kept -unless you pass `--remove-all`. +unless you pass `--remove-all`; that includes the display tags -- the ICC colour +profile and the EXIF orientation -- which also survive every pixel-command +re-encode (`all`, `visible`, `invisible`, `erase`), because colour and geometry +are display fidelity, not provenance. Orientation is re-tagged, not re-applied: +the raster is untouched, so what displayed upright still displays upright. A quiet `--check` or a successful `--remove` is not a clean verdict. The command only inspects and strips embedded AI metadata; a pixel watermark such as SynthID diff --git a/docs/module-internals.md b/docs/module-internals.md index bf69d47..44cf86e 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -2047,15 +2047,30 @@ Contracts: - `read_bgr_and_alpha` reads with `IMREAD_UNCHANGED`, so a 16-bit source stays 16-bit through the pixel path and out through `imwrite`. Anything downstream that needs 8 bits narrows a copy for itself. +- `imwrite` accepts `display_tags_from` naming the file whose decode produced the + pixels, and carries that file's ICC profile and EXIF orientation into the + re-encoded output (issue #98: cv2's encoders write no container metadata, so a + Display P3 portrait used to come back desaturated and sideways). PNG and JPEG + get the tags spliced into the encoded bytes without a second encode; WebP is + re-saved losslessly through Pillow (libwebp hides spliced chunks from readers + unless a hand-built VP8X announces them); HEIF writers bake orientation into + the pixels on save. Orientation is re-tagged only when the raster still has the + source's STORED dimensions, decided against the raster's own shape because cv2 + rotates JPEG and (build-dependent) PNG decodes under `IMREAD_COLOR` but never + `IMREAD_UNCHANGED`, which is what the pixel paths read with -- so a decode that + already turned the raster upright is never tagged into a double rotation. The metadata strip is the other half of that contract: Pillow cannot hold 16-bit colour, so `remove_ai_metadata`'s open+save path silently returned a 16-bit PNG at 8 bits. A PNG whose IHDR declares more than 8 bits now goes through `_strip_png_metadata_lossless`, which walks the chunk list, drops the AI-bearing -text chunks plus `eXIf` and `caBX`, and copies IDAT verbatim -- the PNG analogue of -`_strip_jpeg_metadata_lossless`. A standard `iCCP` colour profile survives the -default strip and is removed by `--remove-all`. The gate is the depth, not the -format: at 8 bits PIL is not lossy, so ordinary PNGs keep the shipped path. +text chunks plus `caBX`, re-emits `eXIf` scrubbed of AI tags, and copies IDAT +verbatim -- the PNG analogue of `_strip_jpeg_metadata_lossless`. A standard +`iCCP` colour profile survives the default strip and is removed by `--remove-all`; +the PIL path forwards the profile and the scrubbed EXIF (orientation included) for +8-bit PNG and WebP outputs too, where they used to be dropped with the rest of +EXIF. The gate is the depth, not the format: at 8 bits PIL is not lossy, so +ordinary PNGs keep the shipped path. Regression coverage: diff --git a/src/remove_ai_watermarks/_internal/watermark_remover.py b/src/remove_ai_watermarks/_internal/watermark_remover.py index c2cd414..336fa3d 100644 --- a/src/remove_ai_watermarks/_internal/watermark_remover.py +++ b/src/remove_ai_watermarks/_internal/watermark_remover.py @@ -199,14 +199,14 @@ class WatermarkRemover: ) return self._qwen_zimage_pipeline - def _write_output(self, image: Image.Image, output_path: Path) -> None: + def _write_output(self, image: Image.Image, output_path: Path, display_tags_from: Path | None = None) -> None: import numpy as np from remove_ai_watermarks import image_io output_path.parent.mkdir(parents=True, exist_ok=True) bgr = np.ascontiguousarray(np.asarray(image.convert("RGB"))[:, :, ::-1]) - if not image_io.imwrite(str(output_path), bgr): + if not image_io.imwrite(str(output_path), bgr, display_tags_from=display_tags_from): image.save(output_path) from remove_ai_watermarks.metadata import remove_ai_metadata @@ -285,5 +285,5 @@ class WatermarkRemover: text_manifest=text_manifest, fidelity_anchor=fidelity_anchor, ) - self._write_output(result, destination) + self._write_output(result, destination, display_tags_from=image_path) return destination diff --git a/src/remove_ai_watermarks/api.py b/src/remove_ai_watermarks/api.py index 7dd5238..31e8809 100644 --- a/src/remove_ai_watermarks/api.py +++ b/src/remove_ai_watermarks/api.py @@ -152,7 +152,7 @@ def _write_visible_result( # and surfaced as a confusing "cannot read image " naming the OUTPUT path # (Tier E, 2026-07-20). Raise here so a library caller and the CLI both get an # accurate message about the write. - if not image_io.write_bgr_with_alpha(out_path, result, loaded.alpha): + if not image_io.write_bgr_with_alpha(out_path, result, loaded.alpha, display_tags_from=source_path): raise OSError(f"failed to write output (is the destination writable?): {out_path}") if strip_metadata: diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 031f526..3e540ed 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -453,7 +453,12 @@ def _pixels_required(func: Callable[..., None]) -> Callable[..., None]: return wrapper -def _write_output_or_exit(output: Path, bgr: NDArray[Any], alpha: NDArray[Any] | None) -> None: +def _write_output_or_exit( + output: Path, + bgr: NDArray[Any], + alpha: NDArray[Any] | None, + display_source: Path | None = None, +) -> None: """Write the final image, or fail with a readable error instead of a traceback. `image_io.imwrite` is contractually NON-RAISING: it returns False when the codec @@ -462,9 +467,11 @@ def _write_output_or_exit(output: Path, bgr: NDArray[Any], alpha: NDArray[Any] | directory, full disk) died with a bare `FileNotFoundError` traceback pointing at the stat, not at the write. Found by the Tier E adversarial sweep 2026-07-20. Regression: `tests/test_cli_robustness.py::TestFailedWriteIsReported`. - """ + + `display_source` is the file whose decode produced `bgr`; its ICC profile and EXIF + orientation ride along into the re-encoded output (issue #98).""" output.parent.mkdir(parents=True, exist_ok=True) - if not image_io.write_bgr_with_alpha(output, bgr, alpha): + if not image_io.write_bgr_with_alpha(output, bgr, alpha, display_tags_from=display_source): console.print(f" Error: failed to write output (is the destination writable?): {output}") raise SystemExit(1) @@ -709,7 +716,7 @@ def _run_visible_explicit( raise SystemExit(1) from e elapsed = time.monotonic() - t0 - _write_output_or_exit(output, result, alpha) + _write_output_or_exit(output, result, alpha, display_source=source) if strip_metadata: try: from remove_ai_watermarks.metadata import remove_ai_metadata @@ -863,7 +870,7 @@ def cmd_erase( raise SystemExit(1) from e elapsed = time.monotonic() - t0 - _write_output_or_exit(output, result, alpha) + _write_output_or_exit(output, result, alpha, display_source=source) if strip_metadata: try: diff --git a/src/remove_ai_watermarks/image_io.py b/src/remove_ai_watermarks/image_io.py index 755c43f..4eab43a 100644 --- a/src/remove_ai_watermarks/image_io.py +++ b/src/remove_ai_watermarks/image_io.py @@ -21,12 +21,17 @@ stays cheap in a bare environment (matching the rest of the package). from __future__ import annotations +import logging from pathlib import Path from typing import TYPE_CHECKING, Any +from remove_ai_watermarks._internal.constants import PNG_SIGNATURE + if TYPE_CHECKING: from numpy.typing import NDArray +log = logging.getLogger(__name__) + def imread(path: str | Path, flags: int | None = None) -> NDArray[Any] | None: """Unicode-safe ``cv2.imread`` with a Pillow fallback for HEIC/AVIF. @@ -135,6 +140,228 @@ def _register_heif() -> None: pillow_heif.register_heif_opener() +# ── Display-tag carry across re-encodes (issue #98) ───────────────────────── +# +# cv2's encoders write no container metadata, so every pixel-path output used to +# come back without the source's ICC colour profile and EXIF orientation: a +# Display P3 portrait lost a chunk of its saturation and displayed sideways even +# though the user never asked the tool to touch metadata. Those two tags are +# display fidelity, not provenance, so they travel from the decode source to the +# re-encoded output. Only those two: wholesale EXIF copying would resurrect the +# AI-provenance tags the metadata strip exists to drop. + +_ORIENT_EXIF_TAG = 0x0112 # EXIF ImageIFD.Orientation +# Orientations that transpose width and height (the 90-degree rotations). Only these +# can be detected by comparing the raster's shape with the source's stored size; +# the mirror/180 values (2, 3, 4) keep today's tagless behaviour because no cheap +# signal distinguishes a rotated decode from a stored one for them. +_TRANSPOSING_ORIENTATIONS = frozenset({5, 6, 7, 8}) + + +def _orientation_exif(orient: int | None, *, prefixed: bool) -> bytes | None: + """EXIF bytes holding ONLY the orientation tag, in the form a container wants. + + ``piexif.dump`` returns the JPEG APP1 payload (``Exif\\x00\\x00`` + TIFF); the + PNG ``eXIf`` chunk and the WebP ``EXIF`` chunk carry the bare TIFF structure, so + the prefix is stripped there. Orientation 1 (upright) is the absence of a tag. + """ + if orient is None or orient == 1: + return None + import piexif + + dumped = piexif.dump({"0th": {piexif.ImageIFD.Orientation: orient}}) + if prefixed: + return dumped + return dumped[6:] if dumped.startswith(b"Exif\x00\x00") else dumped + + +def _read_display_tags(source: Path, raster_shape: tuple[int, int]) -> tuple[bytes | None, int | None]: + """Extract ``(icc_profile, orientation)`` from ``source`` for a raster of ``raster_shape``. + + The contract is that ``source`` is the file whose decode produced the pixels + being written, and ``raster_shape`` is that raster's ``(rows, cols)``. Whether the + decode turned the raster upright cannot be assumed from the container: cv2 + applies EXIF orientation for JPEG and (build-dependent) PNG -- but only under + ``IMREAD_COLOR``, never ``IMREAD_UNCHANGED``, which is what the pixel paths read + with -- and the reporter of issue #98 ran a build that left PNG flat too. So the + decision is made against the ground truth of the raster itself: for the + transposing orientations (5-8) the tag is carried exactly when the raster still + has the source's STORED dimensions, and dropped when it has the upright ones; + anything else (a resized pipeline output, the rare mirror/180 values) drops the + tag and keeps today's behaviour rather than risk a double rotation. ICC travels + in every case: it describes colour, not geometry. Failures read as "no tags": + the restore is an enhancement, never a reason to fail a write. + """ + _register_heif() + try: + from PIL import Image + except Exception: + return None, None + try: + with Image.open(source) as im: + icc = im.info.get("icc_profile") + orient = im.getexif().get(_ORIENT_EXIF_TAG) + stored_w, stored_h = im.size + except Exception: + return None, None + if not (isinstance(icc, bytes) and icc): + icc = None + if orient in _TRANSPOSING_ORIENTATIONS: + rows, cols = raster_shape + # Upright dimensions are the stored ones swapped; only a raster still in + # stored orientation may carry the tag. + if (cols, rows) != (stored_w, stored_h): + orient = None + else: + orient = None + return icc, orient + + +def _png_splice_display_tags(data: bytes, icc: bytes | None, orient: int | None) -> bytes: + """Insert ``iCCP`` / ``eXIf`` chunks after IHDR, leaving IDAT verbatim. + + Replaces same-kind chunks rather than duplicating them, so rewriting an + already-tagged PNG cannot stack profiles. Returns the input unchanged when the + chunk list does not walk cleanly (the file still decodes exactly as written). + """ + import struct + import zlib + + def _chunk(kind: bytes, payload: bytes) -> bytes: + return struct.pack(">I", len(payload)) + kind + payload + struct.pack(">I", zlib.crc32(kind + payload)) + + iccp = _chunk(b"iCCP", b"ICC Profile\x00\x00" + zlib.compress(icc, 9)) if icc else None + exif_payload = _orientation_exif(orient, prefixed=False) + exif_chunk = _chunk(b"eXIf", exif_payload) if exif_payload else None + if iccp is None and exif_chunk is None: + return data + + out = bytearray(data[:8]) + pos, n = 8, len(data) + inserted = False + while pos + 8 <= n: + (length,) = struct.unpack(">I", data[pos : pos + 4]) + kind = data[pos + 4 : pos + 8] + end = pos + 12 + length + if end > n: + return data # malformed length: leave the encoded bytes alone + replaced = (kind == b"iCCP" and iccp is not None) or (kind == b"eXIf" and exif_chunk is not None) + if not replaced: + out += data[pos:end] + if kind == b"IHDR" and not inserted: + out += (iccp or b"") + (exif_chunk or b"") + inserted = True + if kind == b"IEND": + break + pos = end + if not inserted: + return data + return bytes(out) + + +def _jpeg_splice_display_tags(data: bytes, icc: bytes | None, orient: int | None) -> bytes: + """Insert APP1 (EXIF) and APP2 (ICC) markers without touching the DCT scan. + + The new segments go at the end of the existing APP block (after JFIF APP0, ahead + of DQT/SOF), and same-purpose segments are replaced instead of duplicated. A large + ICC profile is chunked across APP2 markers per the ICC-in-JPEG spec: one-byte + sequence and total counters, the layout Pillow's writer and the spec share. + Returns the input unchanged when the marker walk does not reach the scan cleanly. + """ + import struct + + exif = _orientation_exif(orient, prefixed=True) + segments: list[bytes] = [] + if exif: + segments.append(b"\xff\xe1" + struct.pack(">H", len(exif) + 2) + exif) + if icc: + # 65535 per segment minus the 2-byte length and the 14-byte ICC header + # ("ICC_PROFILE\x00" + 1-byte sequence + 1-byte total). + part_size = 65535 - 2 - 14 + parts = [icc[i : i + part_size] for i in range(0, len(icc), part_size)] + for seq, part in enumerate(parts, start=1): + payload = b"ICC_PROFILE\x00" + bytes((seq, len(parts))) + part + segments.append(b"\xff\xe2" + struct.pack(">H", len(payload) + 2) + payload) + if not segments: + return data + + out = bytearray(data[:2]) # SOI + pos, n = 2, len(data) + while pos + 4 <= n: + if data[pos] != 0xFF: + return data # lost the marker boundary: leave the encoded bytes alone + marker = data[pos + 1] + if marker not in range(0xE0, 0xF0): + break # first segment outside the APP block: insertion point found + (seg_len,) = struct.unpack(">H", data[pos + 2 : pos + 4]) + end = pos + 2 + seg_len + if seg_len < 2 or end > n: + return data + payload = data[pos + 4 : end] + replaced = (marker == 0xE1 and exif and payload.startswith(b"Exif\x00\x00")) or ( + marker == 0xE2 and icc and payload.startswith(b"ICC_PROFILE\x00") + ) + if not replaced: + out += data[pos:end] + pos = end + out += b"".join(segments) + out += data[pos:] + return bytes(out) + + +def _webp_restore_display_tags(path: Path, icc: bytes | None, orient: int | None) -> None: + """Attach the display tags to a written WebP by re-saving it losslessly. + + A hand-spliced ``ICCP``/``EXIF`` RIFF chunk is invisible to readers: libwebp + surfaces metadata only when a ``VP8X`` header announces it, and stitching that + header by hand duplicates what Pillow's writer already does correctly (measured: + a spliced file fails to open in both Pillow and cv2). Pillow's WebP save in + lossless mode is pixel-identical, so the container bytes change and the picture + does not. + """ + try: + from PIL import Image + except Exception: + return + exif = _orientation_exif(orient, prefixed=True) + save_kwargs: dict[str, Any] = {"format": "WEBP", "lossless": True, "exact": True} + if icc: + save_kwargs["icc_profile"] = icc + if exif: + save_kwargs["exif"] = exif + try: + with Image.open(path) as im: + im.load() + im.save(path, **save_kwargs) + except Exception: + log.debug("webp display-tag restore failed on %s", path, exc_info=True) + + +def _restore_display_tags(path: Path, icc: bytes | None, orient: int | None) -> None: + """Best-effort splice of the display tags into an already-written output. + + Never raises and never re-encodes pixels: a failure leaves the un-tagged (but + valid) file the codec wrote, which is exactly the pre-fix behaviour. + """ + if icc is None and (orient is None or orient == 1): + return + try: + data = path.read_bytes() + except OSError: + log.debug("display-tag restore could not read back %s", path, exc_info=True) + return + try: + if data.startswith(PNG_SIGNATURE): + path.write_bytes(_png_splice_display_tags(data, icc, orient)) + elif data[:2] == b"\xff\xd8": + path.write_bytes(_jpeg_splice_display_tags(data, icc, orient)) + elif data[:4] == b"RIFF" and data[8:12] == b"WEBP": + _webp_restore_display_tags(path, icc, orient) + # else: a container without a restore path (BMP/TIFF): leave as written + except Exception: + log.debug("display-tag restore failed on %s", path, exc_info=True) + + # Containers cv2 cannot encode -> written via Pillow (pillow-heif). _HEIF_WRITE_EXTS = {".heic", ".heif", ".avif"} @@ -166,9 +393,17 @@ def _encode_params(ext: str) -> list[int]: return [] -def _pil_write(path: str | Path, img: NDArray[Any]) -> bool: +def _pil_write( + path: str | Path, + img: NDArray[Any], + *, + icc_profile: bytes | None = None, + orientation: int | None = None, +) -> bool: """Encode HEIC/AVIF via Pillow (+pillow-heif) at high quality -- cv2 has no encoder - for them. BGR / BGRA in; returns False if Pillow (with the plugin) cannot save.""" + for them. BGR / BGRA in; returns False if Pillow (with the plugin) cannot save. + The display tags, when given, ride along in the save call: Pillow's HEIF encoder + writes both natively, so this branch needs no post-write splice.""" import cv2 import numpy as np from PIL import Image @@ -178,26 +413,46 @@ def _pil_write(path: str | Path, img: NDArray[Any]) -> bool: arr, mode = cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA), "RGBA" else: arr, mode = cv2.cvtColor(to_bgr(img), cv2.COLOR_BGR2RGB), "RGB" + save_kwargs: dict[str, Any] = {"quality": 100} + if icc_profile: + save_kwargs["icc_profile"] = icc_profile + exif = _orientation_exif(orientation, prefixed=True) + if exif: + save_kwargs["exif"] = exif try: - Image.fromarray(np.ascontiguousarray(arr), mode).save(str(path), quality=100) + Image.fromarray(np.ascontiguousarray(arr), mode).save(str(path), **save_kwargs) return True except Exception: return False -def imwrite(path: str | Path, img: NDArray[Any]) -> bool: +def imwrite(path: str | Path, img: NDArray[Any], *, display_tags_from: str | Path | None = None) -> bool: """Unicode-safe image write that PRESERVES the input format at max quality. Format is taken from the path extension. HEIC/AVIF (which cv2 cannot encode) go through Pillow; everything else through cv2 with quality-preserving params (see :func:`_encode_params`) so a lossy re-encode of the untouched pixels stays near- lossless. Returns ``True`` on success, ``False`` if the codec rejects the image or - the path cannot be written (matching ``cv2.imwrite``, never raising).""" + the path cannot be written (matching ``cv2.imwrite``, never raising). + + When ``display_tags_from`` names the file whose decode produced ``img``'s pixels, + the source's ICC profile and EXIF orientation are carried into the re-encoded + output (issue #98): cv2's encoders write no container metadata, and colour and + geometry are display fidelity, not provenance. The tags are read BEFORE the write + so an in-place rewrite (source == output) still finds them, and orientation is + re-tagged only when the raster still has the source's stored dimensions (see + :func:`_read_display_tags`), so a decode that already turned the pixels upright + is never tagged into a second rotation.""" import cv2 + icc, orient = ( + _read_display_tags(Path(display_tags_from), (img.shape[0], img.shape[1])) + if display_tags_from is not None + else (None, None) + ) ext = (Path(path).suffix or ".png").lower() if ext in _HEIF_WRITE_EXTS: - return _pil_write(path, img) + return _pil_write(path, img, icc_profile=icc, orientation=orient) try: ok, buf = cv2.imencode(ext, img, _encode_params(ext)) except cv2.error: @@ -208,6 +463,7 @@ def imwrite(path: str | Path, img: NDArray[Any]) -> bool: buf.tofile(str(path)) except OSError: return False + _restore_display_tags(Path(path), icc, orient) return True @@ -234,7 +490,13 @@ def read_bgr_and_alpha(path: str | Path) -> tuple[NDArray[Any] | None, NDArray[A return image, None -def write_bgr_with_alpha(path: str | Path, bgr: NDArray[Any], alpha: NDArray[Any] | None) -> bool: +def write_bgr_with_alpha( + path: str | Path, + bgr: NDArray[Any], + alpha: NDArray[Any] | None, + *, + display_tags_from: str | Path | None = None, +) -> bool: """Write BGR (with optional alpha) to ``path``. Returns ``imwrite``'s success flag. When ``alpha`` is provided and the output extension supports it, the original @@ -253,5 +515,5 @@ def write_bgr_with_alpha(path: str | Path, bgr: NDArray[Any], alpha: NDArray[Any import numpy as np if alpha is None or Path(path).suffix.lower() not in ALPHA_FORMATS: - return imwrite(path, bgr) - return imwrite(path, np.dstack([bgr, alpha])) + return imwrite(path, bgr, display_tags_from=display_tags_from) + return imwrite(path, np.dstack([bgr, alpha]), display_tags_from=display_tags_from) diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index 4b183df..fa57a32 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/src/remove_ai_watermarks/invisible_engine.py @@ -380,7 +380,7 @@ class InvisibleEngine: progress=self._progress_callback, ) - image_io.imwrite(out_path, out_cv) + image_io.imwrite(out_path, out_cv, display_tags_from=image_path) return out_path finally: diff --git a/src/remove_ai_watermarks/metadata.py b/src/remove_ai_watermarks/metadata.py index e0d55b3..bdccfdb 100644 --- a/src/remove_ai_watermarks/metadata.py +++ b/src/remove_ai_watermarks/metadata.py @@ -1397,13 +1397,17 @@ def _strip_png_metadata_lossless(source_path: Path, output_path: Path, keep_stan the chunks instead copies the pixel stream verbatim, so the depth -- and every pixel-bearing chunk survives untouched. This is the PNG analogue of :func:`_strip_jpeg_metadata_lossless`, and it drops the same chunks the PIL path - drops: AI-bearing text chunks, ``eXIf`` (PNG output drops EXIF there too) and the - ``caBX`` C2PA store. The standard ``iCCP`` colour profile is preserved only when - ``keep_standard`` is true. + drops: AI-bearing text chunks, ``eXIf`` and the ``caBX`` C2PA store. The standard + ``iCCP`` colour profile is preserved only when ``keep_standard`` is true, and + ``eXIf`` is re-emitted scrubbed of AI tags rather than dropped, so a portrait's + display orientation survives the default strip (issue #98); ``--remove-all`` + drops both. Returns False when the bytes are not a walkable PNG, so the caller falls back to the PIL re-save. """ + import zlib + from remove_ai_watermarks.forensic_metadata import png_text_decode data = source_path.read_bytes() @@ -1419,7 +1423,19 @@ def _strip_png_metadata_lossless(source_path: Path, output_path: Path, keep_stan if length > n or end > n: return False # malformed length: defer to the PIL re-encode fallback payload = data[pos + 8 : pos + 8 + length] if chunk_type in (b"tEXt", b"zTXt", b"iTXt") else b"" - if not _png_chunk_should_be_dropped(chunk_type, payload, keep_standard, png_text_decode): + if chunk_type == b"eXIf": + # Dropped wholesale before, which took the display orientation with it + # (issue #98). keep_standard re-emits a scrubbed payload; an unparseable + # one still goes, fail-safe as before. + scrubbed = _scrubbed_png_exif(data[pos + 8 : pos + 8 + length]) if keep_standard else None + if scrubbed is not None: + out += ( + struct.pack(">I", len(scrubbed)) + + b"eXIf" + + scrubbed + + struct.pack(">I", zlib.crc32(b"eXIf" + scrubbed)) + ) + elif not _png_chunk_should_be_dropped(chunk_type, payload, keep_standard, png_text_decode): out += view[pos:end] if chunk_type == b"IEND": break @@ -1432,6 +1448,25 @@ def _strip_png_metadata_lossless(source_path: Path, output_path: Path, keep_stan return True +def _scrubbed_png_exif(payload: bytes) -> bytes | None: + """Scrub AI tags from a PNG ``eXIf`` payload, returning the bytes to keep. + + The payload is the bare TIFF structure, which piexif loads directly; the shared + AI scrub runs and the result is re-dumped in the same bare form. Returns None when + the payload does not parse as EXIF, so the caller drops it exactly as before. + """ + import piexif + + try: + exif = piexif.load(payload) + _scrub_ai_exif(exif) + dumped = piexif.dump(exif) + return dumped[6:] if dumped.startswith(b"Exif\x00\x00") else dumped + except Exception: + logger.debug("unparseable eXIf payload; dropped", exc_info=True) + return None + + def _keep_standard_text_metadata(key: str, value: object, keep_standard: bool) -> bool: """Whether a text metadata item is standard and free of AI provenance.""" return ( @@ -1722,15 +1757,33 @@ def remove_ai_metadata( pnginfo.add_text(k, v) save_kwargs["pnginfo"] = pnginfo - if exif_data and save_kwargs["format"] == "JPEG": + if keep_standard and isinstance(img.info.get("icc_profile"), bytes) and img.info["icc_profile"]: + # Forward the profile explicitly rather than trusting each saver's + # im.info fallback (PNG has one; the others are not contractual), so + # the default strip keeps the image's colour for every output format. + save_kwargs["icc_profile"] = img.info["icc_profile"] + if ( + exif_data + and save_kwargs["format"] in ("JPEG", "PNG", "WEBP") + and (save_kwargs["format"] == "JPEG" or keep_standard) + ): # Scrub AI-provenance EXIF tags (xAI/Grok signature, generator tokens) - # while keeping genuine camera/editor EXIF; PNG output drops EXIF entirely. + # while keeping genuine camera/editor EXIF, orientation included. PNG + # and WebP outputs used to drop EXIF entirely, which took the display + # orientation with it: a portrait came back tagged horizontal (issue + # #98). --remove-all still drops EXIF wholesale on those two. if removed := _scrub_ai_exif(exif_data): logger.info("Scrubbed AI EXIF tag(s): %s", ", ".join(removed)) with contextlib.suppress(Exception): save_kwargs["exif"] = piexif.dump(exif_data) output_path.parent.mkdir(parents=True, exist_ok=True) + if not keep_standard: + # Pillow's PNG/WebP/JPEG savers fall back to im.info["icc_profile"], so a + # profile would survive a strip that promised to remove everything. + # --remove-all drops it; the default strip keeps it (display fidelity, + # issue #98), which is also what the lossless walkers do. + img.info.pop("icc_profile", None) img.save(output_path, **save_kwargs) logger.info("Stripped AI metadata → %s", output_path) diff --git a/tests/test_cli.py b/tests/test_cli.py index 69c3803..31465cb 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1264,6 +1264,30 @@ class TestEraseCommand: assert result.exit_code == 0, result.output assert output.exists() + def test_reencode_keeps_icc_profile_and_orientation(self, runner, tmp_path): + """Issue #98: erase re-encodes through cv2, which writes no container + metadata, so the output used to come back with no colour profile and no + orientation tag. Both must survive the write and the default strip.""" + import piexif + from PIL import ImageCms + + icc = ImageCms.ImageCmsProfile(ImageCms.createProfile("sRGB")).tobytes() + src = tmp_path / "portrait.png" + Image.new("RGB", (96, 64), "red").save( + src, + icc_profile=icc, + exif=piexif.dump({"0th": {piexif.ImageIFD.Orientation: 6}}), + ) + output = tmp_path / "erased.png" + result = runner.invoke( + main, + ["erase", str(src), "--region", "8,8,16,16", "--backend", "cv2", "-o", str(output)], + ) + assert result.exit_code == 0, result.output + with Image.open(output) as im: + assert im.info.get("icc_profile") == icc + assert im.getexif().get(0x0112) == 6 + def test_erase_two_regions(self, runner, sample_png, tmp_path): output = tmp_path / "erased2.png" result = runner.invoke( diff --git a/tests/test_image_io.py b/tests/test_image_io.py index d39e33b..379d64a 100644 --- a/tests/test_image_io.py +++ b/tests/test_image_io.py @@ -207,3 +207,200 @@ class TestQualityPreservingWrite: def test_unencodable_ext_returns_false_not_raises(self, tmp_path: Path) -> None: # a bogus extension cv2 can't encode returns False (never raises the cv2.error). assert image_io.imwrite(tmp_path / "x.zzz", _make_bgr()) is False + + +class TestDisplayTagCarry: + """A re-encode must not eat the source's ICC profile and EXIF orientation + (issue #98): cv2's encoders write no container metadata, so erase/visible/ + invisible outputs used to come back colourless and, on a build that leaves the + raster stored-orientation, sideways. The pixel paths decode with + IMREAD_UNCHANGED (which never rotates), so the raster keeps the source's stored + orientation and the tag must ride along to the output.""" + + ORIENT = 6 # Rotate 90 CW: a portrait display of a landscape-stored raster. + + @staticmethod + def _icc() -> bytes: + from PIL import ImageCms + + return ImageCms.ImageCmsProfile(ImageCms.createProfile("sRGB")).tobytes() + + @classmethod + def _source(cls, path: Path, *, orient: int | None = 6, icc: bool = True) -> Path: + """A 96x64 stored raster carrying an sRGB profile and Orientation=6.""" + import piexif + from PIL import Image + + kwargs: dict[str, object] = {"icc_profile": cls._icc()} if icc else {} + if orient is not None: + exif = piexif.dump({"0th": {piexif.ImageIFD.Orientation: orient, piexif.ImageIFD.Make: "TestCam"}}) + if path.suffix in (".png", ".webp"): + kwargs["exif"] = exif # Pillow writes the eXIf chunk / WebP EXIF chunk + Image.new("RGB", (96, 64), "red").save(path, **kwargs) + if orient is not None and path.suffix in (".jpg", ".jpeg"): + piexif.insert(piexif.dump({"0th": {piexif.ImageIFD.Orientation: orient}}), str(path)) + return path + + @staticmethod + def _tags(path: Path) -> tuple[object, object]: + from PIL import Image + + with Image.open(path) as im: + return im.info.get("icc_profile"), im.getexif().get(0x0112) + + def _stored_raster(self, path: Path) -> np.ndarray: + # What the pixel paths hold after reading the source: UNCHANGED never rotates. + raster = image_io.imread(path, cv2.IMREAD_UNCHANGED) + assert raster is not None + return raster + + def test_png_output_keeps_profile_and_orientation(self, tmp_path: Path) -> None: + src = self._source(tmp_path / "src.png") + raster = self._stored_raster(src) + out = tmp_path / "out.png" + assert image_io.imwrite(out, raster, display_tags_from=src) is True + icc, orient = self._tags(out) + assert icc == self._icc() + assert orient == self.ORIENT + # The splice must not touch pixels: Pillow decodes without rotating. + from PIL import Image + + with Image.open(out) as im: + assert np.array_equal(np.asarray(im)[:, :, ::-1], raster) + + def test_jpeg_output_carries_both_tags(self, tmp_path: Path) -> None: + src = self._source(tmp_path / "src.png") + raster = self._stored_raster(src) + out = tmp_path / "out.jpg" + assert image_io.imwrite(out, raster, display_tags_from=src) is True + icc, orient = self._tags(out) + assert icc == self._icc() + assert orient == self.ORIENT + + def test_webp_output_carries_both_tags_losslessly(self, tmp_path: Path) -> None: + src = self._source(tmp_path / "src.png") + raster = self._stored_raster(src) + out = tmp_path / "out.webp" + assert image_io.imwrite(out, raster, display_tags_from=src) is True + icc, orient = self._tags(out) + assert icc == self._icc() + assert orient == self.ORIENT + # The lossless Pillow re-save must stay pixel-identical. + back = image_io.imread(out) + assert back is not None + assert np.array_equal(back, raster) + + def test_webp_source_orientation_travels(self, tmp_path: Path) -> None: + # cv2 never rotates a WebP decode, so even a fresh PNG output must be + # tagged, or the portrait source would display sideways. + src = self._source(tmp_path / "src.webp") + raster = self._stored_raster(src) + out = tmp_path / "out.png" + assert image_io.imwrite(out, raster, display_tags_from=src) is True + assert self._tags(out)[1] == self.ORIENT + + def test_an_upright_raster_is_not_tagged_into_a_second_rotation(self, tmp_path: Path) -> None: + # A raster the decode already turned upright (this cv2 build does that for + # JPEG/PNG under IMREAD_COLOR) must not receive the tag again. Stored 96x64 + # with Orientation 6 displays as 64x96, so the upright raster is 64 wide. + src = self._source(tmp_path / "src.png") + upright = np.zeros((96, 64, 3), np.uint8) # 64 cols x 96 rows: stored dims swapped + out = tmp_path / "out.png" + assert image_io.imwrite(out, upright, display_tags_from=src) is True + icc, orient = self._tags(out) + assert orient is None + assert icc == self._icc() + + def test_a_raster_of_unexplained_shape_is_not_tagged(self, tmp_path: Path) -> None: + # A resized pipeline output matches neither the stored nor the upright + # dimensions; carrying the tag there would be a guess, so it is dropped. + src = self._source(tmp_path / "src.png") + resized = np.zeros((31, 17, 3), np.uint8) + out = tmp_path / "out.png" + assert image_io.imwrite(out, resized, display_tags_from=src) is True + assert self._tags(out)[1] is None + + def test_an_untagged_source_changes_nothing(self, tmp_path: Path) -> None: + # No tags to carry: the output must be byte-identical to a plain write. + src = self._source(tmp_path / "src.png", orient=None, icc=False) + raster = self._stored_raster(src) + tagged, plain = tmp_path / "tagged.png", tmp_path / "plain.png" + assert image_io.imwrite(tagged, raster, display_tags_from=src) is True + assert image_io.imwrite(plain, raster) is True + assert tagged.read_bytes() == plain.read_bytes() + + def test_a_missing_decode_source_still_writes(self, tmp_path: Path) -> None: + raster = np.zeros((8, 8, 3), np.uint8) + out = tmp_path / "out.png" + assert image_io.imwrite(out, raster, display_tags_from=tmp_path / "gone.png") is True + assert out.exists() + assert self._tags(out) == (None, None) + + def test_in_place_rewrite_keeps_the_tags(self, tmp_path: Path) -> None: + # Tags are read BEFORE the write, so rewriting a file over itself cannot + # lose what it is about to restore. + path = self._source(tmp_path / "in.png") + raster = self._stored_raster(path) + assert image_io.imwrite(path, raster, display_tags_from=path) is True + icc, orient = self._tags(path) + assert icc == self._icc() + assert orient == self.ORIENT + + def test_png_splice_replaces_rather_than_stacks(self, tmp_path: Path) -> None: + # Rewriting an already-tagged PNG must not duplicate the chunks. + from remove_ai_watermarks.image_io import _png_splice_display_tags + + base = tmp_path / "base.png" + assert image_io.imwrite(base, np.zeros((8, 8, 3), np.uint8)) is True + once = _png_splice_display_tags(base.read_bytes(), self._icc(), self.ORIENT) + twice = _png_splice_display_tags(once, self._icc(), self.ORIENT) + + kinds: list[bytes] = [] + i = 8 + while i + 8 <= len(twice): + length = int.from_bytes(twice[i : i + 4], "big") + kinds.append(twice[i + 4 : i + 8]) + i += 12 + length + assert kinds.count(b"iCCP") == 1 + assert kinds.count(b"eXIf") == 1 + + @pytest.mark.skipif(not _heif_writable("HEIF"), reason="no HEIC encoder in this env") + def test_heif_output_carries_the_profile_and_displays_upright(self, tmp_path: Path) -> None: + # Pillow's HEIF writer bakes the orientation into the pixels on save (it + # transposes the raster and writes Orientation 1), so the display result, + # not the raw tag, is what to assert: upright portrait, exact profile. + src = self._source(tmp_path / "src.png") + raster = self._stored_raster(src) + out = tmp_path / "out.heic" + assert image_io.imwrite(out, raster, display_tags_from=src) is True + from PIL import Image + + with Image.open(out) as im: + assert im.info.get("icc_profile") == self._icc() + assert im.size == (64, 96) # the upright portrait, not the stored landscape + assert im.getexif().get(0x0112) in (None, 1) + + +class TestRemoverWriteCarriesDisplayTags: + """The diffusion path's write must thread ``display_tags_from`` through to + ``imwrite``: the invisible pipeline decodes via Pillow (never rotated), saves a + temp PNG that carries the ICC profile but no orientation (``exif_transpose`` + consumed it upstream), and ``_write_output`` restores from that input path.""" + + def test_the_output_mirrors_its_input_path(self, tmp_path: Path) -> None: + from PIL import Image, ImageCms + + from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover + + icc = ImageCms.ImageCmsProfile(ImageCms.createProfile("sRGB")).tobytes() + # Exactly what the invisible engine's temp save produces: ICC, no orientation. + temp = tmp_path / "temp.png" + Image.new("RGB", (32, 16), "blue").save(temp, icc_profile=icc) + + remover = WatermarkRemover.__new__(WatermarkRemover) # _write_output uses no state + out = tmp_path / "out.png" + remover._write_output(Image.new("RGB", (32, 16), "blue"), out, display_tags_from=temp) + + with Image.open(out) as im: + assert im.info.get("icc_profile") == icc + assert im.getexif().get(0x0112) is None diff --git a/tests/test_metadata.py b/tests/test_metadata.py index e269cc1..ba20168 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -2099,6 +2099,54 @@ class TestSixteenBitPngStrip: assert b"caBX" not in out.read_bytes() + @staticmethod + def _exif_payload(fields: dict[int, object]) -> bytes: + """A bare-TIFF ``eXIf`` chunk payload (what the chunk format carries).""" + dumped = piexif.dump({"0th": fields}) + return dumped[6:] if dumped.startswith(b"Exif\x00\x00") else dumped + + def test_the_orientation_survives_the_chunk_walk(self, tmp_path: Path): + # eXIf used to be dropped wholesale, taking the display orientation with it + # (issue #98); a scrubbed payload survives so a portrait stays a portrait. + src, out = tmp_path / "src.png", tmp_path / "out.png" + self._write_png16( + src, + [(b"eXIf", self._exif_payload({piexif.ImageIFD.Orientation: 6, piexif.ImageIFD.Make: "TestCam"}))], + ) + + remove_ai_metadata(src, out) + + with Image.open(out) as im: + assert im.getexif().get(0x0112) == 6 + assert im.getexif().get(0x010F) == "TestCam" + + def test_the_walk_scrubs_an_ai_token_from_the_kept_exif(self, tmp_path: Path): + # The kept eXIf is scrubbed, not trusted: a generator name goes, the + # orientation stays. + src, out = tmp_path / "src.png", tmp_path / "out.png" + self._write_png16( + src, + [(b"eXIf", self._exif_payload({piexif.ImageIFD.Orientation: 6, piexif.ImageIFD.Make: "Midjourney"}))], + ) + + remove_ai_metadata(src, out) + + with Image.open(out) as im: + assert im.getexif().get(0x0112) == 6 + assert im.getexif().get(0x010F) is None + + def test_remove_all_drops_the_orientation(self, tmp_path: Path): + src, out = tmp_path / "src.png", tmp_path / "out.png" + self._write_png16( + src, + [(b"eXIf", self._exif_payload({piexif.ImageIFD.Orientation: 6}))], + ) + + remove_ai_metadata(src, out, keep_standard=False) + + with Image.open(out) as im: + assert im.getexif().get(0x0112) is None + def test_an_eight_bit_png_keeps_the_shipped_path(self, tmp_path: Path): # The chunk walk is scoped to the depth Pillow cannot hold; an ordinary PNG # must keep going through PIL, which rewrites the file. @@ -2111,3 +2159,61 @@ class TestSixteenBitPngStrip: assert b"parameters" not in out.read_bytes() assert Image.open(out).size == (4, 4) + + +class TestDisplayTagsSurviveTheStrip: + """The default strip keeps the display-fidelity tags (ICC profile, EXIF + orientation) that the re-encode path now carries (issue #98); --remove-all + drops them. Before this, PNG and WebP outputs lost the orientation with the + rest of EXIF, and Pillow's saver fallback kept the ICC profile alive even + under --remove-all, contradicting the documented "removed by --remove-all".""" + + ORIENT = 6 + + @staticmethod + def _icc() -> bytes: + from PIL import ImageCms + + return ImageCms.ImageCmsProfile(ImageCms.createProfile("sRGB")).tobytes() + + @classmethod + def _png(cls, path: Path, *, orient: int = 6, make: str = "TestCam") -> Path: + exif = {"0th": {piexif.ImageIFD.Orientation: orient, piexif.ImageIFD.Make: make}} + Image.new("RGB", (96, 64), "red").save(path, icc_profile=cls._icc(), exif=piexif.dump(exif)) + return path + + def test_png_default_strip_keeps_icc_and_orientation(self, tmp_path: Path): + src, out = tmp_path / "s.png", tmp_path / "o.png" + self._png(src) + remove_ai_metadata(src, out) + with Image.open(out) as im: + assert im.info.get("icc_profile") == self._icc() + assert im.getexif().get(0x0112) == self.ORIENT + assert im.getexif().get(0x010F) == "TestCam" # Make: scrubbed of AI only + + def test_png_remove_all_drops_icc_and_orientation(self, tmp_path: Path): + src, out = tmp_path / "s.png", tmp_path / "o.png" + self._png(src) + remove_ai_metadata(src, out, keep_standard=False) + with Image.open(out) as im: + assert "icc_profile" not in im.info + assert im.getexif().get(0x0112) is None + + def test_png_strip_scrubs_an_ai_token_but_keeps_the_orientation(self, tmp_path: Path): + # The kept EXIF is scrubbed, not trusted: a generator name goes, the + # orientation stays. + src, out = tmp_path / "s.png", tmp_path / "o.png" + self._png(src, make="Midjourney") + remove_ai_metadata(src, out) + with Image.open(out) as im: + assert im.getexif().get(0x0112) == self.ORIENT + assert im.getexif().get(0x010F) is None + + def test_webp_default_strip_keeps_icc_and_orientation(self, tmp_path: Path): + src, out = tmp_path / "s.webp", tmp_path / "o.webp" + exif = piexif.dump({"0th": {piexif.ImageIFD.Orientation: self.ORIENT, piexif.ImageIFD.Make: "TestCam"}}) + Image.new("RGB", (96, 64), "red").save(src, icc_profile=self._icc(), exif=exif, lossless=True) + remove_ai_metadata(src, out) + with Image.open(out) as im: + assert im.info.get("icc_profile") == self._icc() + assert im.getexif().get(0x0112) == self.ORIENT diff --git a/typings/piexif/__init__.pyi b/typings/piexif/__init__.pyi index 54a30c2..3237ef6 100644 --- a/typings/piexif/__init__.pyi +++ b/typings/piexif/__init__.pyi @@ -13,6 +13,7 @@ class ImageIFD: Make: int Artist: int ImageDescription: int + Orientation: int class ExifIFD: UserComment: int