################################################################################
# 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 transformation operators.
Transform operators accept either raw image bytes (``DataType.binary()``)
or a built-in ``IMAGE`` value (from ``image_decode()``). Pixel transforms
return built-in ``IMAGE`` values so chained operators can avoid repeated
JPEG/PNG encode-decode cycles. Boundary operators such as
``image_encode()``, ``image_compress()``, and ``image_convert_format()``
return encoded image bytes.
Typical pipeline::
from pyflink.multimodal.operators.image_transform import (
image_decode, image_encode, image_resize, image_blur,
)
decode = image_decode()
resize = image_resize(width=512, height=512)
blur = image_blur(radius=2)
encode = image_encode(format="JPEG")
df = df.with_columns("img", decode(col("raw_bytes")))
df = df.with_columns("img", resize(col("img")))
df = df.with_columns("img", blur(col("img")))
df = df.with_columns("jpg", encode(col("img")))
"""
import base64
import os
import tempfile
import math
from numbers import Integral, Real
from typing import TYPE_CHECKING
import numpy as np
from pyflink.dataframe import udf, DataType
from pyflink.table.udf import ScalarFunction
from pyflink.multimodal.codec import (
DEFAULT_IMAGE_ENCODE_QUALITY,
_IMAGE_DATA_ERRORS,
_PILImage,
_normalize_image_mode,
decode_image as decode_image_value,
decode_image_input,
detect_image_format,
encode_image_input,
image_array_to_pil,
image_array_to_pil_compatible,
image_mode_channels,
ndarray_to_image,
normalize_image_format,
pil_to_image,
validate_image_quality,
)
from pyflink.model.cache_manager import check_dependencies, prepare_and_load_model_handle
from pyflink.multimodal.utils import _build_or_apply_udf, _udf_runtime_kwargs
__all__ = [
# Codec boundary
"image_decode",
"image_encode",
"image_compress",
"image_convert_format",
"image_to_tensor",
# Pixel transforms
"image_convert_mode",
"image_resize",
"image_rescale",
"image_crop",
"image_crop_black_border",
"image_flip",
"image_blur",
"image_adjust_color",
# Model-backed
"image_remove_background",
]
if TYPE_CHECKING:
import pandas as pd
# Default configuration constants
_VALID_RESIZE_METHODS = {"nearest", "bilinear", "bicubic", "lanczos"}
_VALID_BLUR_TYPES = {"gaussian", "box", "mean"}
_VALID_CROP_TYPES = {"coordinate", "ratio", "center"}
_VALID_BLACK_BORDER_ALGORITHMS = {"threshold", "edge", "histogram", "auto"}
_VALID_FLIP_MODES = {"horizontal", "vertical", "rotate180"}
_VALID_TENSOR_LAYOUTS = {"CHW", "HWC"}
_REMBG_INPUT_ERRORS = _IMAGE_DATA_ERRORS + (RuntimeError,)
_VALID_DECODE_ON_ERROR = {"raise", "null"}
_VALID_IMAGE_ENCODE_OUTPUTS = {"bytes", "data_url"}
_IMAGE_DATA_URL_MIME_TYPES = {
"JPEG": "image/jpeg",
"PNG": "image/png",
"WEBP": "image/webp",
"BMP": "image/bmp",
"GIF": "image/gif",
"TIFF": "image/tiff",
}
def _normalize_decode_on_error(on_error):
if on_error not in _VALID_DECODE_ON_ERROR:
raise ValueError(
f"on_error must be one of {sorted(_VALID_DECODE_ON_ERROR)}, "
f"got {on_error!r}"
)
return on_error
def _normalize_image_encode_output(output):
if output not in _VALID_IMAGE_ENCODE_OUTPUTS:
raise ValueError(
f"output must be one of {sorted(_VALID_IMAGE_ENCODE_OUTPUTS)}, "
f"got {output!r}"
)
return output
def _image_data_url(encoded, output_format=None):
image_format = detect_image_format(encoded) or output_format or "JPEG"
mime_type = _IMAGE_DATA_URL_MIME_TYPES[image_format]
payload = base64.b64encode(encoded).decode("ascii")
return f"data:{mime_type};base64,{payload}"
def _scale_size(size, ratio):
return tuple(max(1, int(round(edge * ratio))) for edge in size)
def _validate_positive_dimension(name, value):
value = _validate_integer(name, value)
if value <= 0:
raise ValueError(f"{name} must be a positive integer, got {value!r}")
return value
def _validate_finite_number(name, value):
if (
not isinstance(value, Real)
or isinstance(value, bool)
or not math.isfinite(value)
):
raise ValueError(f"{name} must be a finite number, got {value!r}")
return float(value)
def _validate_positive_number(name, value):
value = _validate_finite_number(name, value)
if value <= 0:
raise ValueError(f"{name} must be positive, got {value!r}")
return value
def _validate_non_negative_number(name, value):
value = _validate_finite_number(name, value)
if value < 0:
raise ValueError(f"{name} must be >= 0, got {value!r}")
return value
def _validate_pair(name, value):
if value is None:
raise ValueError(f"{name} must be a sequence of two values")
try:
first, second = value
except (TypeError, ValueError) as e:
raise ValueError(f"{name} must be a sequence of two values") from e
return first, second
def _validate_integer(name, value):
if not isinstance(value, Integral) or isinstance(value, bool):
raise ValueError(f"{name} must be an integer, got {value!r}")
return int(value)
def _normalize_resize_method(method):
if not isinstance(method, str):
raise ValueError(f"method must be a string, got {method!r}")
normalized = method.lower()
if normalized not in _VALID_RESIZE_METHODS:
raise ValueError(
f"method must be one of {sorted(_VALID_RESIZE_METHODS)}, "
f"got {method!r}"
)
return normalized
def _normalize_tensor_layout(layout):
if not isinstance(layout, str):
raise ValueError(f"layout must be a string, got {layout!r}")
normalized = layout.upper()
if normalized not in _VALID_TENSOR_LAYOUTS:
raise ValueError(
f"layout must be one of {sorted(_VALID_TENSOR_LAYOUTS)}, "
f"got {layout!r}"
)
return normalized
def _image_array_to_float32(pixel_array):
if pixel_array.dtype == np.uint8:
return pixel_array.astype(np.float32) / 255.0
if pixel_array.dtype == np.uint16:
return pixel_array.astype(np.float32) / 65535.0
return pixel_array.astype(np.float32, copy=False)
def _crop_pixel_array(pixel_array, box):
left, top, right, bottom = box
h, w = pixel_array.shape[:2]
if left >= 0 and top >= 0 and right <= w and bottom <= h:
return pixel_array[top:bottom, left:right]
out_h = bottom - top
out_w = right - left
if pixel_array.ndim == 2:
output = np.zeros((out_h, out_w), dtype=pixel_array.dtype)
else:
output = np.zeros(
(out_h, out_w, pixel_array.shape[2]), dtype=pixel_array.dtype
)
src_left = max(left, 0)
src_top = max(top, 0)
src_right = min(right, w)
src_bottom = min(bottom, h)
if src_left >= src_right or src_top >= src_bottom:
return output
dst_left = src_left - left
dst_top = src_top - top
output[
dst_top:dst_top + (src_bottom - src_top),
dst_left:dst_left + (src_right - src_left),
] = pixel_array[src_top:src_bottom, src_left:src_right]
return output
def _center_crop_box(image_size, crop_size):
width, height = image_size
crop_width, crop_height = crop_size
left = int(round((width - crop_width) / 2.0))
top = int(round((height - crop_height) / 2.0))
return left, top, left + crop_width, top + crop_height
def _crop_size_from_ratio(image_size, crop_ratio):
width, height = image_size
width_ratio, height_ratio = _validate_pair("crop_ratio", crop_ratio)
if not (0 < width_ratio <= 1 and 0 < height_ratio <= 1):
raise ValueError(f"crop_ratio must contain values in (0, 1], got {crop_ratio}")
return (
max(1, int(round(width * width_ratio))),
max(1, int(round(height * height_ratio))),
)
# DecodeImage - binary bytes -> native ``IMAGE`` (pipeline entry point)
class _DecodeImage(ScalarFunction):
"""
Decode encoded image bytes (JPEG/PNG/...) into a decoded image value.
This is the entry point for zero-copy image pipelines. Place it
once at the beginning; downstream transform operators will pass
raw pixel data without re-encoding.
Usage::
df.with_column("img", image_decode()(col("raw_bytes")))
"""
def __init__(
self,
on_error="raise",
mode=None,
pixel_limit=None,
):
super().__init__()
self.on_error = _normalize_decode_on_error(on_error)
if mode is not None:
mode = _normalize_image_mode(mode)
if pixel_limit is not None:
pixel_limit = _validate_positive_dimension("pixel_limit", pixel_limit)
self.mode = mode
self.pixel_limit = pixel_limit
def open(self, function_context):
check_dependencies("PIL")
def _open_and_convert(self, image_bytes):
image = decode_image_value(
image_bytes, mode=self.mode, max_pixels=self.pixel_limit
)
if image is None:
raise ValueError("Failed to decode image bytes")
return image
def eval(self, image_bytes):
if image_bytes is None:
return None
if self.on_error == "raise":
# Fail fast by letting decode exceptions propagate.
return self._open_and_convert(image_bytes)
else:
# Null mode treats decode failures as null rows.
try:
return self._open_and_convert(image_bytes)
except _IMAGE_DATA_ERRORS:
return None
[docs]def image_decode(
*columns,
on_error="raise",
mode=None,
pixel_limit=None,
concurrency=None,
):
"""
Decode encoded image bytes into a decoded image.
This is the pipeline entry point for image data quality handling.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
on_error: Handling strategy for corrupt or undecodable images:
- ``"raise"`` (default): raise an exception and stop processing.
- ``"null"``: output null so downstream operators skip the row.
mode: Optional target mode to convert to, e.g. "RGB" or "L16".
``None`` preserves the decoded mode when it can be represented.
Set ``mode="RGB"`` explicitly before model operators that require
normalized RGB input.
pixel_limit: Decompression bomb guard. A small compressed file can
expand to billions of pixels and OOM the TaskManager. When set,
image dimensions are read from the file header (zero-decode cost)
and images with ``width * height > pixel_limit`` are rejected
before any pixel allocation. Follows the ``on_error`` strategy:
raises on ``"raise"``, returns null on ``"null"``.
``None`` (default) falls back to PIL's global
``MAX_IMAGE_PIXELS`` (~178 million).
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing decoded image values.
Example::
>>> # As a reusable variable
>>> decode = image_decode(on_error="null")
>>> df = df.with_columns("img", decode(col("raw_bytes")))
>>>
>>> # Inline
>>> df = df.with_columns("img", image_decode(col("raw_bytes"), on_error="null"))
"""
wrapper = udf(
_DecodeImage(on_error=on_error, mode=mode, pixel_limit=pixel_limit),
return_dtype=DataType.image(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# EncodeImage - native ``IMAGE`` to image binary bytes
class _EncodeImage(ScalarFunction):
"""
Encode a decoded image into compressed image bytes.
This is the exit point for zero-copy image pipelines. Place it
at the end when you need JPEG/PNG bytes for storage or output.
Args:
output_format: Target format, e.g. "JPEG", "PNG". If omitted, encoded
byte inputs keep their source format; decoded image values default
to JPEG because decoded images do not retain source container
metadata.
quality: Compression quality (1-100, JPEG/WebP only). Default 85.
Notes:
JPEG output normalizes unsupported modes to 8-bit RGB because JPEG has
no alpha/high-precision storage. Non-JPEG output preserves mode only
when Pillow can encode it directly; unsupported combinations raise.
"""
def __init__(self, output_format=None, quality=DEFAULT_IMAGE_ENCODE_QUALITY,
output="bytes"):
super().__init__()
self.output_format = normalize_image_format(
output_format, param_name="output_format", allow_none=True
)
self.quality = validate_image_quality(quality)
self.output = _normalize_image_encode_output(output)
def open(self, function_context):
check_dependencies("PIL")
def eval(self, image_input):
encoded = encode_image_input(
image_input, output_format=self.output_format, quality=self.quality
)
if encoded is None or self.output == "bytes":
return encoded
return _image_data_url(encoded, output_format=self.output_format)
[docs]def image_encode(
*columns,
format=None,
quality=DEFAULT_IMAGE_ENCODE_QUALITY,
output="bytes",
concurrency=None,
):
"""
Encode an image into compressed bytes or a base64 data URL.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
format: Target image format, e.g. ``"JPEG"``, ``"PNG"``, ``"WEBP"``,
``"BMP"``. ``None`` (default) keeps the source format for encoded
byte inputs; decoded image values default to ``"JPEG"`` because
decoded images do not retain original container metadata.
quality: Compression quality (1-100). Only affects lossy formats
such as JPEG and WebP. Default 85.
output: Output representation. ``"bytes"`` (default) preserves the
existing binary contract. ``"data_url"`` returns a string such as
``data:image/jpeg;base64,...`` for AI functions that accept image
URLs or inline data URLs.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing encoded image bytes or a data URL string.
Example::
>>> # As a reusable variable
>>> encode = image_encode(format="JPEG")
>>> df = df.with_columns("jpg_bytes", encode(col("img")))
>>>
>>> # Inline
>>> df = df.with_columns("jpg_bytes", image_encode(col("img"), format="JPEG"))
"""
output = _normalize_image_encode_output(output)
wrapper = udf(
_EncodeImage(output_format=format, quality=quality, output=output),
return_dtype=DataType.string() if output == "data_url" else DataType.binary(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageConvertMode - convert to a target image mode
class _ImageConvertMode(ScalarFunction):
"""
Convert an image to a target mode.
Args:
mode: Target mode, e.g. "RGB", "RGBA", "L", or "L16".
Usage::
df.with_column("gray", image_convert_mode(mode="L")(col("img")))
"""
def __init__(self, mode):
super().__init__()
self.mode = _normalize_image_mode(mode)
def open(self, function_context):
check_dependencies("PIL")
def eval(self, image_input):
try:
pixel_array = decode_image_input(image_input, mode=self.mode)
if pixel_array is None:
return None
return ndarray_to_image(pixel_array)
except _IMAGE_DATA_ERRORS:
return None
[docs]def image_convert_mode(*columns, mode, concurrency=None):
"""
Convert an image to a target mode.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
mode: Target image mode, e.g. ``"RGB"``, ``"RGBA"``, ``"L"``, or
``"L16"``. The image is converted to this mode before output.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing decoded image values.
Example::
>>> # As a reusable variable
>>> convert_mode = image_convert_mode(mode="L")
>>> df = df.with_columns("gray", convert_mode(col("img")))
>>>
>>> # Inline
>>> df = df.with_columns("gray", image_convert_mode(col("img"), mode="L"))
"""
wrapper = udf(
_ImageConvertMode(mode),
return_dtype=DataType.image(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageToTensor - fixed-shape image tensor boundary
class _ImageToTensor(ScalarFunction):
"""
Convert an image to a fixed-shape float32 tensor.
This is a pure conversion boundary — no resize or mode conversion is
performed. Use ``image_resize`` and ``image_convert_mode`` upstream
to prepare the input.
Args:
width: Expected input width in pixels (shape declaration).
height: Expected input height in pixels (shape declaration).
mode: Expected input mode, e.g. "RGB" or "L" (shape declaration).
layout: Tensor layout, "CHW" or "HWC".
"""
def __init__(self, width, height, mode="RGB", layout="CHW"):
super().__init__()
width = _validate_positive_dimension("width", width)
height = _validate_positive_dimension("height", height)
if mode is None:
raise ValueError("mode is required for image_to_tensor")
self.width = width
self.height = height
self.mode = _normalize_image_mode(mode)
self.channels = image_mode_channels(self.mode)
self.layout = _normalize_tensor_layout(layout)
def open(self, function_context):
check_dependencies("PIL")
def eval(self, image_input):
try:
pixel_array = decode_image_input(image_input)
if pixel_array is None:
return None
if pixel_array.ndim == 2:
pixel_array = pixel_array[:, :, np.newaxis]
h, w, c = pixel_array.shape
if w != self.width or h != self.height or c != self.channels:
return None
tensor = _image_array_to_float32(pixel_array)
if self.layout == "CHW":
tensor = tensor.transpose(2, 0, 1)
return np.ascontiguousarray(tensor)
except _IMAGE_DATA_ERRORS:
return None
[docs]def image_to_tensor(
*columns,
width,
height,
mode="RGB",
layout="CHW",
concurrency=None,
):
"""
Convert an image to a fixed-shape float32 tensor.
Pure conversion boundary — no resize or mode conversion is performed.
The output type is a fixed-shape ``TensorType``. Pixel values are converted
to ``float32`` and scaled to ``[0, 1]`` for uint8/uint16 inputs. Inputs
whose dimensions or channel count do not match the declared shape return
``None``.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
width: Expected input width in pixels. The input image must already
be this width (use ``image_resize`` upstream if needed).
height: Expected input height in pixels.
mode: Expected input image mode, e.g. ``"RGB"`` (3 channels) or
``"L"`` (1 channel). Determines the channel dimension of the
output tensor. Default ``"RGB"``.
layout: Tensor axis layout. ``"CHW"`` (default) produces shape
``(channels, height, width)``; ``"HWC"`` produces
``(height, width, channels)``. Accepted case-insensitively and
normalized to uppercase.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing a fixed-shape ``float32`` tensor.
Example::
>>> df = df.with_columns("img", image_decode(col("raw_bytes")))
>>> df = df.with_columns("img", image_convert_mode(col("img"), mode="RGB"))
>>> df = df.with_columns("img", image_resize(col("img"), width=224, height=224))
>>> df = df.with_columns("t", image_to_tensor(col("img"), width=224, height=224))
"""
width = _validate_positive_dimension("width", width)
height = _validate_positive_dimension("height", height)
if mode is None:
raise ValueError("mode is required for image_to_tensor")
normalized_layout = _normalize_tensor_layout(layout)
normalized_mode = _normalize_image_mode(mode)
channels = image_mode_channels(normalized_mode)
shape = (
(channels, height, width)
if normalized_layout == "CHW"
else (height, width, channels)
)
wrapper = udf(
_ImageToTensor(width=width, height=height, mode=normalized_mode, layout=layout),
return_dtype=DataType.tensor(DataType.float32(), shape=shape),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageResize - resize to fixed width and height
class _ImageResize(ScalarFunction):
"""
Resize an image to the specified output width and height.
Args:
width: Output width in pixels.
height: Output height in pixels.
method: Resize interpolation method.
"""
def __init__(self, width, height, method="lanczos"):
super().__init__()
width = _validate_positive_dimension("width", width)
height = _validate_positive_dimension("height", height)
self.width = width
self.height = height
self.method = _normalize_resize_method(method)
def open(self, function_context):
check_dependencies("PIL")
resampling = getattr(_PILImage, "Resampling", _PILImage)
self._resize_filter = {
"nearest": resampling.NEAREST,
"bilinear": resampling.BILINEAR,
"bicubic": resampling.BICUBIC,
"lanczos": resampling.LANCZOS,
}[self.method]
def eval(self, image_input):
try:
pixel_array = decode_image_input(image_input)
if pixel_array is None:
return None
size = (pixel_array.shape[1], pixel_array.shape[0])
new_size = (self.width, self.height)
if size == new_size and not isinstance(image_input, (bytes, bytearray)):
return image_input
pil_image = image_array_to_pil_compatible(pixel_array)
if size != new_size:
pil_image = pil_image.resize(new_size, self._resize_filter)
return pil_to_image(pil_image)
except _IMAGE_DATA_ERRORS:
return None
[docs]def image_resize(*columns, width, height, method="lanczos", concurrency=None):
"""
Resize an image to fixed output dimensions.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
width: Output width in pixels (positive integer).
height: Output height in pixels (positive integer).
method: Resize interpolation method — ``"lanczos"`` (default),
``"bilinear"``, ``"bicubic"``, or ``"nearest"``. Accepted
case-insensitively and normalized to lowercase.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing decoded image values.
Example::
>>> # As a reusable variable
>>> resize = image_resize(width=512, height=512)
>>> df = df.with_columns("resized", resize(col("img")))
>>>
>>> # Inline
>>> df = df.with_columns("resized", image_resize(col("img"), width=512, height=512))
"""
wrapper = udf(
_ImageResize(width=width, height=height, method=method),
return_dtype=DataType.image(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageBlur - Gaussian / Box blur
class _ImageBlur(ScalarFunction):
"""
Apply blur to an image.
Aligned with Data-Juicer ``ImageBlurMapper``.
Args:
radius: Blur radius in pixels. Default 2.
blur_type: Blur method - "gaussian", "box", or "mean". Default
"gaussian". ``box`` uses the configured radius; ``mean`` maps to
PIL's fixed mean blur kernel, matching Data-Juicer semantics.
"""
def __init__(self, radius=2, blur_type="gaussian"):
super().__init__()
radius = _validate_non_negative_number("blur radius", radius)
blur_type = blur_type.lower() if isinstance(blur_type, str) else blur_type
if blur_type not in _VALID_BLUR_TYPES:
raise ValueError(
f"blur_type must be one of {sorted(_VALID_BLUR_TYPES)}, "
f"got {blur_type!r}"
)
self.radius = radius
self.blur_type = blur_type
def open(self, function_context):
check_dependencies("PIL")
from PIL import ImageFilter
self._ImageFilter = ImageFilter
def eval(self, image_input):
try:
pixel_array = decode_image_input(image_input)
if pixel_array is None:
return None
pil_image = image_array_to_pil_compatible(pixel_array)
if self.blur_type == "mean":
pil_image = pil_image.filter(self._ImageFilter.BLUR)
elif self.blur_type == "box":
pil_image = pil_image.filter(self._ImageFilter.BoxBlur(self.radius))
else:
pil_image = pil_image.filter(
self._ImageFilter.GaussianBlur(self.radius)
)
return pil_to_image(pil_image)
except _IMAGE_DATA_ERRORS:
return None
[docs]def image_blur(*columns, radius=2, blur_type="gaussian", concurrency=None):
"""
Apply blur to an image.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
radius: Blur radius in pixels. Default 2.
blur_type: Blur algorithm — ``"gaussian"`` (default), ``"box"``, or
``"mean"``. ``"box"`` applies a box blur with the given radius;
``"mean"`` uses PIL's fixed-kernel mean blur (Data-Juicer
compatible).
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing decoded image values.
Example::
>>> # As a reusable variable
>>> blur = image_blur(radius=3)
>>> df = df.with_columns("blurred", blur(col("img")))
>>>
>>> # Inline
>>> df = df.with_columns("blurred", image_blur(col("img"), radius=3))
"""
wrapper = udf(
_ImageBlur(radius=radius, blur_type=blur_type),
return_dtype=DataType.image(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageCrop - coordinate / ratio / center crop
class _ImageCrop(ScalarFunction):
"""
Crop an image by coordinates, center ratio, or center size.
Args:
crop_type: Crop mode: "coordinate", "ratio", or "center".
crop_coords: Coordinate crop box [x1, y1, x2, y2].
crop_ratio: Center crop ratio [width_ratio, height_ratio].
crop_size: Center crop output size [width, height].
"""
def __init__(
self,
crop_type="center",
crop_coords=None,
crop_ratio=(0.8, 0.8),
crop_size=None,
):
super().__init__()
if crop_coords is not None:
crop_type = "coordinate"
if crop_type not in _VALID_CROP_TYPES:
raise ValueError(
f"crop_type must be one of {sorted(_VALID_CROP_TYPES)}, "
f"got {crop_type!r}"
)
self.crop_type = crop_type
self.crop_coords = crop_coords
self.crop_ratio = crop_ratio
self.crop_size = crop_size
if crop_type == "coordinate":
if crop_coords is None or len(crop_coords) != 4:
raise ValueError("crop_coords must be [x1, y1, x2, y2]")
left, top, right, bottom = (
_validate_integer("x1", crop_coords[0]),
_validate_integer("y1", crop_coords[1]),
_validate_integer("x2", crop_coords[2]),
_validate_integer("y2", crop_coords[3]),
)
self._validate_crop_box(left, top, right, bottom)
self.box = (left, top, right, bottom)
elif crop_type == "ratio":
# Validate early; the concrete crop box depends on each input size.
_crop_size_from_ratio((100, 100), crop_ratio)
self.box = None
else:
if crop_size is not None:
width, height = _validate_pair("crop_size", crop_size)
width = _validate_positive_dimension("crop_size width", width)
height = _validate_positive_dimension("crop_size height", height)
crop_size = (width, height)
_crop_size_from_ratio((100, 100), crop_ratio)
self.crop_size = crop_size
self.box = None
@staticmethod
def _validate_crop_box(left, top, right, bottom):
if right <= left or bottom <= top:
raise ValueError(
"crop box must satisfy right > left and bottom > top, "
f"got ({left}, {top}, {right}, {bottom})"
)
def open(self, function_context):
check_dependencies("PIL")
def _crop_box(self, pixel_array):
if self.crop_type == "coordinate":
return self.box
image_size = (pixel_array.shape[1], pixel_array.shape[0])
if self.crop_type == "ratio":
crop_size = _crop_size_from_ratio(image_size, self.crop_ratio)
return _center_crop_box(image_size, crop_size)
crop_size = self.crop_size
if crop_size is None:
crop_size = _crop_size_from_ratio(image_size, self.crop_ratio)
return _center_crop_box(image_size, crop_size)
def eval(self, image_input):
try:
pixel_array = decode_image_input(image_input)
if pixel_array is None:
return None
cropped = _crop_pixel_array(pixel_array, self._crop_box(pixel_array))
return ndarray_to_image(cropped)
except _IMAGE_DATA_ERRORS:
return None
[docs]def image_crop(
*columns,
crop_coords=None,
crop_type="center",
crop_ratio=(0.8, 0.8),
crop_size=None,
concurrency=None,
):
"""
Crop an image by coordinates, center ratio, or center size.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
crop_coords: Absolute pixel box ``(x1, y1, x2, y2)``. When provided,
``crop_type`` is set to ``"coordinate"`` automatically. Regions
outside the input image are padded with zero-valued pixels.
crop_type: Crop strategy — ``"center"`` (default), ``"ratio"``, or
``"coordinate"``. ``"center"`` crops ``crop_size`` pixels from the
center (falls back to ``crop_ratio`` if ``crop_size`` is ``None``).
``"ratio"`` center-crops by ``crop_ratio``. ``"coordinate"`` uses
the absolute box from ``crop_coords``.
crop_ratio: Fraction of width and height to keep, as
``(width_ratio, height_ratio)``. Values in ``(0, 1]``.
Default ``(0.8, 0.8)``.
crop_size: Fixed center-crop output size ``(width, height)`` in pixels.
Only used when ``crop_type="center"``. ``None`` (default) falls
back to ``crop_ratio``.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing decoded image values.
Example::
>>> # As a reusable variable
>>> crop = image_crop(crop_coords=(100, 100, 500, 500))
>>> df = df.with_columns("cropped", crop(col("img")))
>>>
>>> # Inline
>>> df = df.with_columns(
... "cropped", image_crop(col("img"), crop_coords=(100, 100, 500, 500))
... )
"""
if crop_coords is not None:
if len(crop_coords) != 4:
raise TypeError("crop_coords must have exactly 4 elements: x1, y1, x2, y2")
if crop_type not in ("center", "coordinate"):
raise ValueError("crop_coords cannot be combined with crop_type")
crop_type = "coordinate"
wrapper = udf(
_ImageCrop(
crop_type=crop_type,
crop_coords=crop_coords if crop_coords is not None else None,
crop_ratio=crop_ratio,
crop_size=crop_size,
),
return_dtype=DataType.image(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageCompress - encode bytes with a target quality
class _ImageCompress(ScalarFunction):
"""
Encode an image to compressed bytes with the specified format and quality.
Built-in IMAGE values store decoded pixels only. Compression is an
output-boundary operation and therefore returns encoded bytes instead of a
decoded image value. This
is a convenience wrapper around image encoding with a JPEG-oriented
default.
Args:
quality: Compression quality (1-100). Default 85.
output_format: Output format. Default "JPEG".
Corrupt encoded bytes return ``None``. Malformed decoded image input is
treated as invalid operator input and may fail fast.
"""
def __init__(self, quality=DEFAULT_IMAGE_ENCODE_QUALITY, output_format="JPEG"):
super().__init__()
self.quality = validate_image_quality(quality)
self.output_format = normalize_image_format(
output_format, param_name="output_format"
)
def open(self, function_context):
check_dependencies("PIL")
def eval(self, image_input):
return encode_image_input(
image_input, output_format=self.output_format, quality=self.quality
)
[docs]def image_compress(
*columns,
quality=DEFAULT_IMAGE_ENCODE_QUALITY,
format="JPEG",
concurrency=None,
):
"""
Compress an image to encoded bytes with the specified format and quality.
The input may be raw encoded image bytes or a decoded image value. The
output is always newly encoded image bytes. Corrupt encoded bytes return
``None``; malformed decoded image input may fail fast.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
quality: Compression quality (1-100). Default 85.
format: Output image format, e.g. ``"JPEG"``, ``"PNG"``, ``"WEBP"``.
Default ``"JPEG"``.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing encoded image bytes.
Example::
>>> # As a reusable variable
>>> compress = image_compress(quality=60)
>>> df = df.with_columns("compressed", compress(col("img")))
>>>
>>> # Inline
>>> df = df.with_columns("compressed", image_compress(col("img"), quality=60))
"""
wrapper = udf(
_ImageCompress(quality=quality, output_format=format),
return_dtype=DataType.binary(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageFormatConvert - format conversion (PNG -> JPEG, etc.)
class _ImageFormatConvert(ScalarFunction):
"""
Convert an image to encoded bytes in a different format.
Built-in IMAGE values store decoded pixels only. Format conversion is an
output-boundary operation and therefore returns encoded bytes.
Args:
output_format: Target format, e.g. "PNG", "JPEG", "WEBP", "BMP".
Corrupt encoded bytes return ``None``. Malformed decoded image input is
treated as invalid operator input and may fail fast.
"""
def __init__(self, output_format):
super().__init__()
self.output_format = normalize_image_format(
output_format, param_name="output_format"
)
def open(self, function_context):
check_dependencies("PIL")
def eval(self, image_input):
return encode_image_input(image_input, output_format=self.output_format)
# ImageBrightnessAdjust - brightness / contrast / saturation
class _ImageBrightnessAdjust(ScalarFunction):
"""
Adjust brightness, contrast, and saturation of an image.
Each factor is a multiplier: 1.0 = no change, < 1.0 = decrease, > 1.0 = increase.
Args:
brightness: Brightness factor. Default 1.0.
contrast: Contrast factor. Default 1.0.
saturation: Saturation factor. Default 1.0.
"""
def __init__(self, brightness=1.0, contrast=1.0, saturation=1.0):
super().__init__()
brightness = _validate_non_negative_number("brightness", brightness)
contrast = _validate_non_negative_number("contrast", contrast)
saturation = _validate_non_negative_number("saturation", saturation)
self.brightness = brightness
self.contrast = contrast
self.saturation = saturation
def open(self, function_context):
check_dependencies("PIL")
from PIL import ImageEnhance
self._ImageEnhance = ImageEnhance
def eval(self, image_input):
try:
pixel_array = decode_image_input(image_input)
if pixel_array is None:
return None
pil_image = image_array_to_pil_compatible(pixel_array)
if self.brightness != 1.0:
pil_image = self._ImageEnhance.Brightness(pil_image).enhance(
self.brightness
)
if self.contrast != 1.0:
pil_image = self._ImageEnhance.Contrast(pil_image).enhance(
self.contrast
)
if self.saturation != 1.0:
pil_image = self._ImageEnhance.Color(pil_image).enhance(
self.saturation
)
return pil_to_image(pil_image)
except _IMAGE_DATA_ERRORS:
return None
[docs]def image_adjust_color(
*columns,
brightness=1.0,
contrast=1.0,
saturation=1.0,
concurrency=None,
):
"""
Adjust image brightness, contrast, and saturation.
Each factor is a multiplier: ``1.0`` = no change, ``< 1.0`` = decrease,
``> 1.0`` = increase.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
brightness: Brightness multiplier (>= 0). Default ``1.0``.
contrast: Contrast multiplier (>= 0). Default ``1.0``.
saturation: Saturation (color intensity) multiplier (>= 0).
Default ``1.0``.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing decoded image values.
Example::
>>> # As a reusable variable
>>> adjust = image_adjust_color(brightness=1.5)
>>> df = df.with_columns("bright", adjust(col("img")))
>>>
>>> # Inline
>>> df = df.with_columns("bright", image_adjust_color(col("img"), brightness=1.5))
"""
wrapper = udf(
_ImageBrightnessAdjust(
brightness=brightness, contrast=contrast, saturation=saturation
),
return_dtype=DataType.image(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageFlip - horizontal / vertical / both
class _ImageFlip(ScalarFunction):
"""
Flip an image horizontally, vertically, or rotate 180 degrees.
Args:
mode: "horizontal", "vertical", or "rotate180". Default "horizontal".
"""
def __init__(self, mode="horizontal"):
super().__init__()
mode = mode.lower() if isinstance(mode, str) else mode
if mode not in _VALID_FLIP_MODES:
raise ValueError(
f"mode must be one of {sorted(_VALID_FLIP_MODES)}, got {mode!r}"
)
self.mode = mode
def open(self, function_context):
check_dependencies("PIL")
def eval(self, image_input):
try:
pixel_array = decode_image_input(image_input)
if pixel_array is None:
return None
if self.mode == "rotate180":
flipped = np.flip(pixel_array, axis=(0, 1)).copy()
elif self.mode == "horizontal":
flipped = pixel_array[:, ::-1].copy()
else:
flipped = pixel_array[::-1, :].copy()
return ndarray_to_image(flipped)
except _IMAGE_DATA_ERRORS:
return None
[docs]def image_flip(*columns, mode="horizontal", concurrency=None):
"""
Flip an image horizontally, vertically, or rotate 180 degrees.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
mode: Flip direction — ``"horizontal"`` (default), ``"vertical"``,
or ``"rotate180"`` (equivalent to flipping both axes).
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing decoded image values.
Example::
>>> # As a reusable variable
>>> flip = image_flip(mode="horizontal")
>>> df = df.with_columns("flipped", flip(col("img")))
>>>
>>> # Inline
>>> df = df.with_columns("flipped", image_flip(col("img"), mode="horizontal"))
"""
wrapper = udf(
_ImageFlip(mode=mode),
return_dtype=DataType.image(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageBlackBorderCrop - detect and remove black borders
class _ImageBlackBorderCrop(ScalarFunction):
"""
Detect and crop black borders from an image.
Supports threshold, histogram, edge, and auto detection modes. If no
meaningful border is detected, the original decoded image value is returned.
Args:
detect_algorithm: "threshold", "histogram", "edge", or "auto".
black_threshold: Pixel value below which content is treated as black.
edge_sensitivity: Edge detector sensitivity in [0.1, 2.0].
min_border_size: Minimum border width in pixels.
Detection converts the image to 8-bit luminance first. ``black_threshold``
is therefore interpreted on a 0-255 scale, including for L16/RGB16 inputs.
"""
def __init__(
self,
detect_algorithm="auto",
black_threshold=10,
edge_sensitivity=1.0,
min_border_size=1,
):
super().__init__()
if (
not isinstance(detect_algorithm, str)
or detect_algorithm not in _VALID_BLACK_BORDER_ALGORITHMS
):
raise ValueError(
"detect_algorithm must be one of "
f"{sorted(_VALID_BLACK_BORDER_ALGORITHMS)}, "
f"got {detect_algorithm!r}"
)
black_threshold = _validate_integer("threshold", black_threshold)
if black_threshold < 0 or black_threshold > 255:
raise ValueError(
f"threshold must be in [0, 255], got {black_threshold!r}"
)
edge_sensitivity = _validate_finite_number(
"edge_sensitivity", edge_sensitivity
)
if edge_sensitivity < 0.1 or edge_sensitivity > 2.0:
raise ValueError(
f"edge_sensitivity must be in [0.1, 2.0], "
f"got {edge_sensitivity!r}"
)
min_border_size = _validate_positive_dimension(
"min_border_size", min_border_size
)
self.detect_algorithm = detect_algorithm
self.black_threshold = black_threshold
self.edge_sensitivity = edge_sensitivity
self.min_border_size = min_border_size
def open(self, function_context):
check_dependencies("PIL")
@staticmethod
def _box_from_masks(row_mask, col_mask):
if not row_mask.any() or not col_mask.any():
return None
rows = np.where(row_mask)[0]
cols = np.where(col_mask)[0]
return int(cols[0]), int(rows[0]), int(cols[-1]) + 1, int(rows[-1]) + 1
def _threshold_box(self, gray):
content = gray > self.black_threshold
return self._box_from_masks(content.any(axis=1), content.any(axis=0))
def _histogram_box(self, gray):
row_mask = np.percentile(gray, 90, axis=1) > self.black_threshold
col_mask = np.percentile(gray, 90, axis=0) > self.black_threshold
return self._box_from_masks(row_mask, col_mask)
def _edge_box(self, gray):
gray_f = gray.astype(np.float32)
edge_threshold = max(1.0, 10.0 / self.edge_sensitivity)
row_signal = np.abs(np.diff(gray_f.mean(axis=1), prepend=0.0))
col_signal = np.abs(np.diff(gray_f.mean(axis=0), prepend=0.0))
row_edges = np.where(row_signal > edge_threshold)[0]
col_edges = np.where(col_signal > edge_threshold)[0]
if len(row_edges) < 2 or len(col_edges) < 2:
return None
return (
int(col_edges[0]),
int(row_edges[0]),
int(col_edges[-1]) + 1,
int(row_edges[-1]) + 1,
)
def _is_effective_box(self, box, width, height):
if box is None:
return False
left, top, right, bottom = box
if left < 0 or top < 0 or right > width or bottom > height:
return False
if right <= left or bottom <= top:
return False
borders = (left, top, width - right, height - bottom)
return max(borders) >= self.min_border_size
def _detect_box(self, gray):
algorithms = (
("threshold", "histogram", "edge")
if self.detect_algorithm == "auto"
else (self.detect_algorithm,)
)
height, width = gray.shape[:2]
detectors = {
"threshold": self._threshold_box,
"histogram": self._histogram_box,
"edge": self._edge_box,
}
for algorithm in algorithms:
box = detectors[algorithm](gray)
if self._is_effective_box(box, width, height):
return box
return None
def eval(self, image_input):
try:
pixel_array = decode_image_input(image_input)
if pixel_array is None:
return None
pil_image = image_array_to_pil_compatible(pixel_array)
gray = np.array(pil_image.convert("L"))
box = self._detect_box(gray)
if box is None:
return pil_to_image(pil_image)
pil_image = pil_image.crop(box)
return pil_to_image(pil_image)
except _IMAGE_DATA_ERRORS:
return None
[docs]def image_crop_black_border(
*columns,
threshold=None,
detect_algorithm="auto",
black_threshold=None,
edge_sensitivity=1.0,
min_border_size=1,
concurrency=None,
):
"""
Detect and remove black borders from an image.
If every pixel is treated as border (e.g. an all-black image), the
original decoded image value is returned unchanged.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
threshold: Legacy alias for ``black_threshold``. Cannot be combined
with an explicit ``black_threshold``.
detect_algorithm: Detection strategy — ``"auto"`` (default),
``"threshold"``, ``"histogram"``, or ``"edge"``. ``"auto"`` tries
threshold, histogram, and edge detection in order and uses the
first effective crop box.
black_threshold: Pixel luminance value (0-255) at or below which
content is treated as black border. Default ``10``. Detection
converts the image to 8-bit luminance first, so this scale
applies even for L16/RGB16 inputs.
edge_sensitivity: Sensitivity for the ``"edge"`` detector, in
``[0.1, 2.0]``. Higher values detect subtler borders. Default
``1.0``.
min_border_size: Minimum border width in pixels before a crop is
applied. Default ``1``.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing decoded image values.
Example::
>>> # As a reusable variable
>>> crop_border = image_crop_black_border(black_threshold=15)
>>> df = df.with_columns("clean", crop_border(col("img")))
>>>
>>> # Inline
>>> df = df.with_columns("clean", image_crop_black_border(col("img"), black_threshold=15))
"""
if threshold is not None and black_threshold is not None:
raise ValueError(
"threshold cannot be combined with black_threshold"
)
if black_threshold is None:
black_threshold = threshold if threshold is not None else 10
wrapper = udf(
_ImageBlackBorderCrop(
detect_algorithm=detect_algorithm,
black_threshold=black_threshold,
edge_sensitivity=edge_sensitivity,
min_border_size=min_border_size,
),
return_dtype=DataType.image(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageRescale - resize by a multiplicative scale
class _ImageRescale(ScalarFunction):
"""
Rescale an image by a multiplicative factor while preserving aspect ratio.
Args:
scale: Positive scale factor.
method: Resize interpolation method.
"""
def __init__(self, scale, method="lanczos"):
super().__init__()
scale = _validate_positive_number("scale", scale)
self.scale = scale
self.method = _normalize_resize_method(method)
def open(self, function_context):
check_dependencies("PIL")
resampling = getattr(_PILImage, "Resampling", _PILImage)
self._resize_filter = {
"nearest": resampling.NEAREST,
"bilinear": resampling.BILINEAR,
"bicubic": resampling.BICUBIC,
"lanczos": resampling.LANCZOS,
}[self.method]
def eval(self, image_input):
try:
pixel_array = decode_image_input(image_input)
if pixel_array is None:
return None
size = (pixel_array.shape[1], pixel_array.shape[0])
new_size = _scale_size(size, self.scale)
if size == new_size and not isinstance(image_input, (bytes, bytearray)):
return image_input
pil_image = image_array_to_pil_compatible(pixel_array)
if size != new_size:
pil_image = pil_image.resize(new_size, self._resize_filter)
return pil_to_image(pil_image)
except _IMAGE_DATA_ERRORS:
return None
[docs]def image_rescale(*columns, scale, method="lanczos", concurrency=None):
"""
Rescale an image by a multiplicative factor, preserving aspect ratio.
Each dimension is scaled by ``scale`` and rounded to the nearest pixel
(minimum 1).
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
scale: Positive scale factor. ``0.5`` halves each dimension;
``2.0`` doubles them.
method: Resize interpolation method — ``"lanczos"`` (default),
``"bilinear"``, ``"bicubic"``, or ``"nearest"``. Accepted
case-insensitively and normalized to lowercase.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF producing decoded image values.
Example::
>>> # As a reusable variable
>>> rescale = image_rescale(scale=0.5)
>>> df = df.with_columns("small", rescale(col("img")))
>>>
>>> # Inline
>>> df = df.with_columns("small", image_rescale(col("img"), scale=0.5))
"""
wrapper = udf(
_ImageRescale(scale=scale, method=method),
return_dtype=DataType.image(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageRemoveBackground - background removal (rembg)
class _ImageRemoveBackground(ScalarFunction):
"""
Remove background from an image using rembg (U2-Net based).
Returns an image with transparent or colored background.
Aligned with Data-Juicer ``ImageRemoveBackgroundMapper``.
Args:
alpha_matting: Enable alpha matting for edge refinement. Default False.
alpha_matting_foreground_threshold: Foreground threshold for alpha
matting. Default 240.
alpha_matting_background_threshold: Background threshold for alpha
matting. Default 10.
alpha_matting_erode_size: Erode size for alpha matting. Default 10.
bgcolor: Background color as (R, G, B, A) tuple. Default None
(transparent).
model_sharing: Model sharing mode for the rembg session.
Requires: ``pip install rembg``
Usage::
rembg = image_remove_background(alpha_matting=True)
df.with_column("nobg", rembg(col("img")))
"""
@staticmethod
def _prepare_rembg_import_environment():
import builtins
import inspect
import sys
active_print = getattr(builtins, "print", None)
if not inspect.ismethod(active_print):
return
module_name = getattr(active_print, "__module__", None)
function_name = getattr(active_print, "__name__", None)
if not module_name or not function_name or not function_name.isidentifier():
return
module = sys.modules.get(module_name) if module_name else None
if module is None:
return
current = getattr(module, function_name, None)
if current is not None and current is not active_print:
# PyFlink's Beam worker can replace builtins.print with a bound
# method. During rembg import, numba resolves globals through
# __module__/__name__, so register only that redirected method.
setattr(module, function_name, active_print)
def __init__(
self,
alpha_matting=False,
alpha_matting_foreground_threshold=240,
alpha_matting_background_threshold=10,
alpha_matting_erode_size=10,
bgcolor=None,
model_sharing=None,
num_gpus=None,
gpu_type=None,
):
super().__init__()
if not isinstance(alpha_matting, bool):
raise ValueError(
f"alpha_matting must be a boolean, got {alpha_matting!r}"
)
if bgcolor is not None:
if (not isinstance(bgcolor, (tuple, list))
or len(bgcolor) != 4
or not all(isinstance(c, Integral)
and not isinstance(c, bool)
and 0 <= c <= 255
for c in bgcolor)):
raise ValueError(
"bgcolor must be a 4-tuple (R, G, B, A) with values "
f"in [0, 255], got {bgcolor!r}"
)
bgcolor = tuple(int(c) for c in bgcolor)
for name, value in (
("alpha_matting_foreground_threshold",
alpha_matting_foreground_threshold),
("alpha_matting_background_threshold",
alpha_matting_background_threshold),
):
value = _validate_integer(name, value)
if not 0 <= value <= 255:
raise ValueError(
f"{name} must be an integer in [0, 255], got {value!r}"
)
if name == "alpha_matting_foreground_threshold":
alpha_matting_foreground_threshold = value
else:
alpha_matting_background_threshold = value
alpha_matting_erode_size = _validate_positive_dimension(
"alpha_matting_erode_size", alpha_matting_erode_size
)
self.alpha_matting = alpha_matting
self.alpha_matting_foreground_threshold = alpha_matting_foreground_threshold
self.alpha_matting_background_threshold = alpha_matting_background_threshold
self.alpha_matting_erode_size = alpha_matting_erode_size
self.bgcolor = bgcolor
self.model_sharing = model_sharing
self._num_gpus = num_gpus
self._gpu_type = gpu_type
self._rembg_options = {
"alpha_matting": self.alpha_matting,
"alpha_matting_foreground_threshold": (
self.alpha_matting_foreground_threshold
),
"alpha_matting_background_threshold": (
self.alpha_matting_background_threshold
),
"alpha_matting_erode_size": self.alpha_matting_erode_size,
"bgcolor": self.bgcolor,
}
def open(self, function_context):
os.environ.setdefault(
"NUMBA_CACHE_DIR",
os.path.join(tempfile.gettempdir(), "pyflink-multimodal-numba-cache"),
)
import pandas as pd
from pyflink.model.backends.rembg import RembgModelAdapter
self._pd = pd
self._prepare_rembg_import_environment()
self._model_handle = prepare_and_load_model_handle(
adapter_cls=RembgModelAdapter,
config={},
function_context=function_context,
model_sharing=self.model_sharing,
dependencies=("rembg",),
model_id="rembg",
requested_num_gpus=self._num_gpus,
requested_gpu_type=self._gpu_type,
)
self._model_handle.register_operation(
"image_remove_background", _ImageRemoveBackground._predict_batch
)
@staticmethod
def _pixel_array_to_pil(pixel_array):
# rembg supports PIL input. Keep pixels decoded at this model boundary
# to avoid an unnecessary PNG encode before inference.
return image_array_to_pil(pixel_array, mode="RGB")
@staticmethod
def _result_to_image_value(result):
if result is None:
return None
if isinstance(result, (bytes, bytearray)):
pixel_array = decode_image_input(result)
if pixel_array is None:
return None
return ndarray_to_image(pixel_array)
if isinstance(result, np.ndarray):
return ndarray_to_image(result)
return pil_to_image(result)
@staticmethod
def _predict_batch(model, pil_images, options):
from rembg import remove
results = []
for pil_image in pil_images:
try:
result = remove(pil_image, session=model, **options)
results.append(_ImageRemoveBackground._result_to_image_value(result))
except _IMAGE_DATA_ERRORS:
results.append(None)
return results
def _scatter_results_as_series(self, results, valid_idx, total_len):
from pyflink.multimodal.utils import scatter_results
return self._pd.Series(
scatter_results(results, valid_idx, total_len, default=None)
)
def close(self):
if getattr(self, "_model_handle", None) is not None:
self._model_handle.release()
self._model_handle = None
def eval(self, image_batch: "pd.Series") -> "pd.Series":
pd = getattr(self, "_pd", None)
if pd is None:
import pandas as pd
self._pd = pd
pil_images = []
valid_idx = []
total_len = len(image_batch)
for idx, image_input in enumerate(image_batch):
if image_input is None:
continue
try:
pixel_array = decode_image_input(image_input, mode="RGB")
if pixel_array is None:
continue
pil_images.append(self._pixel_array_to_pil(pixel_array))
valid_idx.append(idx)
except _REMBG_INPUT_ERRORS:
continue
if pil_images:
results = self._model_handle.call(
"image_remove_background", pil_images, self._rembg_options
)
else:
results = []
return self._scatter_results_as_series(results, valid_idx, total_len)
[docs]def image_remove_background(
*columns,
alpha_matting=False,
alpha_matting_foreground_threshold=240,
alpha_matting_background_threshold=10,
alpha_matting_erode_size=10,
bgcolor=None,
model_sharing=None,
concurrency=None,
batch_size=None,
num_gpus=None,
gpu_type=None,
):
"""
Remove the background from an image using rembg (U2-Net based).
Requires ``pip install rembg``. This is a pandas batch UDF, but rembg
inference is currently invoked per image inside each batch. ``num_gpus`` /
``gpu_type`` are forwarded as Flink resource hints; the rembg runtime
decides whether the session actually uses GPU execution.
Args:
*columns: Optional image columns. When provided, the UDF is applied
directly instead of returning a factory.
alpha_matting: Enable alpha matting for smoother edge refinement.
Default ``False``.
alpha_matting_foreground_threshold: Foreground confidence threshold
for alpha matting (0-255). Default ``240``.
alpha_matting_background_threshold: Background confidence threshold
for alpha matting (0-255). Default ``10``.
alpha_matting_erode_size: Erosion kernel size for alpha matting
refinement in pixels. Default ``10``.
bgcolor: Background fill color as ``(R, G, B, A)`` tuple. ``None``
(default) produces a transparent background.
model_sharing: Model sharing mode for the rembg session across
parallel subtasks.
concurrency: UDF concurrency. ``None`` uses the framework default.
batch_size: Pandas batch size. ``None`` uses the framework default.
num_gpus: Fractional GPU count per subtask, e.g. ``0.5``. ``None``
runs on CPU.
gpu_type: Required GPU type, e.g. ``"A10"``. ``None`` accepts any
available GPU.
Returns:
A UDF producing decoded image values.
Example::
>>> # As a reusable variable
>>> remove_bg = image_remove_background(alpha_matting=True)
>>> df = df.with_columns("nobg", remove_bg(col("img")))
>>>
>>> # Inline
>>> df = df.with_columns("nobg", image_remove_background(col("img")))
"""
wrapper = udf(
_ImageRemoveBackground(
alpha_matting=alpha_matting,
alpha_matting_foreground_threshold=alpha_matting_foreground_threshold,
alpha_matting_background_threshold=alpha_matting_background_threshold,
alpha_matting_erode_size=alpha_matting_erode_size,
bgcolor=bgcolor,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
return_dtype=DataType.image(),
func_type="pandas",
**_udf_runtime_kwargs(
concurrency=concurrency,
batch_size=batch_size,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
)
return _build_or_apply_udf(wrapper, *columns)
# ═══════════════════════════════════════════════════════════════════════════
# GenerateImage — seed → built-in ``IMAGE`` (synthetic source for demo / tests)
# ═══════════════════════════════════════════════════════════════════════════
_GENERATE_IMAGE_FORMATS = {"PNG", "JPEG", "BMP"}
class _GenerateImage(ScalarFunction):
"""
Generate a synthetic solid-color image deterministically from ``seed``.
Produces a decoded image value. The value stores decoded pixels only and
does not retain encoded format or quality metadata. The fill color is
derived from ``seed`` so that identical seeds produce identical output —
useful for reproducible demos and tests.
Args:
width: Image width in pixels. Default 512.
height: Image height in pixels. Default 512.
format: Intended downstream encoding hint, one of ``"PNG"``,
``"JPEG"``, ``"BMP"`` (case-insensitive). Default ``"PNG"``.
Validated for forward compatibility but NOT stored on the
produced decoded image value.
"""
def __init__(self, width=512, height=512, format="PNG"):
super().__init__()
if not isinstance(width, int) or width <= 0:
raise ValueError(f"width must be a positive integer, got {width!r}")
if not isinstance(height, int) or height <= 0:
raise ValueError(f"height must be a positive integer, got {height!r}")
if not isinstance(format, str):
raise ValueError(f"format must be a string, got {format!r}")
fmt_upper = format.upper()
if fmt_upper not in _GENERATE_IMAGE_FORMATS:
raise ValueError(
f"format must be one of {sorted(_GENERATE_IMAGE_FORMATS)}, "
f"got {format!r}"
)
self.width = width
self.height = height
self.format = fmt_upper
def open(self, function_context):
check_dependencies("PIL")
def eval(self, seed):
if seed is None:
return None
# Coerce seed to int. Strings hash deterministically; numerics cast.
if isinstance(seed, int):
s = seed
elif isinstance(seed, (bytes, bytearray)):
s = int.from_bytes(bytes(seed)[:8] or b"\x00", "big", signed=False)
else:
try:
s = int(seed)
except (TypeError, ValueError):
import hashlib
s = int.from_bytes(
hashlib.sha256(str(seed).encode("utf-8")).digest()[:8], "big"
)
s = abs(s)
color = (s % 256, (s * 37) % 256, (s * 73) % 256)
pil_image = _PILImage.new("RGB", (self.width, self.height), color)
return pil_to_image(pil_image)
def generate_image(width=512, height=512, format="PNG"):
"""
Create an image generation UDF (seed → built-in ``IMAGE`` value).
Produces a deterministic solid-color image whose color is derived from
the input ``seed``. Intended for demos, end-to-end tests, and synthetic
datasets.
Usage::
gen = generate_image(width=256, height=256, format="JPEG")
df.with_column("img", gen(col("seed")))
"""
wrapper = udf(
_GenerateImage(width=width, height=height, format=format),
return_dtype=DataType.image(),
)
return _build_or_apply_udf(wrapper)