stages.image.filters.base#
Module Contents#
Classes#
Base class for image filtering stages. |
API#
- class stages.image.filters.base.BaseFilterStage#
Bases:
nemo_curator.stages.base.ProcessingStage[nemo_curator.tasks.ImageBatch,nemo_curator.tasks.ImageBatch]Base class for image filtering stages.
This class provides a base class for image filtering stages.
- inputs() tuple[list[str], list[str]]#
Define stage input requirements.
Returns (tuple[list[str], list[str]]): Tuple of (required_attributes, required_columns) where: - required_top_level_attributes: List of task attributes that must be present - required_data_attributes: List of attributes within the data that must be present
- model_dir: str#
None
- model_inference_batch_size: int#
32
- name: str#
‘image_filter’
- num_gpus_per_worker: float#
0.25
- outputs() tuple[list[str], list[str]]#
Define stage output specification.
Returns (tuple[list[str], list[str]]): Tuple of (output_attributes, output_columns) where: - output_top_level_attributes: List of task attributes this stage adds/modifies - output_data_attributes: List of attributes within the data that this stage adds/modifies
- abstractmethod process(
- task: nemo_curator.tasks.ImageBatch,
Process an image batch to generate scores and filter by threshold.
Args: task: ImageBatch containing list of ImageObject instances with pre-computed embeddings
Returns: ImageBatch with filtered images that have scores below the threshold
- score_threshold: float#
0.5
- abstractmethod setup(
- _worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None,
Initialize the base filter stage.
- verbose: bool#
False
- yield_next_batch(
- task: nemo_curator.tasks.ImageBatch,
Yields a generator of model inputs for the next batch.
Args: task (ImageBatch): The ImageBatch to process.
Yields: Generator[dict[str, torch.Tensor]]: A generator of model inputs for the next batch.