Source code for pyflink.multimodal.operators.image_info

################################################################################
#  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.
################################################################################
"""
Image information extraction operators.

Operators in this module take raw encoded image bytes or decoded ``Image``
values as input and return numeric values, strings, or booleans - they do
**not** modify the image.

Without business params (factory function, no args)::

    from pyflink.multimodal.operators.image_info import image_aspect_ratio
    df.with_column("ratio", image_aspect_ratio()(col("img")))

With business params (factory function)::

    from pyflink.multimodal.operators.image_info import image_hash
    phash = image_hash(method="phash")
    df.with_column("hash", phash(col("img")))

Runtime args:
    ``concurrency`` is forwarded to the DataFrame UDF runtime.
"""

import io

from pyflink.common import Row
from pyflink.common.image import Image, validate_image
from pyflink.dataframe import udf, DataType
from pyflink.table.udf import ScalarFunction
from pyflink.multimodal.codec import (
    _pil_open,
    _read_image_dimensions,
    decode_image_input,
    detect_image_format,
    image_array_to_pil,
    image_mode_channels,
    safe_decode_image_input,
)
from pyflink.multimodal.utils import _build_or_apply_udf, _udf_runtime_kwargs
from pyflink.model.cache_manager import check_dependencies
from pyflink.model.dependencies import dependency_install_requirement

__all__ = [
    # Metadata extraction
    "image_metadata",
    "image_aspect_ratio",
    "image_sharpness",
    "image_hash",
    # Filters
    "is_valid_image",
    "image_size_filter",
    "image_shape_filter",
    "image_file_size_filter",
]

_SUPPORTED_HASH_METHODS = ("phash", "dhash", "ahash")


# Shared metadata and dimension helpers
def _is_native_image_input(image_input):
    return isinstance(image_input, Image)


def _metadata_from_image(image_input):
    validate_image(image_input)
    return Row(
        width=image_input.width,
        height=image_input.height,
        channels=image_input.channels,
        mode=image_input.mode.canonical_name,
        format="UNKNOWN",
    )


def _metadata_from_encoded_bytes(image_bytes):
    if not isinstance(image_bytes, (bytes, bytearray)):
        return None
    if _pil_open is None:
        raise ImportError(
            "Pillow is required for image operations. "
            f"Install it with: pip install {dependency_install_requirement('PIL')}"
        )
    # Use the original Pillow opener captured by codec to avoid global
    # Image.open monkey patches from optional model libraries.
    with _pil_open(io.BytesIO(image_bytes)) as pil_image:
        width, height = pil_image.size
        image_format = pil_image.format or detect_image_format(image_bytes)
        # getbands() reads Pillow's mode metadata and does not materialize
        # the full pixel array, keeping image_metadata header-only for raw
        # encoded bytes.
        channels = len(pil_image.getbands())
        # Pillow's mode string (e.g. "L", "RGB", "RGBA", "P", "I;16") describes
        # the encoded image as stored; it is reported as-is for raw bytes.
        pil_mode = pil_image.mode
    return Row(
        width=width,
        height=height,
        channels=channels,
        mode=pil_mode,
        format=image_format or "UNKNOWN",
    )


def _dimensions_from_header_or_image(image_input):
    if _is_native_image_input(image_input):
        metadata = _metadata_from_image(image_input)
        return metadata.width, metadata.height
    if isinstance(image_input, (bytes, bytearray)):
        w, h = _read_image_dimensions(image_input)
        if w is not None and h is not None:
            return w, h
    return None, None


def _aspect_ratio_from_dimensions(w, h):
    if h is None or h <= 0 or w is None:
        return None
    return float(w) / float(h)


def _within_optional_bounds(value, lower_bound, upper_bound):
    if lower_bound is not None and value < lower_bound:
        return False
    return upper_bound is None or value <= upper_bound


def _validate_at_least_one_bound(operator_name, *bounds):
    if all(bound is None for bound in bounds):
        raise ValueError(f"{operator_name} requires at least one bound.")


def _laplacian_variance(np, pixel_array):
    if pixel_array.shape[0] < 3 or pixel_array.shape[1] < 3:
        return 0.0
    pixel_array = pixel_array.astype(np.float32)
    laplacian = (
        pixel_array[:-2, 1:-1]
        + pixel_array[2:, 1:-1]
        + pixel_array[1:-1, :-2]
        + pixel_array[1:-1, 2:]
        - 4 * pixel_array[1:-1, 1:-1]
    )
    return float(laplacian.var())


