stages.image.filters.aesthetic_filter

Module Contents

Classes

ImageAestheticFilterStage

Stage for filtering out images based on aesthetic scores.

API

class stages.image.filters.aesthetic_filter.ImageAestheticFilterStage

Bases: nemo_curator.stages.image.filters.base.BaseFilterStage

Stage for filtering out images based on aesthetic scores.

This class processes image batches through an aesthetic scoring model to generate aesthetic scores for each image. Images with scores below the threshold will be filtered out.

model_dir: str

None

model_inference_batch_size: int

32

name: str

‘image_aesthetic_filter’

num_gpus_per_worker: float

0.25

process(

task: nemo_curator.tasks.ImageBatch,

) → nemo_curator.tasks.ImageBatch

Process an image batch to filter by aesthetic score threshold.

Args: task: ImageBatch containing list of ImageObject instances with aesthetic scores

Returns: ImageBatch with filtered images that meet the aesthetic score threshold.

score_threshold: float

0.5

setup(

_worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None,

) → None

Initialize the aesthetic filtering model.

setup_on_node(

_node_info: nemo_curator.backends.base.NodeInfo | None = None,

_worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None,

) → None

Download aesthetic model weights from HF

verbose: bool

False