Implementation:NVIDIA NeMo Curator ImageFilterBase

Knowledge Sources	NVIDIA NeMo Curator
Domains	Image Processing, Data Filtering, GPU Computing
Last Updated	2026-02-14 00:00 GMT

Overview

BaseFilterStage is the abstract base class for all image filtering stages in NeMo Curator, providing shared configuration parameters, GPU resource allocation, and a reusable batching mechanism for processing images through models.

Description

BaseFilterStage extends ProcessingStage[ImageBatch, ImageBatch] as a Python dataclass. It defines common configuration fields used across concrete image filter implementations such as aesthetic filters and NSFW filters. The class handles GPU resource allocation automatically based on CUDA availability: when a GPU is detected, it assigns fractional GPU resources (defaulting to 0.25 GPUs per worker); otherwise, it falls back to CPU-only resources.

The class provides a yield_next_batch() generator method that chunks the ImageBatch.data list of ImageObject instances into sub-batches of configurable size (controlled by model_inference_batch_size). This batching pattern allows subclasses to efficiently feed images through inference models in manageable groups.

Both setup() and process() methods raise NotImplementedError, requiring all concrete subclasses to provide their own model initialization and scoring logic.

Usage

Use BaseFilterStage as the parent class when implementing new image filtering stages. Subclass it and implement setup() to initialize model weights and process() to score images and filter by threshold. Existing subclasses include ImageAestheticFilterStage and ImageNSFWFilterStage.

Code Reference

Source Location

Repository: NeMo-Curator
File: nemo_curator/stages/image/filters/base.py
Lines: 1-82

Signature

@dataclass
class BaseFilterStage(ProcessingStage[ImageBatch, ImageBatch]):
    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_filter"

    def __post_init__(self) -> None: ...
    def inputs(self) -> tuple[list[str], list[str]]: ...
    def outputs(self) -> tuple[list[str], list[str]]: ...
    def setup(self, _worker_metadata: WorkerMetadata | None = None) -> None: ...
    def yield_next_batch(self, task: ImageBatch) -> Generator[list[ImageObject], None, None]: ...
    def process(self, task: ImageBatch) -> ImageBatch: ...

Import

from nemo_curator.stages.image.filters.base import BaseFilterStage

I/O Contract

Inputs

Name	Type	Required	Description
model_dir	`str`	No	Path to the directory containing model weights for inference
num_gpus_per_worker	`float`	No	Fractional GPU allocation per worker (default: 0.25)
model_inference_batch_size	`int`	No	Number of images to process through the model at once (default: 32)
score_threshold	`float`	No	Threshold for filtering images by score (default: 0.5)
verbose	`bool`	No	Enable verbose logging output (default: False)
task	`ImageBatch`	Yes	The image batch to process, containing a list of `ImageObject` instances

Outputs

Name	Type	Description
result	`ImageBatch`	Filtered image batch containing only images that passed the score threshold

Key Implementation Details

GPU Resource Allocation

The __post_init__ method automatically detects CUDA availability and configures resources accordingly:

def __post_init__(self) -> None:
    if torch.cuda.is_available():
        self.resources = Resources(gpus=self.num_gpus_per_worker)
    else:
        self.resources = Resources()

Batch Yielding

The yield_next_batch method provides a generator that chunks the image data into sub-batches for efficient model inference:

def yield_next_batch(self, task: ImageBatch) -> Generator[list[ImageObject], None, None]:
    for i in range(0, len(task.data), self.model_inference_batch_size):
        yield task.data[i : i + self.model_inference_batch_size]

I/O Declaration

Both inputs() and outputs() declare ["data"] as the required field and no additional subfields, indicating the stage reads and writes the full ImageBatch.data list.

Usage Examples

Subclassing BaseFilterStage

from dataclasses import dataclass
from nemo_curator.stages.image.filters.base import BaseFilterStage
from nemo_curator.backends.base import WorkerMetadata
from nemo_curator.tasks import ImageBatch

@dataclass
class MyCustomFilterStage(BaseFilterStage):
    score_threshold: float = 0.7
    name: str = "my_custom_filter"

    def setup(self, _worker_metadata: WorkerMetadata | None = None) -> None:
        # Load your custom model here
        self.model = load_my_model(self.model_dir)

    def process(self, task: ImageBatch) -> ImageBatch:
        filtered_images = []
        for batch in self.yield_next_batch(task):
            scores = self.model.predict(batch)
            for img, score in zip(batch, scores):
                if score < self.score_threshold:
                    filtered_images.append(img)
        task.data = filtered_images
        return task

Related Pages

Environment:NVIDIA_NeMo_Curator_Python_Linux_Base
NVIDIA_NeMo_Curator_ImageAestheticFilterStage - Concrete aesthetic filter subclass
NVIDIA_NeMo_Curator_ImageNSFWFilterStage - Concrete NSFW filter subclass
NVIDIA_NeMo_Curator_ImageEmbeddingStage - Related image embedding stage

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment