Add 16-bit PNG support to Python package

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

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.  

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"

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',

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')
+

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,
+)

python/trustmark/high_bit_depth.py

@@ -0,0 +1,197 @@
+# 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
+
+    # 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"))
+
+
+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).tolist())
+    return buffer.getvalue()