Create an Image Curation Pipeline

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.

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

2. Create the Pipeline

Instantiate a named pipeline to orchestrate the stages.

1 pipeline = 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.

1 pipeline.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.

1 pipeline.add_stage(ImageReaderStage(
2     dali_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.

1 pipeline.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.

1 pipeline.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 &gt;= 0.5
6     verbose=True,
7 ))

Filter NSFW Content

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

1 pipeline.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 &gt;= 0.5
6     verbose=True,
7 ))

Write Filtered Dataset

Save the curated images and metadata to output directory.

1 pipeline.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)
2 pipeline.run()
3 
4 # Option 2: Explicitly specify executor
5 executor = XennaExecutor()
6 pipeline.run(executor)

Complete Example

Here’s the full pipeline code:

1 from nemo_curator.backends.xenna import XennaExecutor
2 from nemo_curator.pipeline import Pipeline
3 from nemo_curator.stages.file_partitioning import FilePartitioningStage
4 from nemo_curator.stages.image.embedders.clip_embedder import ImageEmbeddingStage
5 from nemo_curator.stages.image.filters.aesthetic_filter import ImageAestheticFilterStage
6 from nemo_curator.stages.image.filters.nsfw_filter import ImageNSFWFilterStage
7 from nemo_curator.stages.image.io.image_reader import ImageReaderStage
8 from nemo_curator.stages.image.io.image_writer import ImageWriterStage
9 
10 def 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         dali_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
72 if __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:

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

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

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

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

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

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

1	pipeline.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	))

1	pipeline.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	))

1	pipeline.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	))

1	pipeline.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	))

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