################################################################################
# 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.
################################################################################
"""
Detection and recognition operators.
Operators in this module detect or recognize structured content within
images: objects, text, segments, and subplots.
Model-based operators (``ImageDetection``, ``ImageSegment``, ``ImageOCR``)
are registered as pandas batch UDFs, so Flink batches rows automatically.
``image_detect_objects`` returns axis-aligned boxes ``x/y/w/h``;
``image_ocr`` returns four-point text polygons; ``image_segment`` returns
indexed PNG mask bytes; ``image_detect_subplot`` returns subplot metadata.
Example::
>>> from pyflink.multimodal.operators.image_detect import (
... image_detect_objects, image_ocr, image_segment, image_detect_subplot,
... )
>>> from pyflink.dataframe import col
>>> detect = image_detect_objects(model="yolov8n", confidence=0.05)
>>> df = df.with_column("objects", detect(col("img")))
>>> ocr = image_ocr(lang=["en", "ch_sim"])
>>> df = df.with_column("text", ocr(col("img")))
Runtime args:
Model-based pandas factories forward keyword-only ``concurrency``,
``batch_size``, ``num_gpus``, and ``gpu_type`` to the DataFrame UDF
runtime. ``image_detect_subplot`` forwards ``concurrency``.
"""
import io
import warnings
from typing import TYPE_CHECKING
from pyflink.dataframe import udf, DataType
from pyflink.table.udf import ScalarFunction
from pyflink.multimodal.codec import decode_image_input
from pyflink.model.backends.easyocr import EasyOcrModelAdapter
from pyflink.model.backends.ultralytics import UltralyticsModelAdapter
from pyflink.model.cache_manager import (
_is_gpu_requested,
check_dependencies,
prepare_and_load_model_handle,
)
from pyflink.multimodal.utils import (
_build_or_apply_udf,
_setup_cv2_threads,
_udf_runtime_kwargs,
run_image_batch_inference,
)
__all__ = [
"image_detect_objects",
"image_segment",
"image_ocr",
"image_detect_subplot",
]
if TYPE_CHECKING:
import pandas as pd
_MASK_ENCODING_MODULES = None
def _get_mask_encoding_modules():
global _MASK_ENCODING_MODULES
if _MASK_ENCODING_MODULES is None:
from pyflink.multimodal.codec import _PILImage
import numpy as np
_MASK_ENCODING_MODULES = (np, _PILImage)
return _MASK_ENCODING_MODULES
# ImageDetection - YOLO object detection (pandas batch UDF)
class _ImageDetection(ScalarFunction):
"""Detect objects in a batch of images using YOLO."""
def __init__(self, model="yolov8n", confidence=0.05, imgsz=640, iou=0.5,
model_sharing=None, num_gpus=None, gpu_type=None):
super().__init__()
if not 0 <= confidence <= 1:
raise ValueError(f"confidence must be in [0, 1], got {confidence!r}")
if not isinstance(imgsz, int) or imgsz <= 0:
raise ValueError(f"imgsz must be a positive integer, got {imgsz!r}")
if not 0 <= iou <= 1:
raise ValueError(f"iou must be in [0, 1], got {iou!r}")
self.model_name = model
# Aligned with Data-Juicer ImageDetectionYoloMapper (conf=0.05).
# Low default favors recall in data-cleaning scenarios; users can raise
# the threshold downstream when precision matters more.
self.confidence = confidence
# imgsz: controls inference resolution. Larger values improve small-object
# detection at the cost of speed and memory. Default 640 aligned with
# ultralytics and Data-Juicer.
self.imgsz = imgsz
# iou: NMS IoU threshold. Overlapping boxes above this ratio are suppressed.
# Default 0.5 aligned with Data-Juicer (ultralytics default is 0.7).
self.iou = iou
self.model_sharing = model_sharing
self._num_gpus = num_gpus
self._gpu_type = gpu_type
@staticmethod
def _parse_result(r):
boxes = r.boxes
if boxes is None or len(boxes) == 0:
return []
xyxy = boxes.xyxy.cpu().numpy()
cls_ids = boxes.cls.cpu().numpy().astype(int)
confs = boxes.conf.cpu().numpy()
return [
{
"label": r.names.get(c, f"class_{c}"),
"x": float(xyxy[i, 0]),
"y": float(xyxy[i, 1]),
"w": float(xyxy[i, 2] - xyxy[i, 0]),
"h": float(xyxy[i, 3] - xyxy[i, 1]),
"confidence": float(confs[i]),
}
for i, c in enumerate(cls_ids)
]
@staticmethod
def _predict_batch(model, images, confidence, imgsz, iou):
results = model(images, conf=confidence, imgsz=imgsz, iou=iou, verbose=False)
return [_ImageDetection._parse_result(r) for r in results]
def open(self, function_context):
self._model_handle = prepare_and_load_model_handle(
adapter_cls=UltralyticsModelAdapter,
config={},
function_context=function_context,
model_sharing=self.model_sharing,
dependencies=("ultralytics",),
model_id=self.model_name,
requested_num_gpus=self._num_gpus,
requested_gpu_type=self._gpu_type,
)
self._model_handle.register_operation(
"detect_objects", _ImageDetection._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":
# YOLO model inputs are normalized to RGB at the operator boundary;
# decoded images may carry L/LA/RGBA modes for non-model operators.
return run_image_batch_inference(
image_series,
lambda image_arrays: self._model_handle.call(
"detect_objects", image_arrays, self.confidence,
self.imgsz, self.iou
),
)
[docs]def image_detect_objects(*columns, model="yolov8n", confidence=0.05,
imgsz=640, iou=0.5, model_sharing=None,
concurrency=None, batch_size=None,
num_gpus=None, gpu_type=None):
"""
Create an object detection UDF (YOLO-based).
Requires ``pip install ultralytics``. 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: YOLO model name, e.g. ``"yolov8n"``, ``"yolov8s"``.
Default ``"yolov8n"``.
confidence: Minimum confidence threshold for detections.
Default ``0.05`` (aligned with Data-Juicer).
imgsz: Inference resolution in pixels. Larger values improve
small-object detection at the cost of speed and memory.
Default ``640``.
iou: NMS IoU threshold. Overlapping boxes above this ratio are
suppressed. Default ``0.5``.
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 list of dicts with ``label``, ``x``, ``y``,
``w``, ``h``, ``confidence`` keys.
Example::
>>> # As a reusable variable
>>> detect = image_detect_objects(model="yolov8n", confidence=0.05)
>>> df = df.with_column("objects", detect(col("img")))
>>>
>>> # Inline
>>> df = df.with_column("objects", image_detect_objects(col("img")))
"""
wrapper = udf(
_ImageDetection(
model=model,
confidence=confidence,
imgsz=imgsz,
iou=iou,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
func_type="pandas",
return_dtype=DataType.list(
DataType.struct(
{
"label": DataType.string(),
"x": DataType.float64(),
"y": DataType.float64(),
"w": DataType.float64(),
"h": DataType.float64(),
"confidence": DataType.float64(),
}
)
),
**_udf_runtime_kwargs(
concurrency=concurrency,
batch_size=batch_size,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageSegment - SAM segmentation mask (pandas batch UDF)
class _ImageSegment(ScalarFunction):
"""Segment images into indexed masks using FastSAM."""
def __init__(self, model="FastSAM-x.pt", confidence=0.05, imgsz=1024, iou=0.5,
model_sharing=None, num_gpus=None, gpu_type=None):
super().__init__()
if not 0 <= confidence <= 1:
raise ValueError(f"confidence must be in [0, 1], got {confidence!r}")
if not isinstance(imgsz, int) or imgsz <= 0:
raise ValueError(f"imgsz must be a positive integer, got {imgsz!r}")
if not 0 <= iou <= 1:
raise ValueError(f"iou must be in [0, 1], got {iou!r}")
self.model_name = model
# FastSAM commonly uses 1024px inference size; larger values preserve
# small-object mask detail at higher CPU/GPU cost.
self.imgsz = imgsz
self.confidence = confidence
self.iou = iou
self.model_sharing = model_sharing
self._num_gpus = num_gpus
self._gpu_type = gpu_type
@staticmethod
def _encode_mask(pixel_array, result):
np, pil_image = _get_mask_encoding_modules()
combined = np.zeros(pixel_array.shape[:2], dtype=np.uint8)
resampling = getattr(pil_image, "Resampling", pil_image)
if result.masks is not None:
for i, mask in enumerate(result.masks.data):
if i == 255:
warnings.warn(
"Segmentation mask has more than 255 segments; labels "
"after 255 are saturated to 255 in the uint8 mask.",
UserWarning,
stacklevel=2,
)
# Keep processing overflow masks so label 255 represents
# their union instead of dropping later segments.
mask_np = mask.cpu().numpy().astype(bool)
if mask_np.shape != combined.shape:
mask_pil = pil_image.fromarray(mask_np.astype(np.uint8) * 255)
mask_pil = mask_pil.resize(
(combined.shape[1], combined.shape[0]),
resampling.NEAREST,
)
mask_np = np.array(mask_pil) > 127
combined[mask_np] = min(i + 1, 255)
# Keep the label mask lossless and portable. Returning encoded PNG bytes
# avoids exposing a decoded image value that users would still need to
# encode before storing, visualizing, or passing to non-mask consumers.
mask_image = pil_image.fromarray(combined, mode="L")
with io.BytesIO() as buf:
mask_image.save(buf, format="PNG")
return buf.getvalue()
@staticmethod
def _predict_batch(model, images, confidence, imgsz, iou):
results = model(
images,
imgsz=imgsz,
conf=confidence,
iou=iou,
verbose=False,
)
return [
_ImageSegment._encode_mask(pixel_array, result)
for pixel_array, result in zip(images, results)
]
def open(self, function_context):
self._model_handle = prepare_and_load_model_handle(
adapter_cls=UltralyticsModelAdapter,
config={},
function_context=function_context,
model_sharing=self.model_sharing,
dependencies=("ultralytics",),
model_id=self.model_name,
requested_num_gpus=self._num_gpus,
requested_gpu_type=self._gpu_type,
)
self._model_handle.register_operation(
"segment_masks", _ImageSegment._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(
"segment_masks",
image_arrays, self.confidence, self.imgsz, self.iou
),
)
[docs]def image_segment(*columns, model="FastSAM-x.pt", confidence=0.05, imgsz=1024, iou=0.5,
model_sharing=None, concurrency=None, batch_size=None,
num_gpus=None, gpu_type=None):
"""
Create a semantic segmentation UDF (FastSAM-based).
Requires ``pip install ultralytics``. 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: FastSAM model name. Default ``"FastSAM-x.pt"``.
confidence: Minimum confidence threshold. Default ``0.05``.
imgsz: Input image size for inference. Default ``1024``.
iou: IoU threshold for NMS. Default ``0.5``.
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 an 8-bit PNG indexed segmentation mask. Pixel value
``0`` is background; values ``1..255`` are segment indices. Later masks
overwrite earlier masks in overlapping pixels. Segment indices above
255 are saturated to 255 because the mask is encoded as ``uint8``.
Example::
>>> # As a reusable variable
>>> segment = image_segment(model="FastSAM-x.pt")
>>> df = df.with_column("mask", segment(col("img")))
>>>
>>> # Inline
>>> df = df.with_column("mask", image_segment(col("img")))
"""
wrapper = udf(
_ImageSegment(
model=model, confidence=confidence, imgsz=imgsz, iou=iou,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
func_type="pandas",
return_dtype=DataType.binary(),
**_udf_runtime_kwargs(
concurrency=concurrency,
batch_size=batch_size,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageOCR - EasyOCR text extraction (pandas batch UDF)
class _ImageOCR(ScalarFunction):
"""Extract text from a batch of images using EasyOCR."""
def __init__(self, lang=None, num_gpus=None, gpu_type=None, model_sharing=None):
super().__init__()
if lang is None:
lang = ["en"]
elif (
isinstance(lang, str)
or not isinstance(lang, (list, tuple))
or not lang
or not all(isinstance(code, str) and code for code in lang)
):
raise ValueError(
"lang must be a non-empty list or tuple of language code strings"
)
self.lang = list(lang)
self._num_gpus = num_gpus
self._gpu_type = gpu_type
self.model_sharing = model_sharing
@staticmethod
def _parse_result(results):
# EasyOCR readtext() returns (bbox, text, confidence) where bbox is
# four corner points [[x1,y1],[x2,y2],[x3,y3],[x4,y4]], supporting
# rotated text regions. We preserve bbox to enable downstream
# position-based filtering (e.g., watermark detection, subtitle removal).
return [
{"text": str(text), "confidence": float(conf),
"bbox": [[float(c) for c in pt] for pt in bbox]}
for bbox, text, conf in results
]
@staticmethod
def _predict_batch(model, images):
# Use per-image readtext to support mixed-size batches; grouping by
# size can reintroduce readtext_batched later as a targeted
# optimization.
return [
_ImageOCR._parse_result(model.readtext(pixel_array))
for pixel_array in images
]
def open(self, function_context):
# EasyOCR accepts a boolean GPU switch rather than an explicit device.
# ``num_gpus`` is still forwarded as the Flink resource request, and
# ``gpu_type`` is validated by cache_manager before the model loads.
use_gpu = _is_gpu_requested(self._num_gpus)
self._model_handle = prepare_and_load_model_handle(
adapter_cls=EasyOcrModelAdapter,
config={"lang": list(self.lang), "gpu": use_gpu},
function_context=function_context,
model_sharing=self.model_sharing,
dependencies=("easyocr",),
model_id="easyocr",
requested_num_gpus=self._num_gpus,
requested_gpu_type=self._gpu_type,
)
self._model_handle.register_operation(
"read_text", _ImageOCR._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":
# EasyOCR expects image-like RGB arrays; normalize supported decoded
# image modes before crossing the model boundary.
return run_image_batch_inference(
image_series,
lambda image_arrays: self._model_handle.call("read_text", image_arrays),
)
[docs]def image_ocr(*columns, lang=None, model_sharing=None, concurrency=None,
batch_size=None, num_gpus=None, gpu_type=None):
"""
Create an OCR text extraction UDF (EasyOCR-based).
Requires ``pip install easyocr``. This is a pandas batch UDF. EasyOCR
accepts only a CPU/GPU boolean switch here: ``num_gpus > 0`` requests a
Flink GPU resource and enables EasyOCR GPU mode, while ``gpu_type`` is
used to validate the assigned Flink GPU label before loading. EasyOCR does
not expose explicit CUDA device placement through this operator. Inference
currently runs per image so mixed-size batches do not need padding;
``batch_size`` mainly controls Python/Arrow batching overhead.
Args:
*columns: Optional image column(s). When provided, the UDF is
applied directly instead of returning a factory.
lang: List of language codes, e.g. ``["en"]``, ``["en", "ch_sim"]``.
``None`` (default) uses ``["en"]``.
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``. Values
greater than ``0`` enable EasyOCR GPU mode; ``None`` and ``0``
run on CPU.
gpu_type: Required Flink GPU resource label, e.g. ``"A10"``. ``None``
accepts any assigned GPU label. EasyOCR itself receives only the
boolean GPU switch.
Returns:
A UDF that returns a list of dicts with ``text``, ``confidence``,
``bbox`` keys. ``bbox`` is a four-point polygon
``[[x1, y1], ...]`` in image coordinates, preserving rotated text
regions returned by EasyOCR.
Example::
>>> # As a reusable variable
>>> ocr = image_ocr(lang=["en", "ch_sim"])
>>> df = df.with_column("text", ocr(col("img")))
>>>
>>> # Inline
>>> df = df.with_column("text", image_ocr(col("img")))
"""
wrapper = udf(
_ImageOCR(
lang=lang,
num_gpus=num_gpus,
gpu_type=gpu_type,
model_sharing=model_sharing,
),
func_type="pandas",
return_dtype=DataType.list(
DataType.struct(
{
"text": DataType.string(),
"confidence": DataType.float64(),
"bbox": DataType.list(DataType.list(DataType.float64())),
}
)
),
**_udf_runtime_kwargs(
concurrency=concurrency,
batch_size=batch_size,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageSubplot - detect composite / mosaic images
class _ImageSubplot(ScalarFunction):
"""Detect subplot / mosaic composites using OpenCV Hough lines."""
def __init__(self, threshold=0.5):
super().__init__()
if not 0 <= threshold <= 1:
raise ValueError(f"threshold must be in [0, 1], got {threshold!r}")
self.threshold = threshold
@staticmethod
def _count_unique_lines(positions, image_extent):
if not positions:
return 0
# Merge nearby Hough detections into one logical separator; broken or
# anti-aliased grid lines often produce several close line segments.
tolerance = max(1.0, image_extent * 0.05)
count = 0
cluster_end = None
for position in sorted(positions):
if cluster_end is None or position - cluster_end > tolerance:
count += 1
cluster_end = position
else:
cluster_end = position
return count
def open(self, function_context):
check_dependencies("cv2")
import cv2
import numpy as np
_setup_cv2_threads(cv2)
self._cv2 = cv2
self._np = np
def eval(self, image_input):
if image_input is None:
return None
np = self._np
# mode="L" gives a 2-D grayscale array for edge detection.
gray_array = decode_image_input(image_input, mode="L")
h, w = gray_array.shape
# Hough transform needs sufficient pixels for reliable line voting
# (threshold=100 requires at least 100 edge pixels along a line).
# Images below 20px cannot produce meaningful subplot detection.
if min(h, w) < 20:
return (False, 1)
# Map public sensitivity [0, 1] to Canny low threshold [50, 200];
# the high threshold is 2x, following the common Canny hysteresis ratio.
canny_thresh = int(50 + self.threshold * 150)
edges = self._cv2.Canny(gray_array, canny_thresh, canny_thresh * 2)
# Require grid separators to span at least half the shorter edge. This
# avoids treating short texture edges as subplot separators.
min_length = min(h, w) * 0.5
# Use a conservative 100-vote Hough threshold to filter weak lines.
# The 10px gap tolerance keeps compressed or anti-aliased grid lines
# connected without merging distant separators.
lines = self._cv2.HoughLinesP(
edges, 1, np.pi / 180, threshold=100,
minLineLength=int(min_length), maxLineGap=10
)
if lines is None:
return (False, 1)
h_positions = []
v_positions = []
for line in lines:
x1, y1, x2, y2 = line[0]
angle = abs(np.arctan2(y2 - y1, x2 - x1) * 180 / np.pi)
# Count near-horizontal/vertical internal separators; ignore the
# outer 10% to avoid treating image borders as subplot dividers.
if angle < 5 or angle > 175:
y_mid = (y1 + y2) / 2
if h * 0.1 < y_mid < h * 0.9:
h_positions.append(y_mid)
elif 85 < angle < 95:
x_mid = (x1 + x2) / 2
if w * 0.1 < x_mid < w * 0.9:
v_positions.append(x_mid)
h_lines = self._count_unique_lines(h_positions, h)
v_lines = self._count_unique_lines(v_positions, w)
# One internal separator can represent a 1x2 or 2x1 composite.
is_subplot = h_lines >= 1 or v_lines >= 1
count = (h_lines + 1) * (v_lines + 1) if is_subplot else 1
return (is_subplot, count)
[docs]def image_detect_subplot(*columns, threshold=0.5, concurrency=None):
"""
Create a subplot / mosaic detection UDF (OpenCV Hough-based).
Requires ``pip install opencv-python-headless``. This is a scalar UDF.
Detects whether an image is a composite (mosaic / subplot / screenshot
collage) by looking for strong horizontal/vertical grid lines. Images
smaller than 20px on either axis always return ``(False, 1)`` because
the Hough line detector cannot accumulate enough votes at that resolution.
Corrupt or undecodable images raise; use ``is_valid_image`` first when
invalid inputs should be filtered instead of failing the task.
Args:
*columns: Optional image column(s). When provided, the UDF is
applied directly instead of returning a factory.
threshold: Edge detection sensitivity in ``[0, 1]``. Higher values
require stronger edges. Default ``0.5``.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF that returns an object with ``is_subplot`` (bool) and
``count`` (int).
Example::
>>> # As a reusable variable
>>> subplot = image_detect_subplot(threshold=0.5)
>>> df = df.with_column("subplot", subplot(col("img")))
>>>
>>> # Inline
>>> df = df.with_column("subplot", image_detect_subplot(col("img")))
"""
wrapper = udf(
_ImageSubplot(threshold=threshold),
return_dtype=DataType.struct(
{
"is_subplot": DataType.boolean(),
"count": DataType.int32(),
}
),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)