################################################################################
# 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.
################################################################################
"""General helper utilities for multimodal UDF development."""
import os
from typing import Callable, TypeVar, overload
import numpy as np
from pyflink.table.expression import Expression
IMAGENET_RGB_MEAN = (0.485, 0.456, 0.406)
IMAGENET_RGB_STD = (0.229, 0.224, 0.225)
# UDF construction helpers
def _is_column_expression(value):
"""Return True when *value* is a PyFlink column/expression object."""
return isinstance(value, Expression)
_WrapperT = TypeVar("_WrapperT")
_ReturnT = TypeVar("_ReturnT")
# Overload 1: no input columns — operator factory returns the wrapper unchanged
# so the caller can apply it later, e.g. ``image_hash("phash")(col("img"))``.
@overload
def _build_or_apply_udf(wrapper: _WrapperT) -> _WrapperT:
...
# Overload 2: input columns supplied — apply the wrapper to them and return
# the resulting Expression. The wrapper's return type is propagated so plain
# test callables that return non-Expression values keep their declared type.
@overload
def _build_or_apply_udf(
wrapper: Callable[..., _ReturnT], *columns: Expression
) -> _ReturnT:
...
def _build_or_apply_udf(wrapper, *columns):
"""Return a UDF wrapper, or apply it when input columns are provided.
Operator constructors consistently build the final PyFlink UDF wrapper
first, then call this helper so no-parameter and parameterized operators
follow the same wrapper construction path.
"""
if not columns:
return wrapper
if not all(_is_column_expression(column) for column in columns):
raise ValueError("UDF input arguments must be column expressions")
return wrapper(*columns)
_UDF_RUNTIME_KEYS = ("concurrency", "batch_size", "num_gpus", "gpu_type")
def _udf_runtime_kwargs(**kwargs):
"""Return non-empty DataFrame UDF runtime options."""
return {key: kwargs[key] for key in _UDF_RUNTIME_KEYS if kwargs.get(key) is not None}
def _setup_cv2_threads(cv2_module):
num_threads = os.environ.get("OMP_NUM_THREADS")
if num_threads is None:
return
try:
num_threads = int(num_threads)
except (TypeError, ValueError):
return
if num_threads <= 0:
return
try:
cv2_module.setNumThreads(num_threads)
except AttributeError:
pass
except Exception:
# OpenCV may reject thread changes in constrained runtimes.
pass
[docs]def scatter_results(results, valid_indices, total_len, default=None):
"""Place batch results back at their original indices.
Batch UDF helpers may skip corrupt or ``None`` inputs before invoking an
expensive model. This helper restores those valid-row results to the
original batch length and fills skipped positions with *default*.
"""
if len(results) != len(valid_indices):
raise ValueError(
"scatter_results length mismatch: "
f"got {len(results)} results for {len(valid_indices)} indices"
)
output = [default] * total_len
seen = set()
for idx, result in zip(valid_indices, results):
# Validate every index to prevent silent row/result misalignment.
if idx < 0 or idx >= total_len:
raise IndexError(
f"scatter_results index out of range: {idx} for length {total_len}"
)
if idx in seen:
raise ValueError(f"Duplicate scatter_results index: {idx}")
seen.add(idx)
output[idx] = result
return output
[docs]def run_image_batch_inference(image_series, predict_fn, mode="RGB"):
"""Decode a native IMAGE batch, run inference, and scatter results.
``mode`` is the model-input normalization mode. Most vision models expect
RGB tensors, so the default converts valid IMAGE rows to RGB before
calling ``predict_fn``. Operators that intentionally consume another image
mode should pass it explicitly.
"""
import pandas as pd
from pyflink.multimodal.codec import decode_image_batch
image_arrays, valid_idx = decode_image_batch(image_series, mode=mode)
if not image_arrays:
return pd.Series([None] * len(image_series))
valid_results = predict_fn(image_arrays)
return pd.Series(scatter_results(valid_results, valid_idx, len(image_series)))
# ---------------------------------------------------------------------------
# Batch processing utilities (resize / normalize / tensor conversion)
# ---------------------------------------------------------------------------
[docs]def resize_batch(images, target_size, strategy="resize", interpolation=None):
"""Resize a batch of differently-sized ndarrays to a uniform size.
Intended for pandas batch UDFs to unify images from
:func:`safe_decode_batch` into a stackable batch.
Args:
images: List of ndarrays (RGB images), each potentially a
different size.
target_size: ``(height, width)`` target dimensions.
strategy: Resize strategy:
- ``"resize"``: Scale directly to target (may change aspect
ratio).
- ``"pad"``: Scale keeping aspect ratio, pad shorter side
with black.
- ``"crop"``: Scale keeping aspect ratio, center-crop to
target.
interpolation: Interpolation method. Defaults to
``cv2.INTER_LINEAR`` (if cv2 is available) or PIL LANCZOS.
Pass ``None`` to auto-select.
Returns:
np.ndarray with shape ``(N, H, W, C)``.
Example::
pixel_arrays, idx = safe_decode_batch(series, mode="RGB")
batch = resize_batch(pixel_arrays, (224, 224))
# batch.shape == (len(pixel_arrays), 224, 224, 3)
"""
if not images:
# Keep empty batches aligned with the default RGB preprocessing contract.
return np.empty((0, target_size[0], target_size[1], 3), dtype=np.uint8)
target_height, target_width = target_size
# Fast path: skip resize if all images already match target size
all_match = all(
pixel_array.shape[0] == target_height
and pixel_array.shape[1] == target_width
for pixel_array in images
)
if all_match:
normalized = [
pixel_array if pixel_array.ndim == 3 else pixel_array[:, :, np.newaxis]
for pixel_array in images
]
return np.stack(normalized)
# Prefer cv2 over PIL for batch resize: cv2.resize operates on ndarray
# in-place (C++) while PIL requires fromarray->resize->asarray round-trips.
# Note: default interpolation differs - cv2 uses INTER_LINEAR, PIL uses
# LANCZOS - so pixel values may differ slightly from _ImageResize (which
# always uses PIL). This is acceptable for ML preprocessing batches.
try:
import cv2 as _cv2
def resize_one(pixel_array, new_height, new_width):
interp = interpolation if interpolation is not None else _cv2.INTER_LINEAR
return _cv2.resize(
pixel_array, (new_width, new_height), interpolation=interp
)
except ImportError:
def resize_one(pixel_array, new_height, new_width):
from PIL import Image
pil_image = Image.fromarray(pixel_array)
interp = interpolation if interpolation is not None else Image.LANCZOS
pil_image = pil_image.resize((new_width, new_height), interp)
return np.asarray(pil_image)
resized = []
for pixel_array in images:
if (
pixel_array.shape[0] == target_height
and pixel_array.shape[1] == target_width
):
resized.append(
pixel_array
if pixel_array.ndim == 3
else pixel_array[:, :, np.newaxis]
)
continue
if strategy == "resize":
resized.append(
_resize_direct(
pixel_array, target_height, target_width, resize_one
)
)
elif strategy == "pad":
resized.append(
_resize_pad(pixel_array, target_height, target_width, resize_one)
)
elif strategy == "crop":
resized.append(
_resize_crop(pixel_array, target_height, target_width, resize_one)
)
else:
raise ValueError(
f"Unknown resize strategy: {strategy!r}. "
f"Expected 'resize', 'pad', or 'crop'."
)
return np.stack(resized)
def _resize_direct(pixel_array, target_height, target_width, resize_one):
"""Direct resize to target size (may change aspect ratio)."""
resized = resize_one(pixel_array, target_height, target_width)
# Ensure consistent 3D output for grayscale (match _resize_pad behavior)
if resized.ndim == 2:
resized = resized[:, :, np.newaxis]
return resized
def _resize_pad(pixel_array, target_height, target_width, resize_one):
"""Resize keeping aspect ratio, pad shorter side with black."""
height, width = pixel_array.shape[:2]
scale = min(target_height / height, target_width / width)
new_height = max(1, int(height * scale))
new_width = max(1, int(width * scale))
resized = resize_one(pixel_array, new_height, new_width)
# Ensure consistent 3D output for grayscale
if resized.ndim == 2:
resized = resized[:, :, np.newaxis]
channels = pixel_array.shape[2] if pixel_array.ndim == 3 else 1
canvas = np.zeros(
(target_height, target_width, channels), dtype=pixel_array.dtype
)
pad_top = (target_height - new_height) // 2
pad_left = (target_width - new_width) // 2
canvas[pad_top:pad_top + new_height, pad_left:pad_left + new_width] = resized
return canvas
def _resize_crop(pixel_array, target_height, target_width, resize_one):
"""Resize keeping aspect ratio, center-crop to target size."""
height, width = pixel_array.shape[:2]
scale = max(target_height / height, target_width / width)
new_height = max(target_height, int(height * scale))
new_width = max(target_width, int(width * scale))
resized = resize_one(pixel_array, new_height, new_width)
crop_top = (new_height - target_height) // 2
crop_left = (new_width - target_width) // 2
cropped = resized[
crop_top:crop_top + target_height,
crop_left:crop_left + target_width,
]
# Ensure consistent 3D output for grayscale (match _resize_pad behavior)
if cropped.ndim == 2:
cropped = cropped[:, :, np.newaxis]
return cropped
[docs]def normalize_batch(batch, mean=IMAGENET_RGB_MEAN, std=IMAGENET_RGB_STD):
"""Normalize a batch of images.
Converts uint8 [0, 255] images to float32 and normalizes with the
given mean and std. Defaults to ImageNet statistics.
Args:
batch: np.ndarray with shape ``(N, H, W, C)``, dtype uint8.
mean: Per-channel means (scaled to [0, 1]). Length must match
the number of channels in *batch*.
std: Per-channel standard deviations (scaled to [0, 1]).
Length must match the number of channels in *batch*.
Returns:
np.ndarray with shape ``(N, H, W, C)``, dtype float32.
Example::
batch = resize_batch(pixel_arrays, (224, 224))
normed = normalize_batch(batch) # ImageNet normalization
"""
if batch.ndim != 4:
raise ValueError(
f"normalize_batch expects a 4-D batch (N, H, W, C), "
f"got shape {batch.shape}"
)
channels = batch.shape[3]
if len(mean) != channels:
raise ValueError(
f"mean has {len(mean)} elements but batch has {channels} channels"
)
if len(std) != channels:
raise ValueError(
f"std has {len(std)} elements but batch has {channels} channels"
)
batch_f = batch.astype(np.float32) / 255.0
_mean = np.array(mean, dtype=np.float32).reshape(1, 1, 1, -1)
_std = np.array(std, dtype=np.float32).reshape(1, 1, 1, -1)
return (batch_f - _mean) / _std
[docs]def batch_to_tensor(batch, device=None):
"""Convert an ndarray batch to a PyTorch tensor.
Performs HWC -> CHW transposition (PyTorch convention) and optional
device transfer. If the input is uint8, it is converted to float32
and scaled to [0, 1] first; float32 input is used as-is.
Args:
batch: np.ndarray with shape ``(N, H, W, C)``.
device: Target device (str or ``torch.device``), e.g.
``"cuda:0"``. ``None`` means CPU.
Returns:
``torch.Tensor`` with shape ``(N, C, H, W)``, dtype float32.
Example::
batch = resize_batch(pixel_arrays, (224, 224))
normed = normalize_batch(batch)
tensor = batch_to_tensor(normed, device="cuda:0")
output = model(tensor)
"""
import torch
if batch.dtype == np.uint8:
batch = batch.astype(np.float32) / 255.0
elif batch.dtype != np.float32:
batch = batch.astype(np.float32)
# HWC -> CHW; .copy() is required because transpose creates
# a non-contiguous view which torch.from_numpy cannot consume.
tensor = torch.from_numpy(batch.transpose(0, 3, 1, 2).copy())
if device is not None:
tensor = tensor.to(device)
return tensor