Add: python API method to add custom embeddings

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Python SDK: `limit` parameter on `ImageDataset.add_samples_from*` methods to index only the first N samples of a dataset.
 - Python dataset queries now support model evaluation queries on the annotation level.
+- Python SDK: `Sample.set_embedding` method to provide custom, precomputed embeddings for image and video samples.
 
 ### Changed
 

lightly_studio/docs/docs/dataset_setup/image_dataset.md

@@ -27,7 +27,8 @@ it will recursively search for images in it. A remote path like `s3://my-bucket/
 supported, see [Using Cloud Storage](cloud_storage.md) for more details.
 
 Added images are automatically embedded so that embedding plot and image search are enabled.
-To skip embedding, pass `embed=False` to the method.
+To skip embedding, pass `embed=False` to the method. See [Custom Embeddings](#custom-embeddings)
+below if you want to provide your own embeddings instead of the built-in model.
 
 The method supports additional arguments, e.g. you can pass `tag_depth=1` to add the image parent
 folder name as a tag to each sample. See the [API reference](../api/dataset.md#lightly_studio.ImageDataset.add_images_from_path) for full details.
@@ -445,6 +446,58 @@ and stores it as annotation confidence.
 
 See the [API reference](../api/dataset.md#lightly_studio.ImageDataset) for `add_annotations_from_coco`, `add_annotations_from_yolo`, and `add_annotations_from_labelformat`.
 
+### Custom Embeddings
+
+Instead of the built-in embedding model, you can provide precomputed embeddings from your own
+model with `sample.set_embedding(...)`. This enables the embedding plot and embedding-based
+sampling for image and video samples.
+
+```python title="Add custom embeddings" hl_lines="11 23"
+import numpy as np
+from PIL import Image
+
+import lightly_studio as ls
+
+dataset_path = ls.utils.download_example_dataset(download_dir="dataset_examples")
+
+# Skip the built-in embedding model.
+dataset = ls.ImageDataset.create()
+dataset.add_images_from_path(
+    path=f"{dataset_path}/coco_subset_128_images/images", embed=False
+)
+
+
+def compute_custom_embedding(file_path: str) -> list[float]:
+    """Stand-in for your own model: the mean RGB color of the image."""
+    with Image.open(file_path) as image:
+        pixels = np.asarray(image.convert("RGB"), dtype=np.float32)
+    return pixels.mean(axis=(0, 1)).tolist()
+
+
+for sample in dataset:
+    sample.set_embedding(compute_custom_embedding(sample.file_path_abs))
+
+ls.start_gui()
+```
+
+See the full runnable version in
+[`example_custom_embeddings.py`](https://github.com/lightly-ai/lightly-studio/blob/main/lightly_studio/src/lightly_studio/examples/example_custom_embeddings.py).
+
+There are a few limitations to be aware of:
+
+- **No text-based search.** Custom embeddings do not support text-based search, since there is
+  no text encoder for a custom embedding space. The search bar is hidden automatically.
+- **`embed=False` only skips image-level embeddings.** If you also import annotations through
+  one of the bulk methods (`add_samples_from_coco`, `add_samples_from_yolo`,
+  `add_samples_from_labelformat`, `add_annotations_from_coco`, `add_annotations_from_yolo`, ...),
+  they compute annotation crop embeddings under a separate `embed_annotations` argument (also
+  `True` by default) — pass `embed_annotations=False` too if you want to avoid that. Attaching
+  annotations one by one with `sample.add_annotation(...)`/`add_annotations(...)` never computes
+  crop embeddings, regardless of any flag.
+- **Image and video samples only.** Custom embeddings via `sample.set_embedding(...)` currently
+  only work on image/video samples, not on annotation crops — there is no equivalent method on
+  `Annotation` yet.
+
 ## Image Dataset in the GUI
 
 Launch the GUI from your terminal:

lightly_studio/src/lightly_studio/api/routes/api/collection.py

@@ -136,9 +136,9 @@ def has_embeddings(
         Path(title="collection Id"),
         Depends(get_and_validate_collection_id),
     ],
-) -> bool:
-    """Check if a collection has embeddings."""
-    return embedding_utils.collection_has_embeddings(
+) -> embedding_utils.CollectionEmbeddingsStatus:
+    """Report which embedding-based features are available for a collection."""
+    return embedding_utils.get_collection_embeddings_status(
         session=session, collection_id=collection.collection_id
     )
 

lightly_studio/src/lightly_studio/core/sample.py

@@ -3,22 +3,27 @@
 from __future__ import annotations
 
 from abc import ABC
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Any, cast
 from uuid import UUID
 
+import numpy as np
 from sqlalchemy.orm import object_session
 from sqlmodel import Session
 
 from lightly_studio.core.annotation import CreateAnnotation
 from lightly_studio.models.annotation.annotation_base import AnnotationType
 from lightly_studio.models.caption import CaptionCreate
+from lightly_studio.models.embedding_model import EmbeddingModelCreate
 from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.sample_embedding import SampleEmbeddingCreate
 from lightly_studio.resolvers import (
     annotation_resolver,
     caption_resolver,
     collection_resolver,
+    embedding_model_resolver,
     metadata_resolver,
+    sample_embedding_resolver,
     tag_resolver,
 )
 
@@ -306,6 +311,41 @@ def add_annotations(
             collection_name=annotation_source,
         )
 
+    def set_embedding(
+        self, embedding: Sequence[float], embedding_model_name: str = "custom"
+    ) -> None:
+        """Set a custom embedding vector for this sample.
+
+        This lets you plug in your own model's embeddings instead of relying on the
+        built-in embedding pipeline. Note that samples embedded this way are not
+        compatible with text-based search: only the embedding plot, image-based
+        similarity, and sampling strategies read these vectors.
+
+        Args:
+            embedding: The embedding vector.
+            embedding_model_name: Name to group these vectors under. Samples that
+                share a name must always use vectors of the same length.
+        """
+        session = self.get_object_session()
+        model = embedding_model_resolver.get_or_create(
+            session=session,
+            embedding_model=EmbeddingModelCreate(
+                name=embedding_model_name,
+                embedding_model_hash=f"custom:{embedding_model_name}:{len(embedding)}",
+                embedding_dimension=len(embedding),
+                collection_id=self.collection_id,
+                supports_text_search=False,
+            ),
+        )
+        sample_embedding_resolver.upsert(
+            session=session,
+            sample_embedding=SampleEmbeddingCreate(
+                sample_id=self.sample_id,
+                embedding_model_id=model.embedding_model_id,
+                embedding=np.asarray(embedding, dtype=np.float32),
+            ),
+        )
+
     def delete_annotation(self, annotation: Annotation) -> None:
         """Delete an annotation from this sample.
 

lightly_studio/src/lightly_studio/database/db_manager.py

@@ -120,6 +120,7 @@ def __init__(
             )
         else:
             SQLModel.metadata.create_all(self._engine)
+            _migrate_duckdb_schema(self._engine)
 
     @contextmanager
     def session(self) -> Generator[Session, None, None]:
@@ -310,6 +311,26 @@ def _initialize_postgres_schema(
     db_migrations.run_migrations(engine=engine, engine_url=engine_url)
 
 
+def _migrate_duckdb_schema(engine: Engine) -> None:
+    """Apply in-place schema migrations for existing DuckDB databases.
+
+    DuckDB has no migration framework (Alembic only covers PostgreSQL) and
+    ``SQLModel.metadata.create_all`` never alters existing tables, so columns added
+    to a model after a database file was created must be added here. All statements
+    must be idempotent.
+    """
+    with engine.connect() as conn:
+        # Added 07/2026. Defaults to true because all rows created before this column
+        # existed came from CLIP-style vision-text generators.
+        conn.execute(
+            statement=text(
+                "ALTER TABLE embedding_model "
+                "ADD COLUMN IF NOT EXISTS supports_text_search BOOLEAN DEFAULT true"
+            )
+        )
+        conn.commit()
+
+
 def _detect_backend_from_url(engine_url: str) -> DatabaseBackend:
     """Detect the database backend from the engine URL.
 

lightly_studio/src/lightly_studio/dataset/embedding_generator.py

@@ -142,6 +142,7 @@ def get_embedding_model_input(self, collection_id: UUID) -> EmbeddingModelCreate
             embedding_model_hash="random_model",
             embedding_dimension=self._dimension,
             collection_id=collection_id,
+            supports_text_search=True,
         )
 
     def embed_text(self, _text: str) -> list[float]:

lightly_studio/src/lightly_studio/dataset/embedding_utils.py

@@ -2,36 +2,63 @@
 
 from uuid import UUID
 
-from sqlmodel import Session
+from sqlmodel import Session, SQLModel
 
-from lightly_studio.dataset.embedding_manager import EmbeddingManagerProvider
 from lightly_studio.resolvers import (
+    embedding_model_resolver,
     sample_embedding_resolver,
 )
 
 
-def collection_has_embeddings(session: Session, collection_id: UUID) -> bool:
-    """Check if there are any embeddings available for the given collection.
+class CollectionEmbeddingsStatus(SQLModel):
+    """Which embedding-based features are available for a collection."""
+
+    has_embeddings: bool
+    """At least one sample embedding is stored, under any embedding model.
+
+    Enables the embedding plot, image similarity, and embedding-based sampling.
+    """
+
+    has_text_search_embeddings: bool
+    """At least one sample embedding is stored under a text-capable embedding model.
+
+    Enables text-based search. Custom embeddings (``Sample.set_embedding``) have no
+    matching text encoder and do not count here.
+    """
+
+
+def get_collection_embeddings_status(
+    session: Session, collection_id: UUID
+) -> CollectionEmbeddingsStatus:
+    """Report which embedding-based features are available for the given collection.
+
+    Reads only the database (``embedding_model`` and ``sample_embedding`` tables); it
+    never registers or loads embedding models as a side effect.
 
     Args:
         session: Database session for resolver operations.
-        collection_id: The ID of the collection to check for embeddings.
+        collection_id: The ID of the collection to check.
 
     Returns:
-        True if embeddings exist for the collection, False otherwise.
+        The embeddings status of the collection.
     """
-    embedding_manager = EmbeddingManagerProvider.get_embedding_manager()
-    model_id = embedding_manager.load_or_get_default_model(
+    embedding_models = embedding_model_resolver.get_all_by_collection_id(
+        session=session, collection_id=collection_id
+    )
+    has_embeddings = sample_embedding_resolver.has_any_embedding(
+        session=session,
+        collection_id=collection_id,
+        embedding_model_ids=[model.embedding_model_id for model in embedding_models],
+    )
+    if not has_embeddings:
+        return CollectionEmbeddingsStatus(has_embeddings=False, has_text_search_embeddings=False)
+    has_text_search_embeddings = sample_embedding_resolver.has_any_embedding(
         session=session,
         collection_id=collection_id,
+        embedding_model_ids=[
+            model.embedding_model_id for model in embedding_models if model.supports_text_search
+        ],
     )
-    if model_id is None:
-        # No default embedding model loaded for this collection.
-        return False
-
-    return (
-        sample_embedding_resolver.get_embedding_count(
-            session=session, collection_id=collection_id, embedding_model_id=model_id
-        )
-        > 0
+    return CollectionEmbeddingsStatus(
+        has_embeddings=True, has_text_search_embeddings=has_text_search_embeddings
     )

lightly_studio/src/lightly_studio/dataset/mobileclip_embedding_generator.py

@@ -65,6 +65,7 @@ def get_embedding_model_input(self, collection_id: UUID) -> EmbeddingModelCreate
             embedding_model_hash=self._model_hash,
             embedding_dimension=EMBEDDING_DIMENSION,
             collection_id=collection_id,
+            supports_text_search=True,
         )
 
     def embed_text(self, text: str) -> list[float]:

lightly_studio/src/lightly_studio/dataset/perception_encoder_embedding_generator.py

@@ -144,6 +144,7 @@ def get_embedding_model_input(self, collection_id: UUID) -> EmbeddingModelCreate
             embedding_model_hash=self._model_hash,
             embedding_dimension=self._model.output_dim,
             collection_id=collection_id,
+            supports_text_search=True,
         )
 
     def embed_text(self, text: str) -> list[float]:

lightly_studio/src/lightly_studio/examples/example_custom_embeddings.py

@@ -0,0 +1,46 @@
+"""Example of adding custom, precomputed embeddings to a dataset.
+
+Use `Sample.set_embedding` to plug in embeddings from your own model instead of
+relying on the built-in embedding pipeline. This enables the embedding plot,
+image-based similarity, and embedding-based sampling strategies. Load the dataset
+with `embed=False` so the built-in model never runs.
+
+Note: text-based search does not work with custom embeddings, since there is no
+text encoder for a custom embedding space. The GUI hides the search bar
+automatically.
+"""
+
+import numpy as np
+from environs import Env
+from PIL import Image
+
+import lightly_studio as ls
+from lightly_studio.database import db_manager
+
+# Read environment variables
+env = Env()
+env.read_env()
+
+# Cleanup an existing database
+db_manager.connect(cleanup_existing=True)
+
+# Define the path to the dataset directory
+dataset_path = env.path("EXAMPLES_DATASET_PATH", "/path/to/your/dataset")
+
+# Create a Dataset from a path, skipping the built-in embedding pipeline.
+dataset = ls.ImageDataset.create()
+dataset.add_images_from_path(path=dataset_path, embed=False)
+
+
+def compute_custom_embedding(file_path: str) -> list[float]:
+    """Stand-in for your own model: the mean RGB color of the image."""
+    with Image.open(file_path) as image:
+        pixels = np.asarray(image.convert("RGB"), dtype=np.float32)
+    return [float(value) for value in pixels.mean(axis=(0, 1))]
+
+
+for sample in dataset:
+    embedding = compute_custom_embedding(sample.file_path_abs)
+    sample.set_embedding(embedding)
+
+ls.start_gui()

lightly_studio/src/lightly_studio/migrations/versions/1783577705_0c5e19664860_add_supports_text_search_to_embedding_.py

@@ -0,0 +1,39 @@
+"""add_supports_text_search_to_embedding_model.
+
+Revision ID: 0c5e19664860
+Revises: a1b2c3d4e5f6
+Create Date: 2026-07-09 10:15:00.000000
+
+"""
+
+from collections.abc import Sequence
+from typing import Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "0c5e19664860"
+down_revision: Union[str, Sequence[str], None] = "a1b2c3d4e5f6"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # Defaults to true because all rows created before this column existed came
+    # from CLIP-style vision-text generators.
+    op.add_column(
+        "embedding_model",
+        sa.Column(
+            "supports_text_search",
+            sa.Boolean(),
+            server_default=sa.true(),
+            nullable=False,
+        ),
+    )
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    op.drop_column("embedding_model", "supports_text_search")

lightly_studio/src/lightly_studio/models/embedding_model.py

@@ -16,6 +16,7 @@ class EmbeddingModelBase(SQLModel):
     embedding_model_hash: str = Field(default="", sa_column=Column(VARCHAR(128)))
     embedding_dimension: int
     collection_id: UUID = Field(default=None, foreign_key="collection.collection_id", index=True)
+    supports_text_search: bool = Field(default=True)
 
 
 class EmbeddingModelCreate(EmbeddingModelBase):