################################################################################
# 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 embedding and similarity operators.
Operators in this module generate vector embeddings from images or
compute image-text similarity scores using CLIP.
All operators in this module use pandas batch UDFs and run CLIP inference
once per input batch.
Example::
>>> from pyflink.multimodal.operators.image_embed import (
... image_embedding, image_text_similarity,
... )
>>> embed = image_embedding(model="ViT-B/32")
>>> df = df.with_column("vector", embed(col("img")))
>>> # Fixed text mode — compare all images against one prompt
>>> sim = image_text_similarity(text="a photo of a cat")
>>> df = df.with_column("score", sim(col("img")))
>>> # Per-row text mode — each row has its own text
>>> sim = image_text_similarity()
>>> df = df.with_column("score", sim(col("img"), col("text")))
Runtime args:
``concurrency``, ``batch_size``, ``num_gpus``, and ``gpu_type`` are
forwarded to the DataFrame UDF runtime.
Requires: ``pip install open_clip_torch torch pillow``
"""
from typing import TYPE_CHECKING
from pyflink.dataframe import udf, DataType
from pyflink.model.backends.open_clip import OpenClipModelAdapter
from pyflink.model.cache_manager import prepare_and_load_model_handle
from pyflink.multimodal.codec import decode_image_batch
from pyflink.multimodal.utils import (
_build_or_apply_udf,
scatter_results,
_udf_runtime_kwargs,
run_image_batch_inference,
)
from pyflink.table.udf import ScalarFunction
if TYPE_CHECKING:
import pandas as pd
__all__ = [
"image_embedding",
"image_text_similarity",
]
class _ImageEmbedding(ScalarFunction):
"""Generate CLIP embedding vectors for image batches."""
def __init__(self, model="ViT-B/32", pretrained="openai", model_sharing=None,
num_gpus=None, gpu_type=None):
super().__init__()
self.model_name = model
self.pretrained = pretrained
self.model_sharing = model_sharing
self._num_gpus = num_gpus
self._gpu_type = gpu_type
@staticmethod
def _predict_batch(model, pixel_arrays):
return model.encode_images(pixel_arrays).cpu().tolist()
def open(self, function_context):
self._model_handle = prepare_and_load_model_handle(
adapter_cls=OpenClipModelAdapter,
config={},
function_context=function_context,
model_sharing=self.model_sharing,
dependencies=("open_clip", "torch", "PIL"),
model_id=OpenClipModelAdapter.build_model_id(
self.model_name, self.pretrained
),
requested_num_gpus=self._num_gpus,
requested_gpu_type=self._gpu_type,
)
self._model_handle.register_operation(
"image_embedding", _ImageEmbedding._predict_batch
)
def close(self):
if getattr(self, "_model_handle", None) is not None:
self._model_handle.release()
self._model_handle = None
def eval(self, image_series: "pd.Series") -> "pd.Series":
return run_image_batch_inference(
image_series,
lambda image_arrays: self._model_handle.call(
"image_embedding", image_arrays
),
)
[docs]def image_embedding(
*columns,
model="ViT-B/32",
pretrained="openai",
model_sharing=None,
concurrency=None,
batch_size=None,
num_gpus=None,
gpu_type=None,
):
"""
Create an image embedding UDF (CLIP / open_clip-based).
Requires ``pip install open_clip_torch torch Pillow``. This is a pandas
batch UDF that supports GPU acceleration via ``num_gpus`` / ``gpu_type``.
Args:
*columns: Optional image column(s). When provided, the UDF is
applied directly instead of returning a factory.
model: CLIP model architecture name. Default ``"ViT-B/32"``.
pretrained: Pretrained weights checkpoint. Default ``"openai"``.
model_sharing: Model sharing mode across parallel subtasks.
``None`` uses per-process caching.
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 that returns a L2-normalized ``list[float32]`` embedding vector,
or ``None`` for null image inputs. The vector dimension depends on the
model (e.g. 512 for ``ViT-B/32``).
Example::
>>> # As a reusable variable
>>> embed = image_embedding(model="ViT-B/32")
>>> df = df.with_column("vector", embed(col("img")))
>>>
>>> # Inline
>>> df = df.with_column("vector", image_embedding(col("img")))
"""
wrapper = udf(
_ImageEmbedding(
model=model,
pretrained=pretrained,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
func_type="pandas",
return_dtype=DataType.list(DataType.float32()),
**_udf_runtime_kwargs(
concurrency=concurrency,
batch_size=batch_size,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
)
return _build_or_apply_udf(wrapper, *columns)
class _ImageFixedTextSimilarity(ScalarFunction):
"""Compute image-text similarity against a fixed text prompt."""
def __init__(self, text, model="ViT-B/32", pretrained="openai",
model_sharing=None, num_gpus=None, gpu_type=None):
super().__init__()
if not isinstance(text, str) or not text.strip():
raise ValueError("image_text_similarity requires non-empty text.")
self.text = text
self.model_name = model
self.pretrained = pretrained
self.model_sharing = model_sharing
self._num_gpus = num_gpus
self._gpu_type = gpu_type
@staticmethod
def _predict_batch(model, pixel_arrays, text):
return model.score_images_against_text(pixel_arrays, text)
def open(self, function_context):
# Reuses the same CLIP model handle as ImageEmbedding.
self._model_handle = prepare_and_load_model_handle(
adapter_cls=OpenClipModelAdapter,
config={},
function_context=function_context,
model_sharing=self.model_sharing,
dependencies=("open_clip", "torch", "PIL"),
model_id=OpenClipModelAdapter.build_model_id(
self.model_name, self.pretrained
),
requested_num_gpus=self._num_gpus,
requested_gpu_type=self._gpu_type,
)
self._model_handle.register_operation(
"image_fixed_text_similarity",
_ImageFixedTextSimilarity._predict_batch,
)
def close(self):
if getattr(self, "_model_handle", None) is not None:
self._model_handle.release()
self._model_handle = None
def eval(self, image_series: "pd.Series") -> "pd.Series":
return run_image_batch_inference(
image_series,
lambda image_arrays: self._model_handle.call(
"image_fixed_text_similarity", image_arrays, self.text
),
)
class _ImageTextSimilarity(ScalarFunction):
"""Compute per-row image-text similarity scores."""
def __init__(self, model="ViT-B/32", pretrained="openai", model_sharing=None,
num_gpus=None, gpu_type=None):
super().__init__()
self.model_name = model
self.pretrained = pretrained
self.model_sharing = model_sharing
self._num_gpus = num_gpus
self._gpu_type = gpu_type
@staticmethod
def _predict_batch(model, pixel_arrays, texts):
image_features = model.encode_images(pixel_arrays)
text_features = model.encode_texts(texts)
# _OpenClipModel encode helpers return L2-normalized vectors,
# so row-wise dot product == cosine similarity for paired inputs.
scores = (image_features * text_features).sum(dim=-1)
return scores.cpu().tolist()
def open(self, function_context):
import pandas as pd
# Per-row text mode filters null text after image decoding, so it cannot
# use run_image_batch_inference directly.
self._pd = pd
self._model_handle = prepare_and_load_model_handle(
adapter_cls=OpenClipModelAdapter,
config={},
function_context=function_context,
model_sharing=self.model_sharing,
dependencies=("open_clip", "torch", "PIL"),
model_id=OpenClipModelAdapter.build_model_id(
self.model_name, self.pretrained
),
requested_num_gpus=self._num_gpus,
requested_gpu_type=self._gpu_type,
)
self._model_handle.register_operation(
"image_text_similarity",
_ImageTextSimilarity._predict_batch,
)
def close(self):
if getattr(self, "_model_handle", None) is not None:
self._model_handle.release()
self._model_handle = None
self._pd = None
def eval(
self, image_series: "pd.Series", text_series: "pd.Series"
) -> "pd.Series":
pd = self._pd
pixel_arrays, valid_idx = decode_image_batch(image_series, mode="RGB")
if not pixel_arrays:
return pd.Series([None] * len(image_series))
filtered_pixel_arrays = []
filtered_idx = []
texts = []
for pixel_array, idx in zip(pixel_arrays, valid_idx):
text = text_series.iloc[idx]
if pd.isna(text):
continue
filtered_pixel_arrays.append(pixel_array)
filtered_idx.append(idx)
texts.append(text)
if not filtered_pixel_arrays:
return pd.Series([None] * len(image_series))
valid_results = self._model_handle.call(
"image_text_similarity", filtered_pixel_arrays, texts
)
return pd.Series(scatter_results(valid_results, filtered_idx, len(image_series)))
[docs]def image_text_similarity(
*columns,
text=None,
model="ViT-B/32",
pretrained="openai",
model_sharing=None,
concurrency=None,
batch_size=None,
num_gpus=None,
gpu_type=None,
):
"""
Create an image-text similarity scoring UDF (CLIP / open_clip-based).
When ``text`` is provided, computes similarity of all images against that
fixed text prompt (single-column UDF). When ``text`` is ``None``, computes
per-row similarity between an image column and a text column (two-column UDF).
Requires ``pip install open_clip_torch torch Pillow``. This is a pandas
batch UDF that supports GPU acceleration via ``num_gpus`` / ``gpu_type``.
Args:
*columns: Optional image column(s). When provided, the UDF is
applied directly instead of returning a factory.
text: Fixed text prompt to compare against all images. ``None``
(default) enables per-row mode where a text column is passed
at call time.
model: CLIP model architecture name. Default ``"ViT-B/32"``.
pretrained: Pretrained weights checkpoint. Default ``"openai"``.
model_sharing: Model sharing mode across parallel subtasks.
``None`` uses per-process caching.
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 that returns a cosine similarity score in ``[-1, 1]``, or
``None`` for null image inputs. Per-row mode also returns ``None`` for
null text inputs.
Example::
>>> # As a reusable variable (fixed text mode)
>>> sim = image_text_similarity(text="a photo of a cat")
>>> df = df.with_column("score", sim(col("img")))
>>>
>>> # As a reusable variable (per-row text mode)
>>> sim = image_text_similarity()
>>> df = df.with_column("score", sim(col("img"), col("text")))
>>>
>>> # Inline
>>> df = df.with_column(
... "score", image_text_similarity(col("img"), col("text"))
... )
"""
runtime_kwargs = _udf_runtime_kwargs(
concurrency=concurrency,
batch_size=batch_size,
num_gpus=num_gpus,
gpu_type=gpu_type,
)
if text is not None:
wrapper = udf(
_ImageFixedTextSimilarity(
text=text,
model=model,
pretrained=pretrained,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
func_type="pandas",
return_dtype=DataType.float64(),
**runtime_kwargs,
)
else:
wrapper = udf(
_ImageTextSimilarity(
model=model,
pretrained=pretrained,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
func_type="pandas",
return_dtype=DataType.float64(),
**runtime_kwargs,
)
return _build_or_apply_udf(wrapper, *columns)