nemo_curator.stages.image.embedders.clip_embedder

Module Contents

Classes

Name	Description
`ImageEmbeddingStage`	Stage for generating image embeddings using CLIP model.

API

class nemo_curator.stages.image.embedders.clip_embedder.ImageEmbeddingStage(
    model_dir: str = None,
    num_gpus_per_worker: float = 0.25,
    model_inference_batch_size: int = 32,
    verbose: bool = False,
    remove_image_data: bool = False,
    name: str = 'image_embedding'
)

Dataclass

Bases: ProcessingStage[ImageBatch, ImageBatch]

Stage for generating image embeddings using CLIP model.

This class processes image batches through a CLIP model to generate embeddings for each image. It assumes image data is already loaded in ImageObject.image_data and stores embeddings in ImageObject.embedding.

model_dir

str = None

model_inference_batch_size

int = 32

name

str = 'image_embedding'

num_gpus_per_worker

float = 0.25

remove_image_data

bool = False

verbose

bool = False

nemo_curator.stages.image.embedders.clip_embedder.ImageEmbeddingStage.__post_init__() -> None

nemo_curator.stages.image.embedders.clip_embedder.ImageEmbeddingStage.inputs() -> tuple[list[str], list[str]]

nemo_curator.stages.image.embedders.clip_embedder.ImageEmbeddingStage.outputs() -> tuple[list[str], list[str]]

nemo_curator.stages.image.embedders.clip_embedder.ImageEmbeddingStage.process(
    task: nemo_curator.tasks.ImageBatch
) -> nemo_curator.tasks.ImageBatch

Process an image batch to generate embeddings.

Parameters:

task

ImageBatch

ImageBatch containing list of ImageObject instances with pre-loaded image_data

Returns: ImageBatch

ImageBatch with embeddings stored in ImageObject.embedding

nemo_curator.stages.image.embedders.clip_embedder.ImageEmbeddingStage.setup(
    _worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None

Initialize the CLIP image embedding model.

nemo_curator.stages.image.embedders.clip_embedder.ImageEmbeddingStage.setup_on_node(
    node_info: nemo_curator.backends.base.NodeInfo,
    worker_metadata: nemo_curator.backends.base.WorkerMetadata
) -> None

Download the weights for the CLIP model on the node.

nemo_curator.stages.image.embedders.clip_embedder.ImageEmbeddingStage.yield_next_batch(
    task: nemo_curator.tasks.ImageBatch
) -> collections.abc.Generator[nemo_curator.tasks.ImageBatch, None, None]

Yield batches of images from the task.

Parameters:

task

ImageBatch

ImageBatch containing list of ImageObject instances with pre-loaded image_data