nemo_curator.stages.image.filters.nsfw_filter

Module Contents

Classes

Name	Description
`ImageNSFWFilterStage`	Stage for filtering out NSFW images using NSFWScorer model.

Data

__all__

API

class nemo_curator.stages.image.filters.nsfw_filter.ImageNSFWFilterStage(
    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_nsfw_filter',
    weights_path: str = None
)

Dataclass

Bases: BaseFilterStage

Stage for filtering out NSFW images using NSFWScorer model.

This class processes image batches through an NSFW scoring model to generate NSFW probability scores for each image. Images with scores above the threshold will be filtered out as NSFW content.

name

str = 'image_nsfw_filter'

weights_path

str = None

nemo_curator.stages.image.filters.nsfw_filter.ImageNSFWFilterStage.process(
    task: nemo_curator.tasks.ImageBatch
) -> nemo_curator.tasks.ImageBatch

Process an image batch to generate NSFW scores and filter by threshold.

Parameters:

task

ImageBatch

ImageBatch containing list of ImageObject instances with pre-computed embeddings

Returns: ImageBatch

ImageBatch with filtered images that have NSFW scores below the threshold

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

Initialize the NSFW filtering model.

nemo_curator.stages.image.filters.nsfw_filter.ImageNSFWFilterStage.setup_on_node(
    _node_info: nemo_curator.backends.base.NodeInfo | None = None,
    _worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None

Download NSFW model weights from LAION repository.

nemo_curator.stages.image.filters.nsfw_filter.__all__ = ['ImageNSFWFilterStage']