# ImageValidityFilter - filter images that can be decoded
class _ImageValidityFilter(ScalarFunction):
    """
    Check if an image input can be decoded by the multimodal image codec.

    Returns ``True`` for raw image bytes or native ``Image`` values that can
    be consumed by downstream image operators, and ``False`` for null,
    corrupt, or oversized image data.
    """

    def __init__(self, mode=None, pixel_limit=None):
        super().__init__()
        if mode is not None:
            image_mode_channels(mode)
        if pixel_limit is not None and pixel_limit <= 0:
            raise ValueError(f"pixel_limit must be positive, got {pixel_limit!r}")
        self.mode = mode
        self.pixel_limit = pixel_limit

    def open(self, function_context):
        pass

    def eval(self, image_input):
        if image_input is None:
            return False
        pixel_array = safe_decode_image_input(
            image_input, mode=self.mode, max_pixels=self.pixel_limit
        )
        return pixel_array is not None


[docs]def is_valid_image(*columns, mode=None, pixel_limit=None, concurrency=None): """ Check if an image can be decoded by the multimodal image codec. The returned predicate is intended for ``DataFrame.filter`` before ``image_decode`` in fetch/decode pipelines. Args: *columns: Optional image columns. When provided, the UDF is applied directly (syntactic sugar). mode: Optional target mode to validate against, e.g. ``"RGB"``. pixel_limit: Maximum allowed ``width * height``. Images exceeding this limit are treated as invalid. concurrency: UDF concurrency. ``None`` uses the framework default. Returns: A UDF producing ``boolean``. Example:: >>> # As a reusable variable >>> valid = is_valid_image(pixel_limit=4096 * 4096) >>> df = df.filter(valid(col("content"))) >>> >>> # Inline >>> df = df.filter(is_valid_image(col("content"), pixel_limit=4096 * 4096)) """ wrapper = udf( _ImageValidityFilter(mode=mode, pixel_limit=pixel_limit), return_dtype=DataType.boolean(), **_udf_runtime_kwargs(concurrency=concurrency), ) return _build_or_apply_udf(wrapper, *columns)
# ImageMetadata - extract width, height, channels, mode, format class _ImageMetadata(ScalarFunction): """ Extract metadata (width, height, channels, mode, encoded format) from an image. Raw encoded bytes are inspected through Pillow headers without converting the full image into an ndarray; ``mode`` is reported as Pillow's mode string (e.g. ``"L"``, ``"RGB"``, ``"P"``, ``"I;16"``). For decoded ``Image`` input, width/height/channels/mode are read from the value - ``mode`` is the canonical mode (``"L"``, ``"LA"``, ``"RGB"``, ``"RGBA"``, ``"L16"``) - and ``format`` is returned as ``"UNKNOWN"`` because decoded images store pixels only, not the original container format. Usage:: df.with_column("meta", image_metadata()(col("img"))) """ def open(self, function_context): pass def eval(self, image_input): if isinstance(image_input, (bytes, bytearray)): return _metadata_from_encoded_bytes(image_input) if _is_native_image_input(image_input): return _metadata_from_image(image_input) return None
[docs]def image_metadata(*columns, concurrency=None): """ Extract metadata (width, height, channels, mode, encoded format) from an image. ``mode`` reports the image's channel/precision mode: Pillow's mode string for raw encoded bytes, or the canonical mode (``"L"``, ``"LA"``, ``"RGB"``, ``"RGBA"``, ``"L16"``) for decoded ``Image`` input. ``format`` describes the encoded container for raw bytes. For decoded ``Image`` input it is ``"UNKNOWN"`` because the internal pixel value intentionally does not retain original encoding metadata. Apply format filters to raw encoded bytes before ``image_decode`` when the original container format matters. Args: *columns: Optional image columns. When provided, the UDF is applied directly (syntactic sugar). concurrency: UDF concurrency. ``None`` uses the framework default. Returns: A UDF producing an object with ``width``, ``height``, ``channels``, ``mode``, and ``format`` fields. Example:: >>> # As a reusable variable >>> metadata = image_metadata() >>> df = df.with_columns("meta", metadata(col("img"))) >>> >>> # Inline >>> df = df.with_columns("meta", image_metadata(col("img"))) """ wrapper = udf( _ImageMetadata(), return_dtype=DataType.struct({ "width": DataType.int32(), "height": DataType.int32(), "channels": DataType.int32(), "mode": DataType.string(), "format": DataType.string(), }), **_udf_runtime_kwargs(concurrency=concurrency), ) return _build_or_apply_udf(wrapper, *columns)
# ImageAspectRatio - width / height ratio class _ImageAspectRatio(ScalarFunction): """ Compute the aspect ratio (width / height) of an image. Usage:: df.with_column("ratio", image_aspect_ratio()(col("img"))) """ def open(self, function_context): pass def eval(self, image_input): w, h = _dimensions_from_header_or_image(image_input) if w is not None and h is not None: return _aspect_ratio_from_dimensions(w, h) pixel_array = decode_image_input(image_input) if pixel_array is None: return None return _aspect_ratio_from_dimensions( pixel_array.shape[1], pixel_array.shape[0] )
[docs]def image_aspect_ratio(*columns, concurrency=None): """ Compute the aspect ratio (width / height) of an image. Args: *columns: Optional image columns. When provided, the UDF is applied directly (syntactic sugar). concurrency: UDF concurrency. ``None`` uses the framework default. Returns: A UDF producing ``float64``. Example:: >>> # As a reusable variable >>> aspect_ratio = image_aspect_ratio() >>> df = df.with_columns("ratio", aspect_ratio(col("img"))) >>> >>> # Inline >>> df = df.with_columns("ratio", image_aspect_ratio(col("img"))) """ wrapper = udf( _ImageAspectRatio(), return_dtype=DataType.float64(), **_udf_runtime_kwargs(concurrency=concurrency), ) return _build_or_apply_udf(wrapper, *columns)
# ImageSharpness - low-cost Laplacian variance focus measure class _ImageSharpness(ScalarFunction): """ Compute a raw sharpness metric using Laplacian variance. The returned value is a focus measure in ``[0, +inf)``: larger values usually indicate sharper edges and stronger high-frequency detail. It is intentionally not normalized because absolute thresholds depend on image scale, texture, contrast, and noise. Use ``max_edge`` or an upstream resize step when comparing images with a fixed threshold. Usage:: df.with_column("sharpness", image_sharpness()(col("img"))) """ def __init__(self, max_edge=None, allow_upscale=False): super().__init__() if max_edge is not None and max_edge <= 0: raise ValueError(f"max_edge must be positive or None, got {max_edge}") self.max_edge = max_edge self.allow_upscale = allow_upscale def open(self, function_context): check_dependencies("numpy", "PIL") import numpy as np from PIL import Image self._np = np self._resample_lanczos = getattr(Image, "Resampling", Image).LANCZOS def _resize_to_max_edge(self, pixel_array): if self.max_edge is None: return pixel_array height, width = pixel_array.shape[:2] longest_edge = max(height, width) if longest_edge <= 0: return pixel_array ratio = float(self.max_edge) / float(longest_edge) if ratio == 1.0 or (ratio > 1.0 and not self.allow_upscale): return pixel_array new_size = ( max(1, int(round(width * ratio))), max(1, int(round(height * ratio))), ) pil_image = image_array_to_pil(pixel_array) resized = pil_image.resize(new_size, self._resample_lanczos) return self._np.asarray(resized) def eval(self, image_input): pixel_array = decode_image_input(image_input, mode="L") if pixel_array is None: return None pixel_array = self._resize_to_max_edge(pixel_array) return _laplacian_variance(self._np, pixel_array)
[docs]def image_sharpness( *columns, max_edge=None, allow_upscale=False, concurrency=None, ): """ Compute a raw sharpness metric using Laplacian variance. The result is in ``[0, +inf)`` and larger values usually indicate sharper images. This is a low-cost focus measure, not an absolute quality score: thresholds depend on image scale, texture, contrast, and noise. Args: *columns: Optional image columns. When provided, the UDF is applied directly (syntactic sugar). max_edge: Cap the longest edge to this many pixels before computing sharpness. ``None`` (default) preserves the original scale. allow_upscale: When ``max_edge`` is set, allow upscaling smaller images. Default ``False``. concurrency: UDF concurrency. ``None`` uses the framework default. Returns: A UDF producing ``float64``. Example:: >>> # As a reusable variable >>> sharpness = image_sharpness(max_edge=512) >>> df = df.with_columns("sharpness", sharpness(col("img"))) >>> >>> # Inline >>> df = df.with_columns("sharpness", image_sharpness(col("img"), max_edge=512)) """ wrapper = udf( _ImageSharpness( max_edge=max_edge, allow_upscale=allow_upscale, ), return_dtype=DataType.float64(), **_udf_runtime_kwargs(concurrency=concurrency), ) return _build_or_apply_udf(wrapper, *columns)
# ImageHash - perceptual / difference / average hash class _ImageHash(ScalarFunction): """ Compute a perceptual hash of an image. Args: method: Hash algorithm - "phash", "dhash", or "ahash". Default "phash". Requires: ``pip install imagehash`` """ def __init__(self, method="phash"): super().__init__() if method not in _SUPPORTED_HASH_METHODS: raise ValueError( f"Unknown hash method '{method}'. " f"Supported: {list(_SUPPORTED_HASH_METHODS)}" ) self.method = method def open(self, function_context): check_dependencies("imagehash") import imagehash self._hash_funcs = { "phash": imagehash.phash, "dhash": imagehash.dhash, "ahash": imagehash.average_hash, } def eval(self, image_input): pixel_array = decode_image_input(image_input) if pixel_array is None: return None try: pil_image = image_array_to_pil(pixel_array) except ValueError: # Perceptual hash libraries operate on 8-bit PIL images; high # precision decoded image input must be normalized here. try: pil_image = image_array_to_pil(pixel_array, mode="RGB") except (TypeError, ValueError): return None return str(self._hash_funcs[self.method](pil_image))
[docs]def image_hash(*columns, method="phash", concurrency=None): """ Compute a perceptual hash of an image. Args: *columns: Optional image columns. When provided, the UDF is applied directly (syntactic sugar). method: Hash algorithm — ``"phash"`` (default), ``"dhash"``, or ``"ahash"``. concurrency: UDF concurrency. ``None`` uses the framework default. Returns: A UDF producing ``string``. Example:: >>> # As a reusable variable >>> hasher = image_hash(method="phash") >>> df = df.with_columns("hash", hasher(col("img"))) >>> >>> # Inline >>> df = df.with_columns("hash", image_hash(col("img"), method="phash")) """ wrapper = udf( _ImageHash(method=method), return_dtype=DataType.string(), **_udf_runtime_kwargs(concurrency=concurrency), ) return _build_or_apply_udf(wrapper, *columns)
# ImageSizeFilter - filter by pixel dimensions class _ImageSizeFilter(ScalarFunction): """ Check if an image's dimensions fall within the specified bounds. Returns True if the image dimensions are within the configured lower and optional upper bounds. ``None`` upper bounds are treated as unbounded. Args: min_w: Minimum width. Default None. min_h: Minimum height. Default None. max_w: Maximum width. Default None. max_h: Maximum height. Default None. """ def __init__(self, min_w=None, min_h=None, max_w=None, max_h=None): super().__init__() _validate_at_least_one_bound( "image_size_filter", min_w, min_h, max_w, max_h ) self.min_w = min_w self.min_h = min_h self.max_w = max_w self.max_h = max_h def open(self, function_context): pass def eval(self, image_input): w, h = _dimensions_from_header_or_image(image_input) if w is not None and h is not None: return ( _within_optional_bounds(w, self.min_w, self.max_w) and _within_optional_bounds(h, self.min_h, self.max_h) ) pixel_array = decode_image_input(image_input) if pixel_array is None: return False h, w = pixel_array.shape[0], pixel_array.shape[1] return ( _within_optional_bounds(w, self.min_w, self.max_w) and _within_optional_bounds(h, self.min_h, self.max_h) )
[docs]def image_size_filter(*columns, min_w=None, min_h=None, max_w=None, max_h=None, concurrency=None): """ Filter images by pixel dimensions. Accepts both raw encoded image bytes and decoded ``Image`` values. For raw bytes, dimensions are read from image headers when possible, so normal JPEG/PNG inputs do not need a full decode. At least one dimension bound must be configured. Args: *columns: Optional image columns. When provided, the UDF is applied directly (syntactic sugar). min_w: Minimum width in pixels. Default ``None`` (unbounded). min_h: Minimum height in pixels. Default ``None`` (unbounded). max_w: Maximum width in pixels. Default ``None`` (unbounded). max_h: Maximum height in pixels. Default ``None`` (unbounded). concurrency: UDF concurrency. ``None`` uses the framework default. Returns: A UDF producing ``boolean``. Example:: >>> # As a reusable variable >>> size_filter = image_size_filter(min_w=100, max_w=4096) >>> df = df.filter(size_filter(col("img"))) >>> >>> # Inline >>> df = df.filter(image_size_filter(col("img"), min_w=100, max_w=4096)) """ wrapper = udf( _ImageSizeFilter(min_w=min_w, min_h=min_h, max_w=max_w, max_h=max_h), return_dtype=DataType.boolean(), **_udf_runtime_kwargs(concurrency=concurrency), ) return _build_or_apply_udf(wrapper, *columns)
# ImageShapeFilter - filter by aspect ratio range class _ImageShapeFilter(ScalarFunction): """ Check if an image's aspect ratio (width/height) falls within a range. Returns True if ``width / height`` is above ``min_ratio`` and below the optional ``max_ratio``. A ``None`` upper bound is treated as unbounded. Args: min_ratio: Minimum aspect ratio. Default None. max_ratio: Maximum aspect ratio. Default None. """ def __init__(self, min_ratio=None, max_ratio=None): super().__init__() _validate_at_least_one_bound( "image_shape_filter", min_ratio, max_ratio ) self.min_ratio = min_ratio self.max_ratio = max_ratio def open(self, function_context): pass def eval(self, image_input): w, h = _dimensions_from_header_or_image(image_input) if w is not None and h is not None: ratio = _aspect_ratio_from_dimensions(w, h) if ratio is None: return False return _within_optional_bounds( ratio, self.min_ratio, self.max_ratio ) pixel_array = decode_image_input(image_input) if pixel_array is None: return False ratio = _aspect_ratio_from_dimensions( pixel_array.shape[1], pixel_array.shape[0] ) if ratio is None: return False return _within_optional_bounds(ratio, self.min_ratio, self.max_ratio)
[docs]def image_shape_filter(*columns, min_ratio=None, max_ratio=None, concurrency=None): """ Filter images by aspect ratio (width / height). At least one aspect-ratio bound must be configured. Args: *columns: Optional image columns. When provided, the UDF is applied directly (syntactic sugar). min_ratio: Minimum aspect ratio. Default ``None`` (unbounded). max_ratio: Maximum aspect ratio. Default ``None`` (unbounded). concurrency: UDF concurrency. ``None`` uses the framework default. Returns: A UDF producing ``boolean``. Example:: >>> # As a reusable variable >>> shape_filter = image_shape_filter(min_ratio=0.5, max_ratio=2.0) >>> df = df.filter(shape_filter(col("img"))) >>> >>> # Inline >>> df = df.filter(image_shape_filter(col("img"), min_ratio=0.5, max_ratio=2.0)) """ wrapper = udf( _ImageShapeFilter(min_ratio=min_ratio, max_ratio=max_ratio), return_dtype=DataType.boolean(), **_udf_runtime_kwargs(concurrency=concurrency), ) return _build_or_apply_udf(wrapper, *columns)
# ImageFileSizeFilter - filter by file size in bytes class _ImageFileSizeFilter(ScalarFunction): """ Check if raw encoded image bytes fall within the specified byte range. Returns True if the byte length is within the configured lower and optional upper bounds. ``None`` bounds are treated as unbounded. This filter must run before ``image_decode`` because native ``Image`` values no longer contain the original encoded byte size. Args: min_bytes: Minimum file size in bytes. Default None. max_bytes: Maximum file size in bytes. Default None. """ def __init__(self, min_bytes=None, max_bytes=None): super().__init__() _validate_at_least_one_bound( "image_file_size_filter", min_bytes, max_bytes ) self.min_bytes = min_bytes self.max_bytes = max_bytes def open(self, function_context): pass def eval(self, image_bytes): if image_bytes is None: return False if not isinstance(image_bytes, (bytes, bytearray)): raise ValueError( "image_file_size_filter expects raw encoded image bytes; " "place it before image_decode()." ) size = len(image_bytes) return _within_optional_bounds(size, self.min_bytes, self.max_bytes)
[docs]def image_file_size_filter(*columns, min_bytes=None, max_bytes=None, concurrency=None): """ Filter images by encoded file size in bytes. This operator filters by ``len(image_bytes)`` of the original encoded payload. It only accepts raw encoded image bytes and must be placed before ``image_decode``. At least one byte-size bound must be configured. Args: *columns: Optional image columns. When provided, the UDF is applied directly (syntactic sugar). min_bytes: Minimum file size in bytes. Default ``None`` (unbounded). max_bytes: Maximum file size in bytes. Default ``None`` (unbounded). concurrency: UDF concurrency. ``None`` uses the framework default. Returns: A UDF producing ``boolean``. Example:: >>> # As a reusable variable >>> file_size_filter = image_file_size_filter(min_bytes=1024, max_bytes=5*1024*1024) >>> df = df.filter(file_size_filter(col("img"))) >>> >>> # Inline >>> df = df.filter( ... image_file_size_filter(col("img"), min_bytes=1024, max_bytes=5*1024*1024) ... ) """ wrapper = udf( _ImageFileSizeFilter(min_bytes=min_bytes, max_bytes=max_bytes), return_dtype=DataType.boolean(), **_udf_runtime_kwargs(concurrency=concurrency), ) return _build_or_apply_udf(wrapper, *columns)