tao-generate-image-embeddings

name	tao-generate-image-embeddings
description	Run TAO Data Services image embedding to turn a parquet of image filepaths into an embedding parquet using CLIP, SigLIP, or a TAO checkpoint. Use when a workflow needs embeddings before nearest-neighbor or unique-neighbor mining, or when the user asks to "embed images", "compute image embeddings", or "generate SigLIP embeddings".
license	Apache-2.0
compatibility	Requires docker, nvidia-container-toolkit, one or more CUDA GPUs, and the TAO data-services container pinned in versions.yaml.
metadata	{"author":"NVIDIA Corporation","version":"0.1.0"}
allowed-tools	Read Bash
tags	["tao","data","embeddings","mining","siglip","clip"]

TAO Generate Image Embeddings

Use this skill to run TAO Data Services image embedding. The skill consumes a parquet of image filepaths and writes a parquet with an embedding column. Downstream mining skills (tao-mine-od-images, tao-mine-nearest-neighbors) consume its output.

The container entrypoint is:

embedding image_embeddings -e /absolute/path/to/image_embeddings.yaml

Inputs

The user can provide either an existing spec or the fields needed to generate one.

Required spec fields:

Field	Meaning
input_parquet	Absolute path to a parquet containing image filepaths.
output_parquet	Absolute path where the embedding parquet is written.
model	CLIP or SigLIP.
model_path	HuggingFace model id, local HF snapshot directory, or a TAO .pth/.ckpt checkpoint. Must match model. SigLIP: google/siglip-base-patch16-224 (768-dim, the template default). CLIP: openai/clip-vit-base-patch32 (512-dim). The two fields are validated independently, so a mismatched pair is accepted here and fails or mis-loads inside the container.

Common optional fields:

Field	Default	Meaning
model_config_path	""	TAO experiment spec path. Required only when model_path is a TAO checkpoint.
batch_size	64	Number of images processed in parallel. Lower it if the GPU runs out of memory.

The input parquet must contain a filepath column. Any additional columns are carried through to the output verbatim, so metadata such as label survives into the embedding parquet.

The default template is assets/default_image_embeddings.yaml.

Encoder Consistency

When embeddings feed a mining step, every parquet compared against another must be produced with the same model and model_path. Embedding dimensionality follows the encoder — 768 for the SigLIP default, 512 for CLIP ViT-B/32 — and nothing in the output parquet records which encoder wrote it. Embeddings from different encoders are not comparable, and mismatched encoders are the most common cause of mining output that looks unrelated to the targets. Reuse one spec across every parquet in a mining run and override only input_parquet / output_parquet.

Quick Start

Run from the tao-skills-external repo root.

Write the spec beside the output parquet. The run does not retain it, so the embeddings otherwise carry no record of the encoder that produced them. That matters here more than elsewhere: every parquet compared against another in a mining step must come from the same model and model_path, and mismatched encoders are the usual cause of mining output that looks unrelated to its targets.

OUT_DIR=/absolute/path/for/this/run              # where output_parquet is written
SPEC="$OUT_DIR/image_embeddings.yaml"            # spec lives beside its output
RUN_ROOT=/absolute/path/that/contains/parquets/images/and/results
GPU_COUNT=1

python3 skills/data/tao-generate-image-embeddings/scripts/verify_image_embeddings_spec.py \
  --spec "$SPEC"

DS_IMAGE=nvcr.io/nvstaging/tao/tao-dataservices:gapanalysis-02  # versions-key: images.tao_toolkit.data_services_od

docker run --rm --gpus "$GPU_COUNT" --ipc=host --network=host \
  -v "$RUN_ROOT:$RUN_ROOT" \
  -v "$HOME/.cache/huggingface:/root/.cache/huggingface" \
  -w "$RUN_ROOT" \
  "$DS_IMAGE" \
  embedding image_embeddings -e "$SPEC"

Do not pass --user $(id -u):$(id -g) to the TAO data-services container; the image imports transformers at startup, which calls getpass.getuser() and fails when the UID is not present in /etc/passwd.

To embed several parquets with one encoder, reuse the same spec and override the two paths per run:

docker run --rm --gpus "$GPU_COUNT" --ipc=host --network=host \
  -v "$RUN_ROOT:$RUN_ROOT" -w "$RUN_ROOT" "$DS_IMAGE" \
  embedding image_embeddings -e "$SPEC" \
  input_parquet=/abs/path/other_input.parquet \
  output_parquet=/abs/path/other_output.parquet

Generate A Spec

If the user provides parquet paths and an encoder instead of a ready spec, copy the template and fill in the nulls. Every tuning value it already carries is the one this stage wants — change one only deliberately.

cp skills/data/tao-generate-image-embeddings/assets/default_image_embeddings.yaml "$SPEC"

Fill input_parquet, output_parquet and — if not using the default encoder — model and model_path, all as absolute paths, then validate:

python3 skills/data/tao-generate-image-embeddings/scripts/verify_image_embeddings_spec.py --spec "$SPEC"

input_parquet: /absolute/path/filepaths.parquet
output_parquet: /absolute/path/results/embeddings.parquet
model: SigLIP
model_path: google/siglip-base-patch16-224
model_config_path: ""          # required only when model_path is a TAO .pth/.ckpt
batch_size: 64

The template is the only place a default value lives, so nothing can disagree with it. verify reports the encoder, since embeddings are only comparable to others produced by the same model and model_path.

Keep the spec, input parquet, image files, and output directory under RUN_ROOT so the same paths resolve inside the container.

Preflight

1. Verify Docker and GPU access:

docker info > /dev/null
nvidia-smi -L

2. Resolve and pull the data-services image if needed:

DS_IMAGE=nvcr.io/nvstaging/tao/tao-dataservices:gapanalysis-02  # versions-key: images.tao_toolkit.data_services_od
docker image inspect "$DS_IMAGE" > /dev/null || docker pull "$DS_IMAGE"

3. Validate the spec:

python3 skills/data/tao-generate-image-embeddings/scripts/verify_image_embeddings_spec.py \
  --spec "$SPEC"

4. Confirm RUN_ROOT contains the spec, the input parquet, the image files its filepath column points at, and the output directory. Mount RUN_ROOT to the same absolute path inside Docker.

Outputs

Artifact	Location
embedding parquet	output_parquet

The output parquet contains filepath, an embedding column of list-like vectors, and every extra column carried through from the input. Print its row count and column list after the run so the caller can confirm the embedding column exists.

Troubleshooting

The subtask image_embeddings requires -e/--experiment_spec_file: rerun with embedding image_embeddings -e "$SPEC".

Input parquet or images not found inside Docker: the filepath values are read verbatim. Use a RUN_ROOT mount where host and container paths are identical, and confirm the images themselves are under that mount — not just the parquet.

Model loading error with a .pth / .ckpt model_path: TAO checkpoints need model_config_path set to the training spec so the architecture can be rebuilt. HuggingFace ids and snapshot directories do not.

CUDA out of memory: lower batch_size (try 32 or 16).

Mined results look unrelated downstream: the parquets compared during mining were embedded with different encoders. Re-embed them with one shared spec — see ## Encoder Consistency.

No GPU available: embedding requires at least one CUDA GPU. Check nvidia-smi -L, the Docker --gpus flag, and the NVIDIA container toolkit installation.