################################################################################
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################
"""Native IMAGE encode/decode conversion helpers for multimodal operators."""
import io
import logging
from dataclasses import dataclass
import numpy as np
from pyflink.common.image import Image, ImageMode, validate_image
from pyflink.model.dependencies import dependency_install_requirement
# Save the original PIL.Image.open before any third-party monkey-patches
# (e.g. ultralytics patches Image.open to add HEIF support, which raises
# ModuleNotFoundError on corrupt input).
# Graceful fallback: module remains importable without Pillow; functions
# that need PIL will fail at call time with a clear error.
try:
from PIL import Image as _PILImage
_pil_open = _PILImage.open
except ImportError:
_PILImage = None
_pil_open = None
# Pillow warns or rejects very large images to protect users from decompression
# bombs. Keep an explicit upper bound for decode safety while allowing common
# high-resolution images; this is about 1.6 GB for RGB uint8 payloads.
_MAX_IMAGE_PIXELS = 178_956_970
if _PILImage is not None:
_PILImage.MAX_IMAGE_PIXELS = _MAX_IMAGE_PIXELS
_logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class _ImageModeSpec:
dtype: np.dtype
channels: int
_IMAGE_MODE_SPECS = {
"L": _ImageModeSpec(np.dtype(np.uint8), 1),
"LA": _ImageModeSpec(np.dtype(np.uint8), 2),
"RGB": _ImageModeSpec(np.dtype(np.uint8), 3),
"RGBA": _ImageModeSpec(np.dtype(np.uint8), 4),
"L16": _ImageModeSpec(np.dtype(np.uint16), 1),
}
_IMAGE_MODE_CHANNELS = {
mode: spec.channels for mode, spec in _IMAGE_MODE_SPECS.items()
}
_IMAGE_MODE_DTYPES = {
mode: spec.dtype for mode, spec in _IMAGE_MODE_SPECS.items()
}
_NATIVE_IMAGE_MODES = {
"L": ImageMode.L,
"LA": ImageMode.LA,
"RGB": ImageMode.RGB,
"RGBA": ImageMode.RGBA,
"L16": ImageMode.L16,
}
_PIL_NATIVE_MODES = {"L", "LA", "RGB", "RGBA"}
_PIL_16BIT_MODES = {"I;16", "I;16L", "I;16B"}
_UINT8_MODES_BY_CHANNELS = {1: "L", 2: "LA", 3: "RGB", 4: "RGBA"}
_UINT16_MODES_BY_CHANNELS = {1: "L16"}
DEFAULT_IMAGE_ENCODE_QUALITY = 85
SUPPORTED_IMAGE_FORMATS = {"JPEG", "PNG", "TIFF", "WEBP", "BMP", "GIF"}
def _normalize_image_mode(mode):
if not isinstance(mode, str):
raise ValueError(f"mode must be a string, got {mode!r}")
normalized = mode.upper()
if normalized not in _IMAGE_MODE_SPECS:
raise ValueError(
f"Unsupported decoded image mode {mode!r}; expected one of "
f"{sorted(_IMAGE_MODE_SPECS)}"
)
return normalized
def image_mode_dtype(mode):
"""Return the numpy dtype used by a native IMAGE mode."""
return _IMAGE_MODE_DTYPES[_normalize_image_mode(mode)]
[docs]def image_mode_channels(mode):
"""Return the channel count used by a native IMAGE mode."""
return _IMAGE_MODE_CHANNELS[_normalize_image_mode(mode)]
def _image_mode_spec(mode):
return _IMAGE_MODE_SPECS[_normalize_image_mode(mode)]
def _infer_image_mode(pixel_array):
dtype = np.dtype(pixel_array.dtype)
if pixel_array.ndim == 2:
channels = 1
elif pixel_array.ndim == 3:
channels = pixel_array.shape[2]
else:
raise ValueError(
"Expected a 2-D or 3-D image array, "
f"got {pixel_array.ndim}-D"
)
if dtype == np.dtype(np.uint8):
mode = _UINT8_MODES_BY_CHANNELS.get(channels)
elif dtype == np.dtype(np.uint16):
mode = _UINT16_MODES_BY_CHANNELS.get(channels)
else:
mode = None
if mode is None:
raise ValueError(
"Unsupported native IMAGE dtype/channel combination: "
f"dtype={pixel_array.dtype}, channels={channels}. "
"Supported native modes are L, LA, RGB, RGBA, and L16."
)
return mode, channels
def _ensure_channel_axis(pixel_array):
if pixel_array.ndim == 2:
return pixel_array[:, :, None]
return pixel_array
def _normalize_native_image_mode(mode):
if isinstance(mode, ImageMode):
return mode
if isinstance(mode, str):
normalized = _normalize_image_mode(mode)
return _NATIVE_IMAGE_MODES[normalized]
raise TypeError(
"Image mode must be an ImageMode or canonical mode string, got %s"
% type(mode).__name__
)
def _native_mode_name(mode):
if isinstance(mode, ImageMode):
return mode.canonical_name
return _normalize_image_mode(mode)
def _ensure_native_image_array(pixel_array):
pixel_array = _ensure_channel_axis(np.asarray(pixel_array))
if not pixel_array.flags["C_CONTIGUOUS"]:
pixel_array = np.ascontiguousarray(pixel_array)
return pixel_array
[docs]def image_to_ndarray(image, mode=None):
"""Convert a native :class:`Image` to an ndarray for operator internals."""
if image is None:
return None
validate_image(image)
pixel_array = image.data
if mode is not None:
target_mode = _native_mode_name(mode)
if image.mode.canonical_name != target_mode:
pixel_array = convert_image_array_to_mode(pixel_array, target_mode)
if pixel_array.ndim == 3 and pixel_array.shape[2] == 1:
return pixel_array[:, :, 0]
return pixel_array
[docs]def ndarray_to_image(pixel_array, mode=None):
"""Convert an ndarray to a native :class:`Image`."""
if pixel_array is None:
return None
if mode is not None:
target_mode = _native_mode_name(mode)
pixel_array = convert_image_array_to_mode(pixel_array, target_mode)
native_mode = _normalize_native_image_mode(target_mode)
else:
inferred_mode, _ = _infer_image_mode(pixel_array)
native_mode = _normalize_native_image_mode(inferred_mode)
normalized = _ensure_native_image_array(pixel_array)
return Image(data=normalized, mode=native_mode)
[docs]def pil_to_image(pil_image, mode=None):
"""Convert a PIL image to a native :class:`Image`."""
if _PILImage is None:
raise ImportError(
"Pillow is required for image operations. "
f"Install it with: pip install {dependency_install_requirement('PIL')}"
)
normalized_mode = _native_mode_name(mode) if mode is not None else None
pixel_array = _canonicalize_decoded_pil_image(pil_image, normalized_mode)
return ndarray_to_image(pixel_array)
def _convert_dtype(pixel_array, target_dtype):
target_dtype = np.dtype(target_dtype)
source_dtype = np.dtype(pixel_array.dtype)
if source_dtype == target_dtype:
return pixel_array
if target_dtype == np.dtype(np.uint8):
if source_dtype == np.dtype(np.uint16):
return (pixel_array / 257).clip(0, 255).astype(np.uint8)
if np.issubdtype(source_dtype, np.integer):
return pixel_array.clip(0, 255).astype(np.uint8)
if source_dtype == np.dtype(np.float32):
max_value = np.nanmax(pixel_array) if pixel_array.size else 1.0
scale = 255.0 if max_value <= 1.0 else 1.0
return (pixel_array * scale).clip(0, 255).astype(np.uint8)
if target_dtype == np.dtype(np.uint16):
if source_dtype == np.dtype(np.uint8):
return pixel_array.astype(np.uint16) * 257
if np.issubdtype(source_dtype, np.integer):
return pixel_array.clip(0, 65535).astype(np.uint16)
if source_dtype == np.dtype(np.float32):
max_value = np.nanmax(pixel_array) if pixel_array.size else 1.0
scale = 65535.0 if max_value <= 1.0 else 1.0
return (pixel_array * scale).clip(0, 65535).astype(np.uint16)
if target_dtype == np.dtype(np.float32):
if source_dtype == np.dtype(np.uint8):
return pixel_array.astype(np.float32) / 255.0
if source_dtype == np.dtype(np.uint16):
return pixel_array.astype(np.float32) / 65535.0
if np.issubdtype(source_dtype, np.integer):
return pixel_array.astype(np.float32)
raise ValueError(f"Unsupported dtype conversion: {source_dtype} -> {target_dtype}")
def _convert_channels(pixel_array, target_channels):
if target_channels not in (1, 2, 3, 4):
raise ValueError(f"Unsupported target channel count: {target_channels}")
if pixel_array.ndim not in (2, 3):
raise ValueError(
f"Expected a 2-D or 3-D image array, got {pixel_array.ndim}-D"
)
src = _ensure_channel_axis(pixel_array)
source_channels = src.shape[2]
if source_channels == target_channels:
return pixel_array if target_channels == 1 else src
max_alpha = np.array(np.iinfo(src.dtype).max, dtype=src.dtype) \
if np.issubdtype(src.dtype, np.integer) else np.array(1.0, dtype=src.dtype)
if source_channels == 1:
luminance = src[:, :, 0]
alpha = np.full(luminance.shape, max_alpha, dtype=src.dtype)
elif source_channels == 2:
luminance = src[:, :, 0]
alpha = src[:, :, 1]
elif source_channels in (3, 4):
rgb = src[:, :, :3].astype(np.float32)
luminance = (
0.299 * rgb[:, :, 0]
+ 0.587 * rgb[:, :, 1]
+ 0.114 * rgb[:, :, 2]
).astype(src.dtype)
alpha = src[:, :, 3] if source_channels == 4 else np.full(
luminance.shape, max_alpha, dtype=src.dtype
)
else:
raise ValueError(f"Unsupported source channel count: {source_channels}")
if target_channels == 1:
return luminance
if target_channels == 2:
return np.stack([luminance, alpha], axis=2)
if source_channels == 1:
rgb = np.repeat(src, 3, axis=2)
elif source_channels == 2:
rgb = np.repeat(src[:, :, :1], 3, axis=2)
else:
rgb = src[:, :, :3]
if target_channels == 3:
return rgb
return np.concatenate([rgb, alpha[:, :, None]], axis=2)
[docs]def convert_image_array_to_mode(pixel_array, mode):
"""Convert an image ndarray to a canonical native IMAGE mode."""
spec = _image_mode_spec(mode)
converted = _convert_channels(pixel_array, spec.channels)
converted = _convert_dtype(converted, spec.dtype)
if spec.channels == 1 and converted.ndim == 3:
converted = converted[:, :, 0]
return converted
def _canonicalize_decoded_pil_image(pil_image, mode):
if pil_image.mode == "P":
pil_image = pil_image.convert(
"RGBA" if pil_image.info.get("transparency") is not None else "RGB"
)
if mode is not None:
normalized_mode = _normalize_image_mode(mode)
pixel_array = np.asarray(pil_image)
if pixel_array.dtype.byteorder in (">", "<"):
pixel_array = pixel_array.astype(
pixel_array.dtype.newbyteorder("="), copy=False
)
return convert_image_array_to_mode(pixel_array, normalized_mode)
if pil_image.mode in _PIL_NATIVE_MODES:
return np.asarray(pil_image)
if pil_image.mode in _PIL_16BIT_MODES:
return np.asarray(pil_image).astype(np.uint16, copy=False)
pixel_array = np.asarray(pil_image)
if pixel_array.dtype.byteorder in (">", "<"):
pixel_array = pixel_array.astype(pixel_array.dtype.newbyteorder("="), copy=False)
_infer_image_mode(pixel_array)
return pixel_array
[docs]def image_array_to_pil(pixel_array, mode=None):
"""Convert ndarray pixels to a PIL Image, converting mode only when requested."""
if _PILImage is None:
raise ImportError(
"Pillow is required for image operations. "
f"Install it with: pip install {dependency_install_requirement('PIL')}"
)
if mode is not None:
pixel_array = convert_image_array_to_mode(pixel_array, mode)
inferred_mode, _ = _infer_image_mode(pixel_array)
if inferred_mode == "L16":
return _PILImage.fromarray(pixel_array)
if inferred_mode in _PIL_NATIVE_MODES:
return _PILImage.fromarray(pixel_array)
raise ValueError(
f"Pillow cannot represent native IMAGE mode {inferred_mode!r} without "
"an explicit mode conversion"
)
def _pil_fallback_mode_for_array(pixel_array):
channels = 1 if pixel_array.ndim == 2 else pixel_array.shape[2]
return {1: "L", 2: "LA", 3: "RGB", 4: "RGBA"}[channels]
[docs]def image_array_to_pil_compatible(pixel_array):
"""Convert ndarray pixels to a PIL Image, normalizing only when required.
Pillow cannot operate directly on high-precision multichannel arrays.
This helper preserves natively supported modes and falls back to the
uint8 mode with the same channel count for PIL-only transforms.
"""
try:
return image_array_to_pil(pixel_array)
except ValueError:
return image_array_to_pil(
pixel_array, mode=_pil_fallback_mode_for_array(pixel_array)
)
[docs]def validate_image_quality(quality):
"""Validate JPEG/WebP encode quality."""
if not isinstance(quality, int) or quality < 1 or quality > 100:
raise ValueError(f"quality must be an integer in [1, 100], got {quality!r}")
return quality
# Image decode helpers
# Data-level exceptions from PIL / cv2 that indicate corrupt input.
# Kept narrow to avoid masking real bugs (e.g. ImportError from a missing
# module). PIL raises OSError for truncated/missing files, ValueError for
# invalid headers, and SyntaxError (via PngImagePlugin.FormatError) for
# corrupt PNG chunks.
_PIL_IMAGE_ERRORS = (
(_PILImage.DecompressionBombError,) if _PILImage is not None else ()
)
_IMAGE_DATA_ERRORS = (OSError, ValueError, SyntaxError) + _PIL_IMAGE_ERRORS
def _read_image_dimensions(image_bytes):
"""Read image dimensions from file header without decoding.
Parses JPEG (SOF0/SOF2), PNG (IHDR), WebP (VP8/VP8L), BMP, and GIF
headers to extract width and height. Returns ``(width, height)`` on
success, or ``(None, None)`` if the format is unrecognized or the
header is too short to parse.
This is a zero-allocation, zero-decode operation - safe to call on
every image before PIL decoding as a memory-bomb guard.
"""
if image_bytes is None or len(image_bytes) < 2:
return None, None
# JPEG: scan for SOF0 (0xFFC0) or SOF2 (0xFFC2) marker
if image_bytes[:2] == b'\xff\xd8':
# Scan markers; each marker is 0xFF + non-zero type byte.
# SOF markers contain height (2B) and width (2B) at offset 5-8
# within the marker payload.
i = 2
while i + 9 <= len(image_bytes):
if image_bytes[i] != 0xFF:
break
marker = image_bytes[i + 1]
# Skip padding 0xFF bytes
if marker == 0xFF:
i += 1
continue
# Standalone markers (no payload): RST0-RST7, SOI, EOI, TEM
if marker == 0xD8 or marker == 0xD9 or marker == 0x01:
i += 2
continue
if 0xD0 <= marker <= 0xD7:
i += 2
continue
# SOS marker - image data follows, stop scanning
if marker == 0xDA:
break
# Read marker payload length (big-endian uint16, includes itself)
seg_len = (image_bytes[i + 2] << 8) | image_bytes[i + 3]
# SOF0, SOF1, SOF2 (progressive): height at offset +5, width at +7
if marker in (0xC0, 0xC1, 0xC2):
h = (image_bytes[i + 5] << 8) | image_bytes[i + 6]
w = (image_bytes[i + 7] << 8) | image_bytes[i + 8]
return w, h
i += 2 + seg_len
return None, None
# PNG: IHDR chunk at offset 16
if len(image_bytes) >= 24 and image_bytes[:4] == b'\x89PNG':
w = int.from_bytes(image_bytes[16:20], 'big')
h = int.from_bytes(image_bytes[20:24], 'big')
return w, h
# WebP
if len(image_bytes) >= 30 and image_bytes[:4] == b'RIFF' and image_bytes[8:12] == b'WEBP':
# Lossy (VP8): bits 0-3 of header are a signature, then
# width (14 bits) at byte 26, height (14 bits) at byte 28
if image_bytes[12:16] == b'VP8 ':
w = int.from_bytes(image_bytes[26:28], 'little') & 0x3FFF
h = int.from_bytes(image_bytes[28:30], 'little') & 0x3FFF
return w, h
# Lossless (VP8L): byte 20 is signature 0x2F, then
# width-1 (14 bits) and height-1 (14 bits) packed in 4 bytes
if image_bytes[12:16] == b'VP8L' and len(image_bytes) >= 25:
if image_bytes[20] == 0x2F:
bits = int.from_bytes(image_bytes[21:25], 'little')
w = (bits & 0x3FFF) + 1
h = ((bits >> 14) & 0x3FFF) + 1
return w, h
# Extended WebP (VP8X): canvas width/height minus 1 are stored
# as 24-bit little-endian integers at bytes 24..29.
if image_bytes[12:16] == b'VP8X':
w = int.from_bytes(image_bytes[24:27], 'little') + 1
h = int.from_bytes(image_bytes[27:30], 'little') + 1
return w, h
return None, None
# BMP: width at offset 18, height at offset 22 (LE int32)
if len(image_bytes) >= 26 and image_bytes[:2] == b'BM':
w = int.from_bytes(image_bytes[18:22], 'little')
h = abs(int.from_bytes(image_bytes[22:26], 'little'))
return w, h
# GIF: width at offset 6, height at offset 8 (LE uint16)
if len(image_bytes) >= 10 and image_bytes[:4] == b'GIF8':
w = int.from_bytes(image_bytes[6:8], 'little')
h = int.from_bytes(image_bytes[8:10], 'little')
return w, h
return None, None
def _convert_for_jpeg(pil_image):
if pil_image.mode not in ("L", "RGB"):
# JPEG only supports 8-bit L/RGB in Pillow; convert explicitly at the
# output boundary instead of letting save() fail with a vague error.
return pil_image.convert("RGB")
return pil_image
[docs]def image_array_to_jpeg_pil(pixel_array):
"""Convert ndarray pixels to a PIL Image compatible with JPEG output."""
try:
pil_image = image_array_to_pil(pixel_array)
except ValueError:
# JPEG cannot preserve high-precision multichannel pixels; normalize
# explicitly at the output boundary.
pil_image = image_array_to_pil(pixel_array, mode="RGB")
return _convert_for_jpeg(pil_image)
[docs]def decode_image(image_input, mode=None, max_pixels=None):
"""Decode encoded bytes or native Image input into a native :class:`Image`."""
pixel_array = decode_image_input(image_input, mode=mode, max_pixels=max_pixels)
if pixel_array is None:
return None
target_mode = _native_mode_name(mode) if mode is not None else None
return ndarray_to_image(pixel_array, mode=target_mode)
[docs]def safe_decode_batch(series, mode=None, max_pixels=None):
"""Decode a batch of bytes/native IMAGE inputs into ndarrays, skipping failures."""
valid_images = []
valid_indices = []
for i, item in enumerate(series):
pixel_array = safe_decode_image_input(item, mode=mode, max_pixels=max_pixels)
if pixel_array is not None:
valid_images.append(pixel_array)
valid_indices.append(i)
return valid_images, valid_indices
[docs]def decode_image_batch(series, mode=None, max_pixels=None):
"""Decode a batch of bytes/native IMAGE inputs into ndarrays (fail-fast)."""
images = []
valid_idx = []
for i, item in enumerate(series):
if item is None:
continue
pixel_array = decode_image_input(item, mode=mode, max_pixels=max_pixels)
if pixel_array is None:
continue
images.append(pixel_array)
valid_idx.append(i)
return images, valid_idx