diff --git a/CHANGELOG.md b/CHANGELOG.md index a9f6757aee..16455dbcb7 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/lightly_studio/docs/docs/dataset_setup/image_dataset.md b/lightly_studio/docs/docs/dataset_setup/image_dataset.md index da53f077a5..85de2b5915 100644 --- a/lightly_studio/docs/docs/dataset_setup/image_dataset.md +++ b/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: diff --git a/lightly_studio/src/lightly_studio/api/routes/api/collection.py b/lightly_studio/src/lightly_studio/api/routes/api/collection.py index 3ed44cb27e..4a70f80734 100644 --- a/lightly_studio/src/lightly_studio/api/routes/api/collection.py +++ b/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 ) diff --git a/lightly_studio/src/lightly_studio/core/sample.py b/lightly_studio/src/lightly_studio/core/sample.py index 57447933bd..a460d259e6 100644 --- a/lightly_studio/src/lightly_studio/core/sample.py +++ b/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. diff --git a/lightly_studio/src/lightly_studio/database/db_manager.py b/lightly_studio/src/lightly_studio/database/db_manager.py index eb45303598..3b37c487e1 100644 --- a/lightly_studio/src/lightly_studio/database/db_manager.py +++ b/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. diff --git a/lightly_studio/src/lightly_studio/dataset/embedding_generator.py b/lightly_studio/src/lightly_studio/dataset/embedding_generator.py index f76a1014e4..c6fde847be 100644 --- a/lightly_studio/src/lightly_studio/dataset/embedding_generator.py +++ b/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]: diff --git a/lightly_studio/src/lightly_studio/dataset/embedding_utils.py b/lightly_studio/src/lightly_studio/dataset/embedding_utils.py index 6f9ae2ec6f..0ae2d4a6e6 100644 --- a/lightly_studio/src/lightly_studio/dataset/embedding_utils.py +++ b/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 ) diff --git a/lightly_studio/src/lightly_studio/dataset/mobileclip_embedding_generator.py b/lightly_studio/src/lightly_studio/dataset/mobileclip_embedding_generator.py index 8ccd326aca..8aa229dbc6 100644 --- a/lightly_studio/src/lightly_studio/dataset/mobileclip_embedding_generator.py +++ b/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]: diff --git a/lightly_studio/src/lightly_studio/dataset/perception_encoder_embedding_generator.py b/lightly_studio/src/lightly_studio/dataset/perception_encoder_embedding_generator.py index e1d9a3c0b5..6dfbed2ba5 100644 --- a/lightly_studio/src/lightly_studio/dataset/perception_encoder_embedding_generator.py +++ b/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]: diff --git a/lightly_studio/src/lightly_studio/examples/example_custom_embeddings.py b/lightly_studio/src/lightly_studio/examples/example_custom_embeddings.py new file mode 100644 index 0000000000..3f07b3d861 --- /dev/null +++ b/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() diff --git a/lightly_studio/src/lightly_studio/migrations/versions/1783577705_0c5e19664860_add_supports_text_search_to_embedding_.py b/lightly_studio/src/lightly_studio/migrations/versions/1783577705_0c5e19664860_add_supports_text_search_to_embedding_.py new file mode 100644 index 0000000000..7e8f03d2a3 --- /dev/null +++ b/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") diff --git a/lightly_studio/src/lightly_studio/models/embedding_model.py b/lightly_studio/src/lightly_studio/models/embedding_model.py index 7813540cdd..2baa8f9cd3 100644 --- a/lightly_studio/src/lightly_studio/models/embedding_model.py +++ b/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): diff --git a/lightly_studio/src/lightly_studio/resolvers/embedding_model_resolver.py b/lightly_studio/src/lightly_studio/resolvers/embedding_model_resolver.py index 48ee806d8b..aec2368350 100644 --- a/lightly_studio/src/lightly_studio/resolvers/embedding_model_resolver.py +++ b/lightly_studio/src/lightly_studio/resolvers/embedding_model_resolver.py @@ -36,6 +36,7 @@ def get_or_create(session: Session, embedding_model: EmbeddingModelCreate) -> Em db_model.name != embedding_model.name or db_model.parameter_count_in_mb != embedding_model.parameter_count_in_mb or db_model.embedding_dimension != embedding_model.embedding_dimension + or db_model.supports_text_search != embedding_model.supports_text_search ): raise ValueError( "An embedding model with the same hash but different parameters already exists." diff --git a/lightly_studio/src/lightly_studio/resolvers/sample_embedding_resolver.py b/lightly_studio/src/lightly_studio/resolvers/sample_embedding_resolver.py index 50fd8053fb..cf18083f8e 100644 --- a/lightly_studio/src/lightly_studio/resolvers/sample_embedding_resolver.py +++ b/lightly_studio/src/lightly_studio/resolvers/sample_embedding_resolver.py @@ -7,7 +7,6 @@ from typing import Any, NamedTuple from uuid import UUID -from sqlalchemy import func from sqlmodel import Session, col, select from lightly_studio.database import db_vector @@ -42,6 +41,26 @@ def create(session: Session, sample_embedding: SampleEmbeddingCreate) -> SampleE return db_sample_embedding +def upsert(session: Session, sample_embedding: SampleEmbeddingCreate) -> SampleEmbeddingTable: + """Create a sample embedding, or update it in place if one already exists. + + Existing rows are looked up by the ``(sample_id, embedding_model_id)`` primary key, so + callers can safely call this repeatedly (e.g. re-running an ingestion script) without + hitting a primary key conflict. + """ + existing = session.get( + SampleEmbeddingTable, + (sample_embedding.sample_id, sample_embedding.embedding_model_id), + ) + if existing is None: + return create(session=session, sample_embedding=sample_embedding) + existing.embedding = sample_embedding.embedding + session.add(existing) + session.commit() + session.refresh(existing) + return existing + + def create_many( session: Session, sample_embeddings: list[SampleEmbeddingCreate], commit: bool = True ) -> None: @@ -205,24 +224,33 @@ def get_hash_by_collection_id( return hasher.hexdigest(), sample_ids -def get_embedding_count(session: Session, collection_id: UUID, embedding_model_id: UUID) -> int: - """Get the number of sample embeddings for samples in a specific collection. +def has_any_embedding( + session: Session, + collection_id: UUID, + embedding_model_ids: Sequence[UUID] | None = None, +) -> bool: + """Check whether at least one sample embedding exists for a collection. Args: session: The database session. collection_id: The collection ID to filter by. - embedding_model_id: The embedding model ID to filter by. + embedding_model_ids: If given, only embeddings under these models count. + An empty sequence always returns False. Returns: - The number of sample embeddings associated with the collection. + True if at least one matching sample embedding exists. """ + if embedding_model_ids is not None and not embedding_model_ids: + return False query = ( - select(func.count(col(SampleEmbeddingTable.sample_id))) + select(SampleEmbeddingTable.sample_id) .join(SampleTable, col(SampleEmbeddingTable.sample_id) == col(SampleTable.sample_id)) .where(SampleTable.collection_id == collection_id) - .where(SampleEmbeddingTable.embedding_model_id == embedding_model_id) + .limit(1) ) - return session.exec(query).one() + if embedding_model_ids is not None: + query = query.where(col(SampleEmbeddingTable.embedding_model_id).in_(embedding_model_ids)) + return session.exec(query).first() is not None def _read_embedding_rows_binary( diff --git a/lightly_studio/tests/api/routes/api/test_collection.py b/lightly_studio/tests/api/routes/api/test_collection.py index 9dd5e04846..f32d82b671 100644 --- a/lightly_studio/tests/api/routes/api/test_collection.py +++ b/lightly_studio/tests/api/routes/api/test_collection.py @@ -4,7 +4,6 @@ import pytest from fastapi.testclient import TestClient -from pytest_mock import MockerFixture from sqlmodel import Session from lightly_studio.api.routes.api.status import ( @@ -14,7 +13,6 @@ HTTP_STATUS_NOT_FOUND, HTTP_STATUS_OK, ) -from lightly_studio.dataset.embedding_manager import EmbeddingManager from lightly_studio.models.collection import SampleType from tests.helpers_resolvers import ( ImageStub, @@ -215,22 +213,16 @@ def test_read_collections_overview(test_client: TestClient, db_session: Session) def test_has_embeddings( test_client: TestClient, db_session: Session, - mocker: MockerFixture, ) -> None: col_id = create_collection(session=db_session).collection_id embedding_model_id = create_embedding_model( session=db_session, collection_id=col_id ).embedding_model_id - mock_get_model = mocker.patch.object( - EmbeddingManager, "load_or_get_default_model", return_value=embedding_model_id - ) # Initially, the collection has no embeddings. response = test_client.get(f"/api/collections/{col_id!s}/has_embeddings") assert response.status_code == HTTP_STATUS_OK - assert response.json() is False - mock_get_model.assert_called_once_with(session=db_session, collection_id=col_id) - mock_get_model.reset_mock() + assert response.json() == {"has_embeddings": False, "has_text_search_embeddings": False} # Add an embedding to the collection. create_samples_with_embeddings( @@ -243,8 +235,27 @@ def test_has_embeddings( # Now, the collection should report having embeddings. response = test_client.get(f"/api/collections/{col_id!s}/has_embeddings") assert response.status_code == HTTP_STATUS_OK - assert response.json() is True - mock_get_model.assert_called_once_with(session=db_session, collection_id=col_id) + assert response.json() == {"has_embeddings": True, "has_text_search_embeddings": True} + + +def test_has_embeddings__custom_model( + test_client: TestClient, + db_session: Session, +) -> None: + col_id = create_collection(session=db_session).collection_id + embedding_model_id = create_embedding_model( + session=db_session, collection_id=col_id, supports_text_search=False + ).embedding_model_id + create_samples_with_embeddings( + session=db_session, + collection_id=col_id, + embedding_model_id=embedding_model_id, + images_and_embeddings=[(ImageStub(), [0.1, 0.2, 0.3])], + ) + + response = test_client.get(f"/api/collections/{col_id!s}/has_embeddings") + assert response.status_code == HTTP_STATUS_OK + assert response.json() == {"has_embeddings": True, "has_text_search_embeddings": False} @pytest.mark.postgres_only # Deep copying is enterprise-only (PostgreSQL-backed). diff --git a/lightly_studio/tests/core/image/test_custom_embeddings_integration.py b/lightly_studio/tests/core/image/test_custom_embeddings_integration.py new file mode 100644 index 0000000000..ce5bc367ab --- /dev/null +++ b/lightly_studio/tests/core/image/test_custom_embeddings_integration.py @@ -0,0 +1,61 @@ +"""Integration test for the "bring your own embeddings" workflow. + +Guards the contract documented in `docs/docs/dataset_setup/image_dataset.md`: loading a +collection with `embed=False` and setting embeddings via `Sample.set_embedding(...)` must +(1) keep the custom embedding model as the collection's default (so the embedding plot +resolves to it) and (2) report `has_embeddings=True` but `has_text_search_embeddings=False` +(so the frontend shows the embedding plot but hides the text-search bar). +""" + +from __future__ import annotations + +from lightly_studio.core.image.image_sample import ImageSample +from lightly_studio.database import db_manager +from lightly_studio.dataset import embedding_utils +from lightly_studio.resolvers import embedding_model_resolver +from lightly_studio.resolvers.twodim_embedding_resolver import get_twodim_embeddings +from tests.helpers_resolvers import create_collection, create_image + + +def test_custom_embeddings_stay_default_and_hide_text_search( + patch_collection: None, # noqa: ARG001 +) -> None: + session = db_manager.persistent_session() + collection = create_collection(session=session) + image_table = create_image(session=session, collection_id=collection.collection_id) + sample = ImageSample(inner=image_table) + + # Simulate the recommended workflow: populate custom embeddings before ever + # touching anything that could auto-register the built-in default model. + sample.set_embedding([0.1, 0.2, 0.3]) + + models = embedding_model_resolver.get_all_by_collection_id( + session=session, collection_id=collection.collection_id + ) + assert [m.name for m in models] == ["custom"] + assert models[0].supports_text_search is False + + # The status the frontend uses to gate features: the embedding plot must be + # available, text search must not. + status = embedding_utils.get_collection_embeddings_status( + session=session, collection_id=collection.collection_id + ) + assert status.has_embeddings is True + assert status.has_text_search_embeddings is False + + # The custom model must be the one the embedding plot resolves to. + default_model = embedding_model_resolver.get_default_by_collection_id( + session=session, collection_id=collection.collection_id + ) + assert default_model is not None + assert default_model.name == "custom" + + # The embedding plot must resolve against the custom model and return one point. + x_values, y_values, sample_ids = get_twodim_embeddings( + session=session, + collection_id=collection.collection_id, + embedding_model_id=default_model.embedding_model_id, + ) + assert len(x_values) == 1 + assert len(y_values) == 1 + assert sample_ids == [sample.sample_id] diff --git a/lightly_studio/tests/core/image/test_image_sample.py b/lightly_studio/tests/core/image/test_image_sample.py index a52cc7ae6e..3f0ca4540a 100644 --- a/lightly_studio/tests/core/image/test_image_sample.py +++ b/lightly_studio/tests/core/image/test_image_sample.py @@ -13,7 +13,11 @@ from lightly_studio.core.annotation.segmentation_mask import SegmentationMaskAnnotation from lightly_studio.core.image.image_sample import ImageSample from lightly_studio.models.annotation.annotation_base import AnnotationType -from lightly_studio.resolvers import image_resolver +from lightly_studio.resolvers import ( + embedding_model_resolver, + image_resolver, + sample_embedding_resolver, +) from tests.helpers_resolvers import ( create_annotation, create_annotation_label, @@ -641,3 +645,42 @@ def test_delete_annotation( # Verify it's gone. assert len(image.annotations) == 0 + + def test_set_embedding(self, db_session: Session) -> None: + collection = create_collection(session=db_session) + image_table = create_image( + session=db_session, + collection_id=collection.collection_id, + ) + sample = ImageSample(inner=image_table) + + sample.set_embedding([0.1, 0.2, 0.3]) + + models = embedding_model_resolver.get_all_by_collection_id( + session=db_session, collection_id=collection.collection_id + ) + assert len(models) == 1 + assert models[0].name == "custom" + assert models[0].embedding_dimension == 3 + + rows = sample_embedding_resolver.get_by_sample_ids( + session=db_session, + sample_ids=[sample.sample_id], + embedding_model_id=models[0].embedding_model_id, + ) + assert len(rows) == 1 + assert list(rows[0].embedding) == pytest.approx([0.1, 0.2, 0.3]) + + # Setting a new vector under the same name updates the row in place. + sample.set_embedding([0.4, 0.5, 0.6]) + models = embedding_model_resolver.get_all_by_collection_id( + session=db_session, collection_id=collection.collection_id + ) + assert len(models) == 1 + rows = sample_embedding_resolver.get_by_sample_ids( + session=db_session, + sample_ids=[sample.sample_id], + embedding_model_id=models[0].embedding_model_id, + ) + assert len(rows) == 1 + assert list(rows[0].embedding) == pytest.approx([0.4, 0.5, 0.6]) diff --git a/lightly_studio/tests/dataset/test_embedding_utils.py b/lightly_studio/tests/dataset/test_embedding_utils.py index 5bfbd49ebc..789a0d73a5 100644 --- a/lightly_studio/tests/dataset/test_embedding_utils.py +++ b/lightly_studio/tests/dataset/test_embedding_utils.py @@ -1,10 +1,8 @@ from __future__ import annotations -from pytest_mock import MockerFixture from sqlmodel import Session from lightly_studio.dataset import embedding_utils -from lightly_studio.dataset.embedding_manager import EmbeddingManager from tests.helpers_resolvers import ( ImageStub, create_collection, @@ -13,27 +11,32 @@ ) -def test_collection_has_embeddings(db_session: Session, mocker: MockerFixture) -> None: +def test_get_collection_embeddings_status__no_models(db_session: Session) -> None: col_id = create_collection(session=db_session).collection_id - embedding_model_id = create_embedding_model( + + status = embedding_utils.get_collection_embeddings_status( session=db_session, collection_id=col_id - ).embedding_model_id - mock_get_model = mocker.patch.object( - EmbeddingManager, "load_or_get_default_model", return_value=embedding_model_id ) + assert status.has_embeddings is False + assert status.has_text_search_embeddings is False - # Initially, the collection has no embeddings. - assert not embedding_utils.collection_has_embeddings( - session=db_session, - collection_id=col_id, - ) - mock_get_model.assert_called_once_with( - session=db_session, - collection_id=col_id, + +def test_get_collection_embeddings_status__model_without_embeddings(db_session: Session) -> None: + col_id = create_collection(session=db_session).collection_id + create_embedding_model(session=db_session, collection_id=col_id) + + status = embedding_utils.get_collection_embeddings_status( + session=db_session, collection_id=col_id ) - mock_get_model.reset_mock() + assert status.has_embeddings is False + assert status.has_text_search_embeddings is False - # Add an embedding to the collection. + +def test_get_collection_embeddings_status__text_capable_model(db_session: Session) -> None: + col_id = create_collection(session=db_session).collection_id + embedding_model_id = create_embedding_model( + session=db_session, collection_id=col_id + ).embedding_model_id create_samples_with_embeddings( session=db_session, collection_id=col_id, @@ -41,12 +44,27 @@ def test_collection_has_embeddings(db_session: Session, mocker: MockerFixture) - images_and_embeddings=[(ImageStub(), [0.1, 0.2, 0.3])], ) - # Now, the collection should report having embeddings. - assert embedding_utils.collection_has_embeddings( - session=db_session, - collection_id=col_id, + status = embedding_utils.get_collection_embeddings_status( + session=db_session, collection_id=col_id ) - mock_get_model.assert_called_once_with( + assert status.has_embeddings is True + assert status.has_text_search_embeddings is True + + +def test_get_collection_embeddings_status__custom_model_only(db_session: Session) -> None: + col_id = create_collection(session=db_session).collection_id + embedding_model_id = create_embedding_model( + session=db_session, collection_id=col_id, supports_text_search=False + ).embedding_model_id + create_samples_with_embeddings( session=db_session, collection_id=col_id, + embedding_model_id=embedding_model_id, + images_and_embeddings=[(ImageStub(), [0.1, 0.2, 0.3])], + ) + + status = embedding_utils.get_collection_embeddings_status( + session=db_session, collection_id=col_id ) + assert status.has_embeddings is True + assert status.has_text_search_embeddings is False diff --git a/lightly_studio/tests/helpers_resolvers.py b/lightly_studio/tests/helpers_resolvers.py index 46408b00f9..bd0dce3edd 100644 --- a/lightly_studio/tests/helpers_resolvers.py +++ b/lightly_studio/tests/helpers_resolvers.py @@ -298,6 +298,7 @@ def create_embedding_model( # noqa: PLR0913 embedding_model_hash: str = "example_hash", parameter_count_in_mb: int = 100, embedding_dimension: int = 128, + supports_text_search: bool = True, ) -> EmbeddingModelTable: """Helper function to create a embedding model.""" return embedding_model_resolver.create( @@ -308,6 +309,7 @@ def create_embedding_model( # noqa: PLR0913 embedding_model_hash=embedding_model_hash, parameter_count_in_mb=parameter_count_in_mb, embedding_dimension=embedding_dimension, + supports_text_search=supports_text_search, ), ) diff --git a/lightly_studio/tests/resolvers/test_sample_embedding_resolver.py b/lightly_studio/tests/resolvers/test_sample_embedding_resolver.py index 4465ca6c2c..98244ce7f1 100644 --- a/lightly_studio/tests/resolvers/test_sample_embedding_resolver.py +++ b/lightly_studio/tests/resolvers/test_sample_embedding_resolver.py @@ -239,7 +239,7 @@ def test_get_all_by_collection_id_with_filter(db_session: Session) -> None: assert embedding_by_id[samples[1].sample_id] == [1.0, 2.0, 3.0] -def test_get_embedding_count(db_session: Session) -> None: +def test_has_any_embedding(db_session: Session) -> None: # Create collections col1_id = create_collection(session=db_session).collection_id col2_id = create_collection(session=db_session, collection_name="col2").collection_id @@ -248,43 +248,47 @@ def test_get_embedding_count(db_session: Session) -> None: images_col1 = create_images( db_session=db_session, collection_id=col1_id, - images=[ImageStub("sample1.png"), ImageStub("sample2.png"), ImageStub("sample3.png")], + images=[ImageStub("sample1.png"), ImageStub("sample2.png")], ) create_images(db_session=db_session, collection_id=col2_id, images=[ImageStub("sample.png")]) - # Create an embedding models + # Create embedding models embedding_model_1 = create_embedding_model(session=db_session, collection_id=col1_id) embedding_model_1_id = embedding_model_1.embedding_model_id embedding_model_2 = create_embedding_model(session=db_session, collection_id=col1_id) embedding_model_2_id = embedding_model_2.embedding_model_id - # Create embeddings for col1 - embedding_inputs = [ - SampleEmbeddingCreate( - sample_id=images_col1[0].sample_id, - embedding_model_id=embedding_model_1_id, - embedding=np.array([0.0, 0.0, 0.0], dtype=np.float32), - ), - SampleEmbeddingCreate( - sample_id=images_col1[1].sample_id, - embedding_model_id=embedding_model_1_id, - embedding=np.array([0.0, 0.0, 0.0], dtype=np.float32), - ), - ] - sample_embedding_resolver.create_many(session=db_session, sample_embeddings=embedding_inputs) + # Create embeddings for col1 under model 1 only. + sample_embedding_resolver.create_many( + session=db_session, + sample_embeddings=[ + SampleEmbeddingCreate( + sample_id=images_col1[0].sample_id, + embedding_model_id=embedding_model_1_id, + embedding=np.array([0.0, 0.0, 0.0], dtype=np.float32), + ), + ], + ) + + # Collection 1 has embeddings, collection 2 does not. + assert sample_embedding_resolver.has_any_embedding(session=db_session, collection_id=col1_id) + assert not sample_embedding_resolver.has_any_embedding( + session=db_session, collection_id=col2_id + ) - # Collection 1 has two embeddings - count = sample_embedding_resolver.get_embedding_count( + # Filtering by model only counts embeddings under the given models. + assert sample_embedding_resolver.has_any_embedding( session=db_session, collection_id=col1_id, - embedding_model_id=embedding_model_1_id, + embedding_model_ids=[embedding_model_1_id], ) - assert count == 2 - - # Collection 2 has no embeddings - count = sample_embedding_resolver.get_embedding_count( + assert not sample_embedding_resolver.has_any_embedding( session=db_session, - collection_id=col2_id, - embedding_model_id=embedding_model_2_id, + collection_id=col1_id, + embedding_model_ids=[embedding_model_2_id], + ) + + # An empty model list never matches. + assert not sample_embedding_resolver.has_any_embedding( + session=db_session, collection_id=col1_id, embedding_model_ids=[] ) - assert count == 0 diff --git a/lightly_studio_view/src/lib/components/AnnotationsGrid/AnnotationsGrid.svelte b/lightly_studio_view/src/lib/components/AnnotationsGrid/AnnotationsGrid.svelte index 48ee21df9c..a9c46ee023 100644 --- a/lightly_studio_view/src/lib/components/AnnotationsGrid/AnnotationsGrid.svelte +++ b/lightly_studio_view/src/lib/components/AnnotationsGrid/AnnotationsGrid.svelte @@ -64,9 +64,11 @@ const plotSelectedRegion = $derived($annotationPlotRegion); // The text embedding search is shared with the images tab and persists across the tab switch. - // Only apply it when this annotation collection actually has embeddings. + // Only apply it when this annotation collection has embeddings from a text-capable model. const hasEmbeddingsQuery = useHasEmbeddings(() => ({ collectionId: collection_id })); - const searchEmbedding = $derived(hasEmbeddingsQuery.data ? $textEmbedding : undefined); + const searchEmbedding = $derived( + hasEmbeddingsQuery.data?.has_text_search_embeddings ? $textEmbedding : undefined + ); // Drag-to-search crop preview. Tiles report only their crop geometry; the blob is // rendered lazily when a drag starts (not per visible tile), and revoked on unmount. diff --git a/lightly_studio_view/src/lib/components/Header/Header.svelte b/lightly_studio_view/src/lib/components/Header/Header.svelte index aee66a3622..a6c8857f25 100644 --- a/lightly_studio_view/src/lib/components/Header/Header.svelte +++ b/lightly_studio_view/src/lib/components/Header/Header.svelte @@ -23,7 +23,7 @@ const { settingsStore } = useSettings(); const hasEmbeddingsQuery = useHasEmbeddings(() => ({ collectionId: collection.collection_id })); - const hasEmbeddings = $derived(!!hasEmbeddingsQuery.data); + const hasEmbeddings = $derived(!!hasEmbeddingsQuery.data?.has_embeddings); const datasetId = $derived(page.params.dataset_id!); const { collection: datasetCollection } = $derived.by(() => diff --git a/lightly_studio_view/src/lib/components/Header/Header.test.ts b/lightly_studio_view/src/lib/components/Header/Header.test.ts index 41ee694164..59e9094288 100644 --- a/lightly_studio_view/src/lib/components/Header/Header.test.ts +++ b/lightly_studio_view/src/lib/components/Header/Header.test.ts @@ -85,7 +85,7 @@ describe('Header', () => { useHasEmbeddings.mockReturnValue( readable({ - data: true, + data: { has_embeddings: true, has_text_search_embeddings: true }, isLoading: false, error: null }) diff --git a/lightly_studio_view/src/lib/hooks/useHasEmbeddings/useHasEmbeddings.ts b/lightly_studio_view/src/lib/hooks/useHasEmbeddings/useHasEmbeddings.ts index d961e374eb..a1cb299b57 100644 --- a/lightly_studio_view/src/lib/hooks/useHasEmbeddings/useHasEmbeddings.ts +++ b/lightly_studio_view/src/lib/hooks/useHasEmbeddings/useHasEmbeddings.ts @@ -1,9 +1,10 @@ import { hasEmbeddingsOptions } from '$lib/api/lightly_studio_local/@tanstack/svelte-query.gen'; +import type { CollectionEmbeddingsStatus } from '$lib/api/lightly_studio_local'; import { createQuery, type CreateQueryResult } from '@tanstack/svelte-query'; export const useHasEmbeddings = ( getParams: () => { collectionId: string } -): CreateQueryResult => { +): CreateQueryResult => { return createQuery(() => hasEmbeddingsOptions({ path: { collection_id: getParams().collectionId } diff --git a/lightly_studio_view/src/routes/datasets/[dataset_id]/[collection_type]/[collection_id]/+layout.svelte b/lightly_studio_view/src/routes/datasets/[dataset_id]/[collection_type]/[collection_id]/+layout.svelte index a93d84cf75..be843161cc 100644 --- a/lightly_studio_view/src/routes/datasets/[dataset_id]/[collection_type]/[collection_id]/+layout.svelte +++ b/lightly_studio_view/src/routes/datasets/[dataset_id]/[collection_type]/[collection_id]/+layout.svelte @@ -259,10 +259,14 @@ }); const hasEmbeddingsQuery = useHasEmbeddings(() => ({ collectionId })); - const hasEmbeddings = $derived(!!hasEmbeddingsQuery.data); + const hasEmbeddings = $derived(!!hasEmbeddingsQuery.data?.has_embeddings); + const hasTextSearchEmbeddings = $derived(!!hasEmbeddingsQuery.data?.has_text_search_embeddings); const hasMediaWithEmbeddings = $derived( (isImages || isVideos || isAnnotations) && hasEmbeddings ); + const hasMediaWithTextSearch = $derived( + (isImages || isVideos || isAnnotations) && hasTextSearchEmbeddings + ); const collectionSearchPlaceholder = $derived( isAnnotations ? 'Search annotations by description or image' @@ -444,7 +448,7 @@ {canSelectAll} isSelectionActive={$selectedCount > 0} {isImages} - {hasMediaWithEmbeddings} + hasMediaWithEmbeddings={hasMediaWithTextSearch} collectionDatasetId={collection.dataset_id} onSelectAll={selectAllHandle.handleSelectAll} onDeselectAll={clearSelection}