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.

Most operators in this module take decoded ``Image`` values and return numeric
values, strings, or booleans - they do **not** modify the image. Lightweight
metadata and filter operators may also inspect encoded bytes when documented.

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.
"""

from pyflink.common import Row
from pyflink.common.image import Image, ImageMode, validate_image
from pyflink.dataframe import udf, DataType
from pyflink.table.udf import ScalarFunction
from pyflink.multimodal.codec import (
    image_array_to_pil,
    image_mode_channels,
    image_to_ndarray,
    is_decodable_image_bytes,
    read_encoded_image_metadata,
)
from pyflink.multimodal.utils import _build_or_apply_udf, _udf_runtime_kwargs
from pyflink.model.cache_manager import check_dependencies

__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")
_OBSERVATION_ERRORS = (TypeError, ValueError)


# 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):
    metadata = read_encoded_image_metadata(image_bytes)
    if metadata is None:
        return None
    width, height, channels, mode, image_format = metadata
    return Row(
        width=width,
        height=height,
        channels=channels,
        mode=mode,
        format=image_format,
    )


def _metadata_from_observation_input(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)
    raise TypeError(f"Expected Image or encoded bytes, got {type(image_input).__name__}")


def _metadata_or_default(image_input, default):
    # Metadata and filter operators are observation gates: corrupt user data
    # should produce the documented default failure value here. Pixel transforms,
    # model operators, and analysis operators still use decoded Image inputs and
    # fail fast when the input is malformed.
    try:
        return _metadata_from_observation_input(image_input)
    except _OBSERVATION_ERRORS:
        return default


def _dimensions_or_default(image_input, default):
    metadata = _metadata_or_default(image_input, None)
    if metadata is None:
        return default
    return metadata.width, metadata.height


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 - lightweight image validity predicate
class _ImageValidityFilter(ScalarFunction):
    """
    Check if an image input passes lightweight validity checks.

    Returns ``True`` for encoded image bytes that Pillow can verify without
    materializing pixels, or decoded ``Image`` values that pass the built-in
    image validation. The encoded-bytes path is header/container oriented and
    does not guarantee that the full pixel stream can later be materialized by
    ``image_decode``. Returns ``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)
            mode = mode.canonical_name if isinstance(mode, ImageMode) else mode.upper()
        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
        if isinstance(image_input, (bytes, bytearray)):
            return is_decodable_image_bytes(
                image_input, mode=self.mode, max_pixels=self.pixel_limit
            )
        if _is_native_image_input(image_input):
            try:
                validate_image(image_input)
                if (
                    self.pixel_limit is not None
                    and image_input.width * image_input.height > self.pixel_limit
                ):
                    return False
                if self.mode is not None:
                    if image_input.mode.canonical_name == self.mode:
                        return True
                    image_to_ndarray(image_input, mode=self.mode)
                return True
            except (TypeError, ValueError):
                return False
        return False


[docs]def is_valid_image(*columns, mode=None, pixel_limit=None, concurrency=None): """ Check if an image passes lightweight validity checks. For encoded bytes this uses header/container verification and does not guarantee that the full pixel stream can be materialized. Use ``image_decode(on_error="null")`` when corrupt bytes should become null instead of failing. 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) from an image. Decoded images store pixels only, not the original container format, so ``format`` is returned as ``"UNKNOWN"`` for decoded input. Encoded bytes are inspected without materializing decoded pixels and corrupt bytes return ``None``. Encoded-byte metadata is lightweight and does not guarantee that the full pixel stream can be materialized. Usage:: df.with_column("meta", image_metadata()(col("img"))) """ def open(self, function_context): pass def eval(self, image_input): if image_input is None: return None return _metadata_or_default(image_input, None)
[docs]def image_metadata(*columns, concurrency=None): """ Extract metadata (width, height, channels, mode) from an image. Accepts decoded Image values and encoded image bytes. For decoded input, ``mode`` reports the canonical Image mode and ``format`` is ``"UNKNOWN"``. For encoded bytes, metadata is read from the container header; corrupt bytes return ``None``. Encoded-byte metadata is lightweight and does not guarantee that the full pixel stream can be materialized. 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. Accepts decoded Image values and encoded image bytes. Invalid or corrupt input returns ``None``. Usage:: df.with_column("ratio", image_aspect_ratio()(col("img"))) """ def open(self, function_context): pass def eval(self, image_input): if image_input is None: return None dimensions = _dimensions_or_default(image_input, None) if dimensions is None: return None w, h = dimensions return _aspect_ratio_from_dimensions(w, h)
[docs]def image_aspect_ratio(*columns, concurrency=None): """ Compute the aspect ratio (width / height) of an image. Accepts decoded Image values and encoded image bytes. Invalid or corrupt input returns ``None``. 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): if image_input is None: return None pixel_array = image_to_ndarray(image_input, mode="L") 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. Input must be a decoded image; encoded bytes raise ``TypeError``. Decode raw bytes first with ``image_decode(on_error="null")`` when invalid inputs should become null. 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): if image_input is None: return None pixel_array = image_to_ndarray(image_input) 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. pil_image = image_array_to_pil(pixel_array, mode="RGB") 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. Input must be a decoded image; encoded bytes raise ``TypeError``. Decode raw bytes first with ``image_decode(on_error="null")`` when invalid inputs should become null. 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. Encoded bytes are inspected without materializing decoded pixels and the header check does not guarantee full pixel materialization. Invalid input returns ``False``. 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): if image_input is None: return False dimensions = _dimensions_or_default(image_input, None) if dimensions is None: return False w, h = dimensions 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 decoded Image values and encoded image bytes. Encoded bytes are checked from lightweight metadata only and may still fail later decode. Invalid or corrupt input evaluates to ``False``. 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``. Encoded bytes are inspected without materializing decoded pixels and the header check does not guarantee full pixel materialization. Invalid input returns ``False``. 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): if image_input is None: return False dimensions = _dimensions_or_default(image_input, None) if dimensions is None: return False w, h = dimensions ratio = _aspect_ratio_from_dimensions( w, h ) 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). Accepts decoded Image values and encoded image bytes. Encoded bytes are checked from lightweight metadata only and may still fail later decode. Invalid or corrupt input evaluates to ``False``. 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. Non-bytes input returns ``False`` because decoded images do not retain the original encoded byte size. This filter must run before ``image_decode`` because decoded ``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)): # File size is an encoded-container property. If this filter is # accidentally placed after image_decode(), treat it as a failed # predicate instead of failing the whole pipeline. return False 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 is meaningful only for encoded image bytes and should be placed before ``image_decode``; decoded Image input evaluates to ``False``. 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)