nemo_curator.stages.image.filters.aesthetic_filter

Module Contents

Classes

Name	Description
`ImageAestheticFilterStage`	Stage for filtering out images based on aesthetic scores.

API

class nemo_curator.stages.image.filters.aesthetic_filter.ImageAestheticFilterStage(
    model_dir: str = None,
    num_gpus_per_worker: float = 0.25,
    model_inference_batch_size: int = 32,
    score_threshold: float = 0.5,
    verbose: bool = False,
    name: str = 'image_aesthetic_filter'
)

Dataclass

Bases: 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

score_threshold

float = 0.5

verbose

bool = False

nemo_curator.stages.image.filters.aesthetic_filter.ImageAestheticFilterStage.process(
    task: nemo_curator.tasks.ImageBatch
) -> nemo_curator.tasks.ImageBatch

Process an image batch to filter by aesthetic score threshold.

Parameters:

task

ImageBatch

ImageBatch containing list of ImageObject instances with aesthetic scores

Returns: ImageBatch

ImageBatch with filtered images that meet the aesthetic score threshold.

nemo_curator.stages.image.filters.aesthetic_filter.ImageAestheticFilterStage.setup(
    _worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None

Initialize the aesthetic filtering model.

nemo_curator.stages.image.filters.aesthetic_filter.ImageAestheticFilterStage.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