stages.image.io.image_writer#

Module Contents#

Classes#

ImageWriterStage

Write images to tar files and corresponding metadata to a Parquet file.

API#

class stages.image.io.image_writer.ImageWriterStage#

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

Write images to tar files and corresponding metadata to a Parquet file.

  • Images are packed into tar archives with at most images_per_tar entries each.

  • Metadata for all written images in the batch is stored in a single Parquet file.

  • Tar filenames are unique across actors via an actor-scoped prefix.

construct_base_name(task: nemo_curator.tasks.image.ImageBatch) → str#

Construct a base name for tar files within this actor.

deterministic_name: bool#

True

images_per_tar: int#

1000

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

output_dir: str#

None

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.image.ImageBatch,
) → nemo_curator.tasks.file_group.FileGroupTask#

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

remove_image_data: bool#

False

verbose: bool#

False