diff --git a/src/remove_ai_watermarks/image_io.py b/src/remove_ai_watermarks/image_io.py index 4eab43a..954958f 100644 --- a/src/remove_ai_watermarks/image_io.py +++ b/src/remove_ai_watermarks/image_io.py @@ -29,6 +29,7 @@ from remove_ai_watermarks._internal.constants import PNG_SIGNATURE if TYPE_CHECKING: from numpy.typing import NDArray + from PIL.Image import Image as PILImage log = logging.getLogger(__name__) @@ -72,7 +73,11 @@ def _pil_read(path: str | Path, flags: int) -> NDArray[Any] | None: """Decode via Pillow (HEIC/AVIF and any other Pillow-readable container) into the cv2 layout ``flags`` implies: grayscale, 3-channel BGR, or BGRA when the source has alpha and ``IMREAD_UNCHANGED`` was requested. Returns None if Pillow (with the - optional HEIF plugin) still cannot open it. No EXIF auto-rotation, matching cv2.""" + optional HEIF plugin) still cannot open it. + + We ask for no EXIF auto-rotation, but Pillow's TIFF reader applies it anyway and + cannot be told not to, so a TIFF raster comes back upright. That is handled where + it matters, in :func:`_read_display_tags`.""" import cv2 import numpy as np @@ -151,6 +156,8 @@ def _register_heif() -> None: # AI-provenance tags the metadata strip exists to drop. _ORIENT_EXIF_TAG = 0x0112 # EXIF ImageIFD.Orientation +_TIFF_IMAGE_WIDTH = 256 # TIFF ImageWidth, the stored raster's columns +_TIFF_IMAGE_LENGTH = 257 # TIFF ImageLength, the stored raster's rows # 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 @@ -175,6 +182,23 @@ def _orientation_exif(orient: int | None, *, prefixed: bool) -> bytes | None: return dumped[6:] if dumped.startswith(b"Exif\x00\x00") else dumped +def _stored_size(im: PILImage) -> tuple[int, int]: + """The size the container STORES, which is not always ``Image.size``. + + Pillow's TIFF reader applies the EXIF orientation itself: it hands back a turned + raster and reports the upright size. Comparing that against a turned raster says + "still in stored orientation" for pixels that are already upright, and the tag + rides along into a second rotation. The TIFF IFD keeps the real stored geometry in + ImageWidth/ImageLength, and no other Pillow plugin exposes ``tag_v2``. + """ + tags = getattr(im, "tag_v2", None) + if tags is not None: + width, height = tags.get(_TIFF_IMAGE_WIDTH), tags.get(_TIFF_IMAGE_LENGTH) + if isinstance(width, int) and isinstance(height, int): + return width, height + return im.size + + 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``. @@ -191,6 +215,9 @@ def _read_display_tags(source: Path, raster_shape: tuple[int, int]) -> tuple[byt 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. + + The stored size cannot be read straight off ``Image``: Pillow's TIFF reader + reports the upright one, so it comes through :func:`_stored_size`. """ _register_heif() try: @@ -201,7 +228,7 @@ def _read_display_tags(source: Path, raster_shape: tuple[int, int]) -> tuple[byt 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 + stored_w, stored_h = _stored_size(im) except Exception: return None, None if not (isinstance(icc, bytes) and icc): diff --git a/tests/test_image_io.py b/tests/test_image_io.py index 3d8457f..1a07d0f 100644 --- a/tests/test_image_io.py +++ b/tests/test_image_io.py @@ -410,3 +410,39 @@ class TestRemoverWriteCarriesDisplayTags: with Image.open(out) as im: assert im.info.get("icc_profile") == icc assert im.getexif().get(0x0112) is None + + +class TestTiffOrientation: + """Pillow's TIFF reader applies the EXIF orientation itself, and reports the + upright size while doing it. A stored-size comparison against ``Image.size`` + therefore matches for a raster that is already upright, and the tag is carried + onto it -- a viewer then turns the picture a second time.""" + + ORIENT = 6 # Rotate 90 CW. + + @staticmethod + def _tagged(path: Path, orient: int) -> Path: + """A 96x64 stored raster carrying ``orient``, written by the container itself.""" + import piexif + from PIL import Image + + exif = piexif.dump({"0th": {piexif.ImageIFD.Orientation: orient}}) + Image.new("RGB", (96, 64), "red").save(path, exif=exif) + return path + + @staticmethod + def _out_orientation(path: Path) -> object: + from PIL import Image + + with Image.open(path) as im: + return im.getexif().get(0x0112) + + def test_a_turned_tiff_decode_is_not_tagged_into_a_second_rotation(self, tmp_path: Path) -> None: + src = self._tagged(tmp_path / "src.tif", self.ORIENT) + raster = image_io.imread(src, cv2.IMREAD_UNCHANGED) + assert raster is not None + # Pillow already applied the rotation: the raster is upright, 64x96. + assert (raster.shape[1], raster.shape[0]) == (64, 96) + out = tmp_path / "out.png" + assert image_io.imwrite(out, raster, display_tags_from=src) is True + assert self._out_orientation(out) is None