Curate ImagesTutorials

Create an Image Curation Pipeline

View as Markdown

Learn the basics of creating an image curation pipeline in Curator by following a complete workflow that filters images by aesthetic quality and NSFW content.

Tutorial Steps:
1:local:
2:depth: 2

Before You Start

  • Follow the Get Started guide to install the package, prepare the model directory, and set up your data paths.

Concepts and Mental Model

Use this overview to understand how stages pass data through the pipeline.

  • Pipeline: An ordered list of stages that process data.
  • Stage: A modular operation (for example, read, embed, filter, write).
  • Executor: Runs the pipeline (XennaExecutor backend).
  • Data units: Input images → embeddings → quality scores → filtered output.
  • Common choices:
    • Embeddings: CLIP ViT-L/14 for semantic understanding
    • Quality filters: Aesthetic predictor and NSFW classifier
    • Thresholds: Configurable scoring thresholds for filtering
  • Outputs: Filtered tar archives with embeddings, quality scores, and image data.

For more information, refer to the Image Concepts section.

1. Define Imports and Paths

Import required classes and define paths used throughout the example.

1from nemo_curator.backends.xenna import XennaExecutor
2from nemo_curator.pipeline import Pipeline
3from nemo_curator.stages.file_partitioning import FilePartitioningStage
4from nemo_curator.stages.image.embedders.clip_embedder import ImageEmbeddingStage
5from nemo_curator.stages.image.filters.aesthetic_filter import ImageAestheticFilterStage
6from nemo_curator.stages.image.filters.nsfw_filter import ImageNSFWFilterStage
7from nemo_curator.stages.image.io.image_reader import ImageReaderStage
8from nemo_curator.stages.image.io.image_writer import ImageWriterStage
9
10INPUT_TAR_DIR = "/path/to/input/tar_archives"
11OUTPUT_DIR = "/path/to/output/dataset"
12MODEL_DIR = "/path/to/models"

2. Create the Pipeline

Instantiate a named pipeline to orchestrate the stages.

1pipeline = Pipeline(
2    name="image_curation", 
3    description="Curate images with embeddings and quality scoring"
4)

3. Define Stages

Add modular stages to partition, read, embed, filter, and write images.

Partition Input Files

Divide tar files across workers for parallel processing.

1pipeline.add_stage(FilePartitioningStage(
2    file_paths=INPUT_TAR_DIR,
3    files_per_partition=1,  # Adjust based on file sizes
4    file_extensions=[".tar"],
5))

Read Images

Load images from tar archives and extract metadata.

1pipeline.add_stage(ImageReaderStage(
2    batch_size=100,  # Images per batch
3    verbose=True,
4    num_threads=16,       # I/O threads
5    num_gpus_per_worker=0.25,
6))

Generate CLIP Embeddings

Create semantic embeddings for each image using CLIP ViT-L/14.

1pipeline.add_stage(ImageEmbeddingStage(
2    model_dir=MODEL_DIR,
3    num_gpus_per_worker=0.25,
4    model_inference_batch_size=32,  # Adjust for GPU memory
5    remove_image_data=False,        # Keep images for filtering
6    verbose=True,
7))

Filter by Aesthetic Quality

Score and filter images based on aesthetic quality using a trained predictor.

1pipeline.add_stage(ImageAestheticFilterStage(
2    model_dir=MODEL_DIR,
3    num_gpus_per_worker=0.25,
4    model_inference_batch_size=32,
5    score_threshold=0.5,  # Keep images with score >= 0.5
6    verbose=True,
7))

Filter NSFW Content

Detect and filter out NSFW (Not Safe For Work) content.

1pipeline.add_stage(ImageNSFWFilterStage(
2    model_dir=MODEL_DIR,
3    num_gpus_per_worker=0.25,
4    model_inference_batch_size=32,
5    score_threshold=0.5,  # Filter images with NSFW score >= 0.5
6    verbose=True,
7))

Write Filtered Dataset

Save the curated images and metadata to output directory.

1pipeline.add_stage(ImageWriterStage(
2    output_dir=OUTPUT_DIR,
3    images_per_tar=1000,   # Images per output tar file (default: 1000)
4    remove_image_data=True, # Remove raw image data to save space
5    verbose=True,
6))

4. Run the Pipeline

Execute the configured pipeline. The pipeline will use XennaExecutor by default if no executor is specified.

1# Option 1: Use default executor (XennaExecutor)
2pipeline.run()
3
4# Option 2: Explicitly specify executor
5executor = XennaExecutor()
6pipeline.run(executor)

Complete Example

Here’s the full pipeline code:

