################################################################################
# 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 quality and safety assessment operators.
Operators in this module score images based on quality, safety, and
watermark detection. All operators return float scores.
Example::
>>> from pyflink.multimodal.operators.image_quality import image_quality_score
>>> quality = image_quality_score()
>>> df = df.with_column("quality", quality(col("img")))
>>> from pyflink.multimodal.operators.image_quality import image_nsfw_score
>>> nsfw = image_nsfw_score()
>>> df = df.with_column("nsfw", nsfw(col("img")))
Runtime args:
Pandas batch UDFs accept ``concurrency``, ``batch_size``, ``num_gpus``,
and ``gpu_type``. Row-level UDFs accept ``concurrency``.
Input errors:
Null image inputs produce ``None`` outputs. Corrupt encoded bytes are not
suppressed; decode errors propagate to the caller. Use
``decode_image(on_error=...)`` or ``is_valid_image`` before these operators
when processing untrusted raw bytes.
"""
from typing import TYPE_CHECKING
import logging
from pyflink.dataframe import udf, DataType
from pyflink.table.udf import ScalarFunction
from pyflink.multimodal.codec import decode_image_input
from pyflink.model.backends.huggingface import (
AestheticScorerAdapter,
DEFAULT_AESTHETIC_SCORER_MODEL,
HfImageClassifierAdapter,
)
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,
run_image_batch_inference,
)
_DEFAULT_NSFW_MODEL = "Falconsai/nsfw_image_detection"
_DEFAULT_WATERMARK_MODEL = "amrul-hzz/watermark_detector"
_SCORE_IMAGES_OP = "score_images"
__all__ = [
"image_nsfw_score",
"image_aesthetic_score",
"image_quality_score",
"image_watermark_score",
]
if TYPE_CHECKING:
import pandas as pd
# Uncalibrated heuristic divisors that bound raw quality metrics to [0, 1].
# Calibrate thresholds on the user's image distribution.
_SHARPNESS_NORM = 1000.0 # Laplacian variance divisor
_CONTRAST_NORM = 80.0 # Grayscale std divisor
_COLORFULNESS_NORM = 64.0 # Channel std divisor
_LUMA_WEIGHTS = (0.299, 0.587, 0.114) # ITU-R 601 luma coefficients
_logger = logging.getLogger(__name__)
_NEGATIVE_PREFIX_TOKENS = (
"not",
"no",
"non",
"without",
"negative",
)
_NEGATIVE_SUFFIX_TOKENS = (
"free",
"absent",
"clean",
"safe",
"sfw",
)
# HfImageClassifierScorer - shared base for NSFW / watermark scoring
class _HfImageClassifierScorer(ScalarFunction):
"""Shared HuggingFace image-classification scorer base class."""
_label_keyword = ""
_missing_label_hint = ""
_missing_label_article = "a"
def __init__(self, model_id, model_sharing=None, num_gpus=None, gpu_type=None):
super().__init__()
self._model_id = model_id
self.model_sharing = model_sharing
self._num_gpus = num_gpus
self._gpu_type = gpu_type
@classmethod
def _is_negative_label(cls, label):
normalized = label.lower().replace("-", "_").replace("/", "_").strip()
tokens = [token for token in normalized.replace("_", " ").split() if token]
for index, token in enumerate(tokens):
if token != cls._label_keyword:
continue
# Only local negation counts. A suffix must be terminal
# ("nsfw_free") so labels like "nsfw_safe_level" stay positive.
previous_token = tokens[index - 1] if index > 0 else None
next_token = tokens[index + 1] if index + 1 < len(tokens) else None
if previous_token in _NEGATIVE_PREFIX_TOKENS:
return True
is_terminal_suffix = index + 1 == len(tokens) - 1
if is_terminal_suffix and next_token in _NEGATIVE_SUFFIX_TOKENS:
return True
return False
@staticmethod
def _label_index(idx, label, model_id):
try:
return int(idx)
except (TypeError, ValueError) as exc:
raise ValueError(
f"Model '{model_id}' label index {idx!r} for "
f"'{label}' must be an integer."
) from exc
@classmethod
def _is_fallback_positive_label(cls, normalized):
if cls._label_keyword not in normalized:
return False
return not cls._is_negative_label(normalized)
@classmethod
def _resolve_positive_label_indices(cls, labels, model_id):
exact_matches = []
fallback_matches = []
for idx, label in labels.items():
normalized = str(label).lower().replace("-", "_").strip()
if normalized == cls._label_keyword:
label_idx = cls._label_index(idx, label, model_id)
exact_matches.append((label_idx, label))
elif cls._is_fallback_positive_label(normalized):
label_idx = cls._label_index(idx, label, model_id)
fallback_matches.append((label_idx, label))
matches = exact_matches or fallback_matches
if not matches:
raise ValueError(
f"Model '{model_id}' does not have "
f"{cls._missing_label_article} "
f"'{cls._label_keyword}' label. "
f"Available labels: {labels}. "
f"Please use a model with {cls._missing_label_hint} capability."
)
if len(matches) > 1:
_logger.warning(
"Model %r has multiple labels matching %r; aggregating labels: %s",
model_id,
cls._label_keyword,
[label for _, label in matches],
)
elif not exact_matches:
_logger.warning(
"Model %r does not expose an exact %r label; using fallback "
"label match: %r",
model_id,
cls._label_keyword,
matches[0][1],
)
return [idx for idx, _ in matches]
@staticmethod
def _predict_batch(model, pixel_arrays, label_indices, multi_label):
probabilities = model.predict_probabilities(pixel_arrays)
return [
_HfImageClassifierScorer._score_positive_labels(
prob, label_indices, multi_label
)
for prob in probabilities
]
@staticmethod
def _score_positive_labels(prob, label_indices, multi_label):
if not label_indices:
return 0.0
prob_len = len(prob)
for idx in label_indices:
if idx < 0 or idx >= prob_len:
raise ValueError(
f"Positive label index {idx} is out of range for "
f"{prob_len} classifier labels."
)
scores = [float(prob[idx]) for idx in label_indices]
score = max(scores) if multi_label else sum(scores)
return max(0.0, min(1.0, float(score)))
def open(self, function_context):
self._model_handle = prepare_and_load_model_handle(
adapter_cls=HfImageClassifierAdapter,
config={},
function_context=function_context,
model_sharing=self.model_sharing,
dependencies=("transformers", "torch"),
model_id=self._model_id,
requested_num_gpus=self._num_gpus,
requested_gpu_type=self._gpu_type,
)
self._model_handle.register_operation(
_SCORE_IMAGES_OP, _HfImageClassifierScorer._predict_batch
)
metadata = self._model_handle.metadata()
labels = metadata.get("labels", {})
label_indices = self._resolve_positive_label_indices(labels, self._model_id)
self._label_indices = label_indices
self._multi_label = (
metadata.get("problem_type") == "multi_label_classification"
)
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 pixel_arrays: self._model_handle.call(
_SCORE_IMAGES_OP,
pixel_arrays,
self._label_indices,
self._multi_label,
),
)
class _ImageNsfwScore(_HfImageClassifierScorer):
"""NSFW content scoring using a HuggingFace image classifier."""
_label_keyword = "nsfw"
_missing_label_hint = "NSFW detection"
_missing_label_article = "an"
def __init__(self, hf_nsfw_model=_DEFAULT_NSFW_MODEL,
model_sharing=None, num_gpus=None, gpu_type=None):
self.hf_nsfw_model = hf_nsfw_model
super().__init__(
hf_nsfw_model,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
)
[docs]def image_nsfw_score(
*columns,
hf_nsfw_model=_DEFAULT_NSFW_MODEL,
model_sharing=None,
concurrency=None,
batch_size=None,
num_gpus=None,
gpu_type=None,
):
"""
Create an NSFW scoring UDF (HuggingFace image classification-based).
Requires ``pip install transformers torch``. This is a pandas batch UDF
that supports GPU acceleration via ``num_gpus`` / ``gpu_type``.
Set ``num_gpus`` to request GPU resources from the DataFrame UDF runtime.
Args:
*columns: Optional image column(s). When provided, the UDF is
applied directly instead of returning a factory.
hf_nsfw_model: HuggingFace model ID for NSFW classification. Defaults
to the built-in NSFW classifier model.
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 float NSFW probability score in ``[0, 1]`` or
``None`` for null inputs. Higher values indicate stronger NSFW
confidence. Use downstream filtering (e.g. ``> 0.5``) to convert to a
boolean decision.
Example::
>>> # As a reusable variable
>>> nsfw = image_nsfw_score()
>>> df = df.with_column("nsfw", nsfw(col("img")))
>>> nsfw_gpu = image_nsfw_score(num_gpus=1.0)
>>> df = df.with_column("nsfw", nsfw_gpu(col("img")))
>>>
>>> # Inline
>>> df = df.with_column("nsfw", image_nsfw_score(col("img")))
>>> df = df.filter(image_nsfw_score(col("img")) > 0.5)
"""
wrapper = udf(
_ImageNsfwScore(
hf_nsfw_model=hf_nsfw_model,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
func_type="pandas",
return_dtype=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)
# ImageAestheticScore - LAION aesthetic predictor (pandas batch UDF)
class _ImageAestheticScore(ScalarFunction):
"""Aesthetic quality scoring using a HuggingFace aesthetics predictor."""
def __init__(
self,
hf_scorer_model=DEFAULT_AESTHETIC_SCORER_MODEL,
model_sharing=None,
num_gpus=None,
gpu_type=None,
):
super().__init__()
self.hf_scorer_model = hf_scorer_model
self.model_sharing = model_sharing
self._num_gpus = num_gpus
self._gpu_type = gpu_type
@staticmethod
def _predict_batch(model, pixel_arrays):
return model.predict_scores(pixel_arrays)
def open(self, function_context):
self._model_handle = prepare_and_load_model_handle(
adapter_cls=AestheticScorerAdapter,
config={},
function_context=function_context,
model_sharing=self.model_sharing,
dependencies=("transformers", "torch"),
model_id=self.hf_scorer_model,
requested_num_gpus=self._num_gpus,
requested_gpu_type=self._gpu_type,
)
self._model_handle.register_operation(
_SCORE_IMAGES_OP, _ImageAestheticScore._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(
_SCORE_IMAGES_OP, image_arrays
),
)
[docs]def image_aesthetic_score(
*columns,
hf_scorer_model=DEFAULT_AESTHETIC_SCORER_MODEL,
model_sharing=None,
concurrency=None,
batch_size=None,
num_gpus=None,
gpu_type=None,
):
"""
Create an aesthetic quality scoring UDF (HuggingFace aesthetics predictor-based).
Requires ``pip install transformers torch``. This is a pandas batch UDF
that supports GPU acceleration via ``num_gpus`` / ``gpu_type``.
Set ``num_gpus`` to request GPU resources from the DataFrame UDF runtime.
The default scorer is a regression head, not a probability model. Its
score is divided by 10 and clamped to ``[0, 1]`` because the default LAION
aesthetic predictor is trained on AVA/SAC-style 1-10 targets. Custom scorer
outputs are expected to already use a ``[0, 1]`` scale and are clamped
defensively. Custom model IDs must provide local safetensors head weights
matching the built-in head architecture
``projection_dim -> 1024 -> 128 -> 64 -> 16 -> 1``; only the built-in
default model may fall back to ``trust_remote_code=True`` when those weights
are unavailable.
Args:
*columns: Optional image column(s). When provided, the UDF is
applied directly instead of returning a factory.
hf_scorer_model: HuggingFace model ID for aesthetic scoring. Default
``"shunk031/aesthetics-predictor-v2-sac-logos-ava1-l14-linearMSE"``.
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 normalized float aesthetic quality score in
``[0, 1]`` or ``None`` for null inputs.
Example::
>>> # As a reusable variable
>>> aesthetic = image_aesthetic_score()
>>> df = df.with_column("aesthetic", aesthetic(col("img")))
>>> aesthetic_gpu = image_aesthetic_score(num_gpus=1.0)
>>> df = df.with_column("aesthetic", aesthetic_gpu(col("img")))
>>>
>>> # Inline
>>> df = df.with_column("aesthetic", image_aesthetic_score(col("img")))
"""
wrapper = udf(
_ImageAestheticScore(hf_scorer_model=hf_scorer_model,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type),
func_type="pandas",
return_dtype=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)
# ImageQualityScore - heuristic no-reference quality assessment
class _ImageQualityScore(ScalarFunction):
"""Image quality scoring using Laplacian variance, contrast, and colorfulness."""
def open(self, function_context):
check_dependencies("numpy")
import numpy as np
self._np = np
def close(self):
self._np = None
def eval(self, image_input):
pixel_array = decode_image_input(image_input, mode="RGB")
if pixel_array is None:
return None
np = self._np
pixel_array = pixel_array.astype(np.float32)
with np.errstate(invalid="ignore", over="ignore"):
gray = np.dot(pixel_array[:, :, :3], _LUMA_WEIGHTS)
# Sharpness: Laplacian variance (higher = sharper)
if gray.shape[0] >= 3 and gray.shape[1] >= 3:
laplacian = gray[:-2, 1:-1] + gray[2:, 1:-1]
laplacian += gray[1:-1, :-2]
laplacian += gray[1:-1, 2:]
laplacian -= 4 * gray[1:-1, 1:-1]
sharpness = min(laplacian.var() / _SHARPNESS_NORM, 1.0)
else:
sharpness = 0.0
# Contrast: standard deviation of grayscale (higher = more contrast)
contrast = min(gray.std() / _CONTRAST_NORM, 1.0)
# Colorfulness: mean std of color channels
color_std = np.mean([pixel_array[:, :, c].std() for c in range(3)])
colorfulness = min(color_std / _COLORFULNESS_NORM, 1.0)
# Heuristic bounded score; tune thresholds on your dataset.
score = 0.5 * sharpness + 0.3 * contrast + 0.2 * colorfulness
return float(np.nan_to_num(score, nan=0.0, posinf=1.0, neginf=0.0))
[docs]def image_quality_score(*columns, concurrency=None):
"""
Create an image quality scoring UDF (Laplacian + contrast + colorfulness-based).
Requires ``pip install numpy``. This is a scalar UDF that processes one
image at a time. It supports direct call syntax: ``image_quality_score(col("img"))``.
Args:
*columns: Optional image column(s). When provided, the UDF is
applied directly instead of returning a factory.
concurrency: UDF concurrency. ``None`` uses the framework default.
Returns:
A UDF that returns a float quality score in ``[0, 1]`` combining
sharpness, contrast, and colorfulness, or ``None`` for null inputs.
Example::
>>> # As a reusable variable
>>> quality = image_quality_score()
>>> df = df.with_column("quality", quality(col("img")))
>>>
>>> # Inline
>>> df = df.with_column("quality", image_quality_score(col("img")))
"""
wrapper = udf(
_ImageQualityScore(),
return_dtype=DataType.float64(),
**_udf_runtime_kwargs(concurrency=concurrency),
)
return _build_or_apply_udf(wrapper, *columns)
# ImageWatermarkScore - watermark detection (pandas batch UDF)
class _ImageWatermarkScore(_HfImageClassifierScorer):
"""Watermark scoring using a HuggingFace image classifier."""
_label_keyword = "watermark"
_missing_label_hint = "watermark detection"
def __init__(self, hf_watermark_model=_DEFAULT_WATERMARK_MODEL,
model_sharing=None, num_gpus=None, gpu_type=None):
self.hf_watermark_model = hf_watermark_model
super().__init__(
hf_watermark_model,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
)
[docs]def image_watermark_score(
*columns,
hf_watermark_model=_DEFAULT_WATERMARK_MODEL,
model_sharing=None,
concurrency=None,
batch_size=None,
num_gpus=None,
gpu_type=None,
):
"""
Create a watermark scoring UDF (HuggingFace image classification-based).
Requires ``pip install transformers torch``. This is a pandas batch UDF
that supports GPU acceleration via ``num_gpus`` / ``gpu_type``.
Set ``num_gpus`` to request GPU resources from the DataFrame UDF runtime.
Args:
*columns: Optional image column(s). When provided, the UDF is
applied directly instead of returning a factory.
hf_watermark_model: HuggingFace model ID for watermark detection.
Defaults to the built-in watermark classifier model.
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 float watermark probability score in ``[0, 1]`` or
``None`` for null inputs. Higher values indicate stronger watermark
confidence. Use downstream filtering (e.g. ``> 0.8``) to convert to a
boolean decision.
Example::
>>> # As a reusable variable
>>> watermark = image_watermark_score()
>>> df = df.with_column("wm", watermark(col("img")))
>>> watermark_gpu = image_watermark_score(num_gpus=1.0)
>>> df = df.with_column("wm", watermark_gpu(col("img")))
>>>
>>> # Inline
>>> df = df.with_column("wm", image_watermark_score(col("img")))
>>> df = df.filter(image_watermark_score(col("img")) > 0.8)
"""
wrapper = udf(
_ImageWatermarkScore(
hf_watermark_model=hf_watermark_model,
model_sharing=model_sharing,
num_gpus=num_gpus,
gpu_type=gpu_type,
),
func_type="pandas",
return_dtype=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)