From 753a7605463f99e6b410d0b673aeedbaccf15fda Mon Sep 17 00:00:00 2001 From: yoraiyanivbria Date: Sun, 30 Aug 2026 10:55:39 +0300 Subject: [PATCH 1/2] Add 16-bit PNG support to Python package PIL.Image.open silently flattens 16-bit-per-channel PNGs to 8-bit on load, so any high-bit-depth source already loses precision before TrustMark.encode() ever runs. TrustMark.encode_high_bit_depth() reads the true 16-bit pixels directly (OpenImageIO for RGB, pypng for RGBA, avoiding OpenImageIO's PNG writer premultiplying alpha on write), runs the real encoder on a throwaway 8-bit copy, and adds back only the watermark's own perturbation onto the untouched full-precision pixels before re-encoding as 16-bit. Both codecs are gated behind the new optional `trustmark[highbitdepth]` extra and lazily imported, so the base package gains no new hard dependency. Decode is intentionally out of scope: PIL's implicit 8-bit flatten on read doesn't materially affect watermark detection, only encode-side precision. --- python/CLAUDE.md | 45 +++++++ python/README.md | 8 ++ python/pyproject.toml | 8 +- python/setup.py | 3 + python/test.py | 67 +++++++++++ python/trustmark/__init__.py | 6 + python/trustmark/high_bit_depth.py | 186 +++++++++++++++++++++++++++++ python/trustmark/trustmark.py | 52 ++++++++ 8 files changed, 374 insertions(+), 1 deletion(-) create mode 100644 python/trustmark/high_bit_depth.py diff --git a/python/CLAUDE.md b/python/CLAUDE.md index 4f62f9a..16297f3 100644 --- a/python/CLAUDE.md +++ b/python/CLAUDE.md @@ -132,6 +132,51 @@ tm.encode( # returns: PIL image (RGB), same resolution as input ``` +## 16-bit PNG support + +`PIL.Image.open` silently flattens a 16-bit-per-channel PNG to 8-bit on load — no +error, no signal anything was lost. Since no TrustMark model is trained on more than +8-bit input anyway, `tm.encode()` alone can never deliver output that's genuinely +16-bit: the source has already lost its extra precision by the time it reaches PIL. + +`encode_high_bit_depth` solves this by reading the true 16-bit pixels directly (never +through PIL), running the real encoder on a throwaway 8-bit copy, and adding back only +the watermark's own perturbation onto the untouched full-precision pixels. It operates +on raw PNG bytes in and out — unlike every other method on this class, it does **not** +take or return a PIL image, since PIL can't represent the data being preserved. + +Requires the optional extra: + +```bash +pip install trustmark[highbitdepth] +``` + +```python +with open('input_16bit.png', 'rb') as f: + raw_png_bytes = f.read() + +watermarked_bytes = tm.encode_high_bit_depth(raw_png_bytes, secret, MODE='binary') + +with open('output_16bit.png', 'wb') as f: + f.write(watermarked_bytes) +``` + +- Supports plain RGB and RGBA 16-bit-per-channel PNGs only. Anything else (8-bit, + grayscale, palette, non-PNG) raises `ValueError` — use `tm.encode()` instead for + ordinary 8-bit sources. +- Raises `ImportError` with an install hint if `raw_png_bytes` is genuinely high-bit-depth + but the required codec (`OpenImageIO` for RGB, `pypng` for RGBA) isn't installed. +- **There is no `decode_high_bit_depth`.** Decode watermarked 16-bit output with the + ordinary `tm.decode()`, on the image as loaded (and silently flattened to 8-bit) by + PIL — precision loss on read doesn't materially affect watermark *detection*, only + the delivered pixel precision on encode, so no separate high-bit-depth decode path + is needed: + +```python +stego = Image.open('output_16bit.png').convert('RGB') # flattened to 8-bit by PIL, that's fine +secret_out, wm_present, wm_schema = tm.decode(stego, MODE='binary') +``` + ## Watermark removal Requires `loadRemover=True` (the default) at construction time. Calling diff --git a/python/README.md b/python/README.md index 9e51f05..2c3ac3c 100644 --- a/python/README.md +++ b/python/README.md @@ -100,6 +100,14 @@ im_recover = tm.remove_watermark(stego) im_recover.save('images/recovered.png') ``` +### 16-bit PNG support + +Watermarking a genuine 16-bit-per-channel PNG through `tm.encode()` loses that extra +precision, since PIL flattens 16-bit PNGs to 8-bit on load. Use +`tm.encode_high_bit_depth()` instead to keep full precision in the delivered image — +see the "16-bit PNG support" section of [`CLAUDE.md`](CLAUDE.md) for details and the +required `pip install trustmark[highbitdepth]` extra. + ## GPU setup TrustMark runs well on CPU hardware. diff --git a/python/pyproject.toml b/python/pyproject.toml index 2868c3e..6e2a450 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -30,7 +30,7 @@ classifiers = [ "Programming Language :: Python :: 3" ] -dependencies = [ +dependencies = [ "omegaconf>=2.1", "numpy>=1.20.0,<2.0.0", "torch>=2.1.2", @@ -40,6 +40,12 @@ dependencies = [ "einops>=0.4.0" ] +[project.optional-dependencies] +highbitdepth = [ + "OpenImageIO>=2.4", + "pypng>=0.20220715.0", +] + [build-system] requires = ["setuptools>=43.0.0", "wheel"] build-backend = "setuptools.build_meta" diff --git a/python/setup.py b/python/setup.py index 9300419..f964cb0 100644 --- a/python/setup.py +++ b/python/setup.py @@ -31,6 +31,9 @@ 'six>=1.9', 'einops>=0.4.0' ], + extras_require={ + 'highbitdepth': ['OpenImageIO>=2.4', 'pypng>=0.20220715.0'], + }, classifiers=[ 'Development Status :: 5 - Production/Stable', diff --git a/python/test.py b/python/test.py index e2e1358..dd823ab 100644 --- a/python/test.py +++ b/python/test.py @@ -7,8 +7,16 @@ from trustmark import TrustMark +from trustmark.high_bit_depth import ( + read_high_bit_depth_rgb, + write_16bit_rgb_png, + read_high_bit_depth_rgba, + write_16bit_rgba_png, + _png_header, +) from PIL import Image from pathlib import Path +import io import math,random import numpy as np @@ -71,3 +79,62 @@ im_recover.putalpha(alpha) im_recover.save('recovered.png', exif=stego.info.get('exif'), icc_profile=stego.info.get('icc_profile'), dpi=stego.info.get('dpi')) +# 16-bit PNG support (requires: pip install trustmark[highbitdepth]) +# Skipped gracefully if OpenImageIO/pypng aren't installed. +try: + import OpenImageIO # noqa: F401 + import png # noqa: F401 + HAVE_HIGHBITDEPTH = True +except ImportError: + HAVE_HIGHBITDEPTH = False + print('Skipping 16-bit PNG tests: pip install trustmark[highbitdepth] not installed') + +if HAVE_HIGHBITDEPTH: + rng = np.random.default_rng(1234) + + # --- pure I/O round trip: RGB, no model involved --- + rgb_pixels = rng.integers(0, 65536, (48, 64, 3)).astype(np.float32) / 65535.0 + rgb_bytes = write_16bit_rgb_png(rgb_pixels) + assert _png_header(rgb_bytes)['bit_depth'] == 16 + read_back = read_high_bit_depth_rgb(rgb_bytes) + max_err = np.max(np.abs(read_back.pixels - rgb_pixels)) + print(f'16-bit RGB I/O round trip max error: {max_err:.6f} (expect near uint16 quantization step)') + + # --- pure I/O round trip: RGBA, no model involved --- + rgba_pixels = rng.integers(0, 65536, (48, 64, 4)).astype(np.float32) / 65535.0 + rgba_bytes = write_16bit_rgba_png(rgba_pixels, gamma=0.45455) + assert _png_header(rgba_bytes)['bit_depth'] == 16 + read_back_rgba = read_high_bit_depth_rgba(rgba_bytes) + max_err_rgba = np.max(np.abs(read_back_rgba.pixels - rgba_pixels)) + print(f'16-bit RGBA I/O round trip max error: {max_err_rgba:.6f}, gamma preserved: {read_back_rgba.gamma}') + + # --- end-to-end: watermark a genuine 16-bit RGB PNG, confirm precision + watermark survive --- + # Uses real photo content (upscaled 8-bit -> 16-bit via *257, the standard bit-depth + # expansion) rather than random noise: TrustMark's detector relies on structural image + # content, and pure per-pixel noise is unreliable to decode regardless of image size + # (confirmed separately -- not specific to the high-bit-depth path). + hb_source_8bit = np.asarray(Image.open('../images/ripley.jpg').convert('RGB').resize((256, 256)), dtype=np.uint16) + hb_rgb_pixels = (hb_source_8bit * 257).astype(np.float32) / 65535.0 + hb_in_bytes = write_16bit_rgb_png(hb_rgb_pixels) + hb_secret = ''.join([random.choice(['0', '1']) for _ in range(capacity)]) + hb_out_bytes = tm.encode_high_bit_depth(hb_in_bytes, hb_secret, MODE='binary') + assert _png_header(hb_out_bytes)['bit_depth'] == 16, 'encode_high_bit_depth must deliver a genuinely 16-bit PNG' + + # No decode_high_bit_depth exists -- ordinary decode() on the PIL-flattened 8-bit + # image is expected to still detect the watermark. + hb_stego_8bit = Image.open(io.BytesIO(hb_out_bytes)).convert('RGB') + hb_secret_out, hb_present, hb_schema = tm.decode(hb_stego_8bit, MODE='binary') + if hb_present and hb_secret_out == hb_secret: + print('16-bit encode_high_bit_depth -> 8-bit decode() round trip: secret recovered correctly') + else: + print(f'16-bit round trip FAILED: present={hb_present} secret_out={hb_secret_out!r} expected={hb_secret!r}') + + # --- negative check: ordinary 8-bit PNG bytes must be rejected, not silently degraded --- + buf = io.BytesIO() + Image.fromarray(np.zeros((16, 16, 3), dtype=np.uint8)).save(buf, format='PNG') + try: + tm.encode_high_bit_depth(buf.getvalue(), hb_secret, MODE='binary') + print('encode_high_bit_depth FAILED to reject an 8-bit PNG') + except ValueError: + print('encode_high_bit_depth correctly rejects 8-bit PNG input') + diff --git a/python/trustmark/__init__.py b/python/trustmark/__init__.py index 49f03a2..0863842 100644 --- a/python/trustmark/__init__.py +++ b/python/trustmark/__init__.py @@ -5,3 +5,9 @@ import numpy as np from .trustmark import TrustMark +from .high_bit_depth import ( + read_high_bit_depth_rgb, + write_16bit_rgb_png, + read_high_bit_depth_rgba, + write_16bit_rgba_png, +) diff --git a/python/trustmark/high_bit_depth.py b/python/trustmark/high_bit_depth.py new file mode 100644 index 0000000..12df54e --- /dev/null +++ b/python/trustmark/high_bit_depth.py @@ -0,0 +1,186 @@ +# Copyright 2026 Adobe +# All Rights Reserved. + +# NOTICE: Adobe permits you to use, modify, and distribute this file in +# accordance with the terms of the Adobe license agreement accompanying +# it. + +"""Bypasses PIL for images with real per-channel precision above 8 bits. + +`PIL.Image.open` silently flattens a 16-bit RGB PNG to 8-bit `"RGB"` mode on load -- no +error, no signal anything was lost (Pillow can't even represent a `(H, W, 3)` uint16 +array -- `Image.fromarray` raises `TypeError` on one). TrustMark itself is only ever +trained on 8-bit input regardless, so this module exists purely to read/write pixels at +their native bit depth around the model -- see `TrustMark.encode_high_bit_depth` in +`trustmark.py` for how the watermark itself gets applied without losing that precision. + +`read_high_bit_depth_rgb`/`read_high_bit_depth_rgba` return `None` for anything that +isn't a genuine 16-bit-per-channel PNG of the matching channel count -- callers should +treat that as "not applicable here", not an error. + +Two separate codecs are used, and neither is a hard dependency of the base package -- +both are lazily imported only once a PNG header confirms they're actually needed, so +`import trustmark` never requires either. Install both with `pip install +trustmark[highbitdepth]`. + +- RGB uses OpenImageIO, which can read/write native 16-bit and carry forward full + source color metadata (ICC profile, gamma, etc.) via `copy_metadata`. Its Python + binding has no in-memory I/O path, so reads/writes here go through a temp file. +- RGBA uses `pypng` (pure Python) instead of OpenImageIO: OpenImageIO's PNG writer + unconditionally premultiplies RGB by alpha on write for any 4-channel PNG, with no + attribute that disables it -- PNG itself only ever stores straight (unassociated) + alpha, so that would be non-conformant output. `pypng` does exact, literal + sample-value round trips instead, natively over in-memory bytes, at the cost of only + carrying the PNG `gAMA` chunk forward as color metadata (no ICC-profile pass-through). +""" + +import io +import os +import struct +import tempfile +from typing import NamedTuple, Optional + +import numpy as np + +PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" + +# PNG IHDR color types (see the PNG spec) relevant to this module. +_COLOR_TYPE_RGB = 2 +_COLOR_TYPE_RGBA = 6 + + +class HighBitDepthImage(NamedTuple): + pixels: np.ndarray # normalized (0..1) float32 RGB, shape (H, W, 3) + source: object # opaque oiio.ImageBuf -- never read pixels from this directly, metadata only + + +class HighBitDepthRGBAImage(NamedTuple): + pixels: np.ndarray # normalized (0..1) float32 RGBA, shape (H, W, 4) + gamma: Optional[float] # from the source PNG's gAMA chunk, if any + + +def _png_header(raw_bytes: bytes) -> Optional[dict]: + """Parses just the PNG signature + IHDR chunk (zero dependencies) to cheaply check + bit depth and color type before deciding whether a codec import is even needed. + Returns None if `raw_bytes` doesn't start with the PNG magic number. + """ + if raw_bytes[:8] != PNG_SIGNATURE: + return None + + # IHDR is always the first chunk: 4-byte length, 4-byte type "IHDR", then 13 bytes + # of data (width, height, bit depth, color type, compression, filter, interlace). + if len(raw_bytes) < 8 + 8 + 13 or raw_bytes[12:16] != b"IHDR": + return None + + width, height, bit_depth, color_type = struct.unpack(">IIBB", raw_bytes[16:26]) + return {"width": width, "height": height, "bit_depth": bit_depth, "color_type": color_type} + + +def read_high_bit_depth_rgb(raw_bytes: bytes) -> Optional[HighBitDepthImage]: + """Returns the normalized (0..1) float32 RGB pixels plus source metadata if `raw_bytes` + decodes to a genuine 16-bit-per-channel, alpha-free, 3-channel PNG -- None otherwise. + """ + header = _png_header(raw_bytes) + if header is None or header["bit_depth"] != 16 or header["color_type"] != _COLOR_TYPE_RGB: + return None + + try: + import OpenImageIO as oiio + except ImportError as e: + raise ImportError( + "Reading/writing 16-bit RGB PNGs requires OpenImageIO. Install with: pip install 'trustmark[highbitdepth]'" + ) from e + + fd, tmp_path = tempfile.mkstemp(suffix=".png") + try: + with os.fdopen(fd, "wb") as tmp: + tmp.write(raw_bytes) + + buf = oiio.ImageBuf(tmp_path) + if buf.has_error: + return None + + spec = buf.spec() + if spec.format.basetype != oiio.BASETYPE.UINT16 or spec.nchannels != 3: + return None + + # Forces full pixel+metadata materialization into memory now, while the backing + # temp file still exists -- `source` (used later by write_16bit_rgb_png's + # copy_metadata()) needs to outlive this file, and ImageBuf's read-on-demand + # behavior isn't a documented guarantee to lean on once the temp file is gone. + if not buf.read(force=True): + return None + + pixels = buf.get_pixels(oiio.FLOAT) + # OpenImageIO normalizes integer formats to 0..1 float on read. + return HighBitDepthImage(pixels=np.asarray(pixels, dtype=np.float32), source=buf) + finally: + os.unlink(tmp_path) + + +def write_16bit_rgb_png(pixels: np.ndarray, source: Optional[object] = None) -> bytes: + """Encodes normalized (0..1) float32 RGB pixels, shape (H, W, 3), as a 16-bit-per-channel + PNG. Pass the `source` from `read_high_bit_depth_rgb` to carry its color metadata forward. + """ + import OpenImageIO as oiio + + height, width, _channels = pixels.shape + spec = oiio.ImageSpec(width, height, 3, oiio.UINT16) + buf = oiio.ImageBuf(spec) + if source is not None: + buf.copy_metadata(source) + + pixels_u16 = (np.clip(pixels, 0.0, 1.0) * 65535.0 + 0.5).astype(np.uint16) + buf.set_pixels(oiio.ROI(0, width, 0, height, 0, 1, 0, 3), pixels_u16) + + fd, tmp_path = tempfile.mkstemp(suffix=".png") + try: + os.close(fd) + if not buf.write(tmp_path): + raise RuntimeError(f"Failed to encode 16-bit PNG: {buf.geterror()}") + with open(tmp_path, "rb") as tmp: + return tmp.read() + finally: + os.unlink(tmp_path) + + +def read_high_bit_depth_rgba(raw_bytes: bytes) -> Optional[HighBitDepthRGBAImage]: + """Returns the normalized (0..1) float32 RGBA pixels plus gamma if `raw_bytes` decodes to a + genuine 16-bit-per-channel, 4-channel (RGBA) PNG -- None otherwise. + """ + header = _png_header(raw_bytes) + if header is None or header["bit_depth"] != 16 or header["color_type"] != _COLOR_TYPE_RGBA: + return None + + try: + import png as pypng + except ImportError as e: + raise ImportError( + "Reading/writing 16-bit RGBA PNGs requires pypng. Install with: pip install 'trustmark[highbitdepth]'" + ) from e + + reader = pypng.Reader(bytes=raw_bytes) + width, height, rows, info = reader.asDirect() + raw = np.array(list(rows), dtype=np.uint16) + + pixels = raw.reshape(height, width, 4).astype(np.float32) / 65535.0 + return HighBitDepthRGBAImage(pixels=pixels, gamma=info.get("gamma")) + + +def write_16bit_rgba_png(pixels: np.ndarray, gamma: Optional[float] = None) -> bytes: + """Encodes normalized (0..1) float32 RGBA pixels, shape (H, W, 4), as a 16-bit-per-channel + PNG. Pass the `gamma` from `read_high_bit_depth_rgba` to carry it forward. + """ + import png as pypng + + height, width, _channels = pixels.shape + raw = (np.clip(pixels, 0.0, 1.0) * 65535.0 + 0.5).astype(np.uint16) + + writer_kwargs = {"width": width, "height": height, "bitdepth": 16, "alpha": True, "greyscale": False} + if gamma is not None: + writer_kwargs["gamma"] = gamma + writer = pypng.Writer(**writer_kwargs) + + buffer = io.BytesIO() + writer.write(buffer, raw.reshape(height, width * 4)) + return buffer.getvalue() diff --git a/python/trustmark/trustmark.py b/python/trustmark/trustmark.py index 35ac492..a8dda6b 100644 --- a/python/trustmark/trustmark.py +++ b/python/trustmark/trustmark.py @@ -15,6 +15,7 @@ from omegaconf import OmegaConf from .datalayer import DataLayer +from . import high_bit_depth as hbd from PIL import Image from torchvision import transforms import numpy as np @@ -506,6 +507,57 @@ def encode(self, in_cover_image, string_secret, MODE='text', WM_STRENGTH=1.0, WM return Image.fromarray(stego.astype(np.uint8)) + @torch.no_grad() + def _embed_watermark_residual(self, pixels_rgb01, string_secret, MODE='text', WM_STRENGTH=1.0, WM_MERGE='bilinear'): + # pixels_rgb01: float32 ndarray (H,W,3), normalized 0..1, full precision. + # Returns: float32 ndarray (H,W,3), normalized 0..1 -- the original pixels plus + # the watermark's own (8-bit-quantized) perturbation, not re-quantized to 8-bit + # itself. Runs the real encoder on a throwaway 8-bit copy since no TrustMark + # model is trained on anything higher, then adds back only the encoder's own + # effect (its output minus that same 8-bit copy) onto the untouched full-precision + # pixels passed in. + img8_arr = np.clip(pixels_rgb01 * 255.0 + 0.5, 0, 255).astype(np.uint8) + img8 = Image.fromarray(img8_arr, mode='RGB') + + # WM_STRENGTH passes straight through: encode() already multiplies it by 1.25 + # internally for model_type=='P' -- doing that here too would double-apply it. + stego8 = self.encode(img8, string_secret, MODE=MODE, WM_STRENGTH=WM_STRENGTH, WM_MERGE=WM_MERGE) + + residual01 = (np.asarray(stego8, dtype=np.float32) - img8_arr.astype(np.float32)) / 255.0 + return np.clip(pixels_rgb01 + residual01, 0.0, 1.0).astype(np.float32) + + @torch.no_grad() + def encode_high_bit_depth(self, raw_png_bytes, string_secret, MODE='text', WM_STRENGTH=1.0, WM_MERGE='bilinear'): + # Watermarks a genuine 16-bit-per-channel PNG (RGB or RGBA) without lossy 8-bit + # flattening of the delivered pixels -- PIL.Image.open silently flattens 16-bit + # RGB PNGs to 8-bit on load, so encode() alone can't preserve source precision. + # Requires the optional 'highbitdepth' extra: pip install trustmark[highbitdepth] + # + # Inputs + # raw_png_bytes: bytes of a 16-bit-per-channel PNG file (RGB or RGBA) + # string_secret: same as encode() + # Outputs: bytes of a 16-bit-per-channel PNG file + # + # There is no decode_high_bit_depth counterpart -- decode ordinary 8-bit-flattened + # PIL images with the existing decode(), since precision loss on read doesn't + # materially affect watermark detection, only the delivered pixel precision here. + rgb = hbd.read_high_bit_depth_rgb(raw_png_bytes) + if rgb is not None: + out_pixels = self._embed_watermark_residual(rgb.pixels, string_secret, MODE, WM_STRENGTH, WM_MERGE) + return hbd.write_16bit_rgb_png(out_pixels, source=rgb.source) + + rgba = hbd.read_high_bit_depth_rgba(raw_png_bytes) + if rgba is not None: + rgb_pixels, alpha = rgba.pixels[..., :3], rgba.pixels[..., 3:4] + out_rgb = self._embed_watermark_residual(rgb_pixels, string_secret, MODE, WM_STRENGTH, WM_MERGE) + out_pixels = np.concatenate([out_rgb, alpha], axis=-1) + return hbd.write_16bit_rgba_png(out_pixels, gamma=rgba.gamma) + + raise ValueError( + "raw_png_bytes is not a genuine 16-bit-per-channel RGB or RGBA PNG. " + "Use TrustMark.encode() for ordinary 8-bit sources instead." + ) + @torch.no_grad() def remove_watermark(self, in_cover_image, WM_STRENGTH=1.0, WM_MERGE='bilinear'): From 93ec92b93f124732d79fde09616f00a014479216 Mon Sep 17 00:00:00 2001 From: yoraiyanivbria Date: Sun, 30 Aug 2026 11:21:54 +0300 Subject: [PATCH 2/2] Make the RGBA reader exception-safe for corrupt PNG bodies pypng raises (e.g. ChunkError) on malformed input that still passes the cheap PNG-header sniff, unlike OpenImageIO which reports decode failures via has_error/return codes rather than a Python exception. Wrap the pypng decode so corrupt-but-header-valid bytes return None like every other non-applicable input, instead of letting the exception escape. Also convert the RGBA writer's pixel array to a plain list before handing it to pypng.Writer.write(), for portability across pypng versions. --- python/trustmark/high_bit_depth.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/python/trustmark/high_bit_depth.py b/python/trustmark/high_bit_depth.py index 12df54e..f4dda98 100644 --- a/python/trustmark/high_bit_depth.py +++ b/python/trustmark/high_bit_depth.py @@ -159,9 +159,20 @@ def read_high_bit_depth_rgba(raw_bytes: bytes) -> Optional[HighBitDepthRGBAImage "Reading/writing 16-bit RGBA PNGs requires pypng. Install with: pip install 'trustmark[highbitdepth]'" ) from e - reader = pypng.Reader(bytes=raw_bytes) - width, height, rows, info = reader.asDirect() - raw = np.array(list(rows), dtype=np.uint16) + # Unlike OpenImageIO (which reports decode failures via has_error/return codes, never + # a Python exception), pypng raises on malformed input -- e.g. a truncated or corrupt + # PNG body that still happens to pass the cheap _png_header sniff above. Caller + # contract is "None for anything not applicable here", not "may raise", so any decode + # failure here is equivalent to "not a genuine 16-bit RGBA PNG". + try: + reader = pypng.Reader(bytes=raw_bytes) + width, height, rows, info = reader.asDirect() + raw = np.array(list(rows), dtype=np.uint16) + except Exception: + return None + + if info.get("bitdepth") != 16 or info.get("planes") != 4: + return None pixels = raw.reshape(height, width, 4).astype(np.float32) / 65535.0 return HighBitDepthRGBAImage(pixels=pixels, gamma=info.get("gamma")) @@ -182,5 +193,5 @@ def write_16bit_rgba_png(pixels: np.ndarray, gamma: Optional[float] = None) -> b writer = pypng.Writer(**writer_kwargs) buffer = io.BytesIO() - writer.write(buffer, raw.reshape(height, width * 4)) + writer.write(buffer, raw.reshape(height, width * 4).tolist()) return buffer.getvalue()