stages.image.io.image_reader#

Module Contents#

Classes#

ImageReaderStage

DALI-based reader that loads images from WebDataset tar shards.

API#

class stages.image.io.image_reader.ImageReaderStage#

Bases: nemo_curator.stages.base.ProcessingStage[nemo_curator.tasks.FileGroupTask, nemo_curator.tasks.ImageBatch]

DALI-based reader that loads images from WebDataset tar shards.

Works with DALI GPU (CUDA) or DALI CPU; decodes on GPU if CUDA is available, otherwise falls back to CPU decoding.

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

num_gpus_per_worker: float#

0.25

num_threads: int#

8

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

process(
task: nemo_curator.tasks.FileGroupTask,
) list[nemo_curator.tasks.ImageBatch]#

Process a task and return the result. Args: task (X): Input task to process Returns (Y | list[Y]): - Single task: For 1-to-1 transformations - List of tasks: For 1-to-many transformations (e.g., readers) - None: If the task should be filtered out

task_batch_size: int#

100

verbose: bool#

True