1from nemo_curator.backends.xenna import XennaExecutor
2from nemo_curator.pipeline import Pipeline
3from nemo_curator.stages.file_partitioning import FilePartitioningStage
4from nemo_curator.stages.image.embedders.clip_embedder import ImageEmbeddingStage
5from nemo_curator.stages.image.filters.aesthetic_filter import ImageAestheticFilterStage
6from nemo_curator.stages.image.filters.nsfw_filter import ImageNSFWFilterStage
7from nemo_curator.stages.image.io.image_reader import ImageReaderStage
8from nemo_curator.stages.image.io.image_writer import ImageWriterStage
9
10def create_image_curation_pipeline():
11    """Create image curation pipeline with quality filtering."""
12    
13    # Define paths
14    INPUT_TAR_DIR = "/path/to/input/tar_archives"
15    OUTPUT_DIR = "/path/to/output/dataset" 
16    MODEL_DIR = "/path/to/models"
17    
18    # Create pipeline
19    pipeline = Pipeline(
20        name="image_curation",
21        description="Curate images with embeddings and quality scoring"
22    )
23    
24    # Add stages
25    pipeline.add_stage(FilePartitioningStage(
26        file_paths=INPUT_TAR_DIR,
27        files_per_partition=1,
28        file_extensions=[".tar"],
29    ))
30    
31    pipeline.add_stage(ImageReaderStage(
32        batch_size=100,
33        verbose=True,
34        num_threads=16,
35        num_gpus_per_worker=0.25,
36    ))
37    
38    pipeline.add_stage(ImageEmbeddingStage(
39        model_dir=MODEL_DIR,
40        num_gpus_per_worker=0.25,
41        model_inference_batch_size=32,
42        remove_image_data=False,
43        verbose=True,
44    ))
45    
46    pipeline.add_stage(ImageAestheticFilterStage(
47        model_dir=MODEL_DIR,
48        num_gpus_per_worker=0.25,
49        model_inference_batch_size=32,
50        score_threshold=0.5,
51        verbose=True,
52    ))
53    
54    pipeline.add_stage(ImageNSFWFilterStage(
55        model_dir=MODEL_DIR,
56        num_gpus_per_worker=0.25,
57        model_inference_batch_size=32,
58        score_threshold=0.5,
59        verbose=True,
60    ))
61    
62    pipeline.add_stage(ImageWriterStage(
63        output_dir=OUTPUT_DIR,
64        images_per_tar=1000,  # Default value
65        remove_image_data=True,
66        verbose=True,
67    ))
68    
69    return pipeline
70
71# Run the pipeline
72if __name__ == "__main__":
73    pipeline = create_image_curation_pipeline()
74    pipeline.run()  # Uses XennaExecutor by default

Performance Tuning

Batch Size Guidelines

The model_inference_batch_size parameter controls the number of images that the model processes at once, directly affecting GPU memory usage and performance. Higher batch sizes improve throughput but require more GPU memory, while lower batch sizes prevent out-of-memory (OOM) errors.

Memory Requirements and Actor Allocation

For the embedding stage (which uses the most GPU memory), each actor requires 0.25 GPU allocation, corresponding to about 20GB of GPU memory on an 80GB GPU (0.25 × 80GB = 20GB). When adjusting batch sizes, ensure memory usage stays below this 20GB threshold per actor.

Performance vs Memory Trade-offs

  • Higher batch sizes: Better GPU usage and faster processing, but risk OOM errors
  • Lower batch sizes: Safer memory usage and OOM prevention, but slower processing
  • Rule of thumb: Start with the recommended batch size for your GPU, then increase gradually while monitoring memory usage

GPU Allocation

The default num_gpus_per_worker=0.25 works well for most scenarios and allows 4 workers per GPU. You typically don’t need to change this unless you have specific requirements:

  • Default (recommended): 0.25 allows multiple workers per GPU for better utilization
  • Single worker per GPU: Set to 1.0 if you want exclusive GPU access per worker
  • CPU processing: Set to 0 (embedding and filtering stages require GPUs for optimal performance)

The system automatically calculates the optimal number of actors based on available resources and the num_gpus_per_worker setting.

Next Steps

After running this basic curation pipeline, you can:

  1. Adjust thresholds: Experiment with different aesthetic and NSFW score thresholds
  2. Add custom filters: Create domain-specific quality filters
  3. Run duplicate removal: Remove similar or duplicate images using semantic duplicate removal
  4. Export for training: Prepare curated data for downstream ML training tasks

For more advanced workflows, refer to the Image Duplicate Removal Tutorial.