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

@@ -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

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:
 

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

src/remove_ai_watermarks/api.py

@@ -152,7 +152,7 @@ def _write_visible_result(
         # and surfaced as a confusing "cannot read image <INPUT>" 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:

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:

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)

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:

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)

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(

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

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

typings/piexif/__init__.pyi

@@ -13,6 +13,7 @@ class ImageIFD:
     Make: int
     Artist: int
     ImageDescription: int
+    Orientation: int
 
 class ExifIFD:
     UserComment: int