Add embedding to video dataset (#274)

Author: JonasWurst

lightly_studio/src/lightly_studio/api/routes/api/text_embedding.py

@@ -16,6 +16,8 @@
     EmbeddingManagerProvider,
     TextEmbedQuery,
 )
+from lightly_studio.db_manager import SessionDep
+from lightly_studio.resolvers import dataset_resolver
 
 text_embedding_router = APIRouter()
 # Define a type alias for the EmbeddingManager dependency
@@ -27,6 +29,7 @@
 
 @text_embedding_router.get("/text_embedding/embed_text", response_model=List[float])
 def embed_text(
+    session: SessionDep,
     embedding_manager: EmbeddingManagerDep,
     query_text: str = Query(..., description="The text to embed."),
     embedding_model_id: Annotated[
@@ -35,9 +38,16 @@ def embed_text(
     ] = None,
 ) -> list[float]:
     """Retrieve embeddings for the input text."""
+    # TODO(Jonas, 12/2025): Remove this hack after dataset_id is provided from frontend
+    # This is a hack, since at the moment, no valid embedding_model_id is passed from the frontend.
+    # so we fetch the root_dataset_id, which will be used inside embed_text to get the default model
+    # for this dataset.
+    root_dataset = dataset_resolver.get_root_dataset(session=session)
+    dataset_id = root_dataset.dataset_id
     try:
         text_embeddings = embedding_manager.embed_text(
-            TextEmbedQuery(query_text, embedding_model_id)
+            dataset_id=dataset_id,
+            text_query=TextEmbedQuery(text=query_text, embedding_model_id=embedding_model_id),
         )
     except ValueError as exc:
         raise HTTPException(

lightly_studio/src/lightly_studio/core/dataset.py

@@ -271,6 +271,7 @@ def add_videos_from_path(
         path: PathLike,
         allowed_extensions: Iterable[str] | None = None,
         num_decode_threads: int | None = None,
+        embed: bool = True,
     ) -> None:
         """Adding video frames from the specified path to the dataset.
 
@@ -281,6 +282,7 @@ def add_videos_from_path(
             uses default VIDEO_EXTENSIONS.
             num_decode_threads: Optional override for the number of FFmpeg decode threads.
                 If omitted, the available CPU cores - 1 (max 16) are used.
+            embed: If True, generate embeddings for the newly added videos.
         """
         # Collect video file paths.
         if allowed_extensions:
@@ -295,13 +297,20 @@ def add_videos_from_path(
         logger.info(f"Found {len(video_paths)} videos in {path}.")
 
         # Process videos.
-        add_videos.load_into_dataset_from_paths(
+        created_sample_ids, _ = add_videos.load_into_dataset_from_paths(
             session=self.session,
             dataset_id=self.dataset_id,
             video_paths=video_paths,
             num_decode_threads=num_decode_threads,
         )
 
+        if embed:
+            _generate_embeddings_video(
+                session=self.session,
+                dataset_id=self.dataset_id,
+                sample_ids=created_sample_ids,
+            )
+
     def add_images_from_path(
         self,
         path: PathLike,
@@ -354,7 +363,7 @@ def add_images_from_path(
             )
 
         if embed:
-            _generate_embeddings(
+            _generate_embeddings_image(
                 session=self.session, dataset_id=self.dataset_id, sample_ids=created_sample_ids
             )
 
@@ -383,7 +392,7 @@ def add_samples_from_labelformat(
         )
 
         if embed:
-            _generate_embeddings(
+            _generate_embeddings_image(
                 session=self.session, dataset_id=self.dataset_id, sample_ids=created_sample_ids
             )
 
@@ -446,7 +455,7 @@ def add_samples_from_yolo(
 
         # Generate embeddings for all samples at once
         if embed:
-            _generate_embeddings(
+            _generate_embeddings_image(
                 session=self.session, dataset_id=self.dataset_id, sample_ids=all_created_sample_ids
             )
 
@@ -512,7 +521,7 @@ def add_samples_from_coco(
             )
 
         if embed:
-            _generate_embeddings(
+            _generate_embeddings_image(
                 session=self.session, dataset_id=self.dataset_id, sample_ids=created_sample_ids
             )
 
@@ -564,7 +573,7 @@ def add_samples_from_coco_caption(
             )
 
         if embed:
-            _generate_embeddings(
+            _generate_embeddings_image(
                 session=self.session, dataset_id=self.dataset_id, sample_ids=created_sample_ids
             )
 
@@ -635,7 +644,11 @@ def compute_similarity_metadata(
         )
 
 
-def _generate_embeddings(session: Session, dataset_id: UUID, sample_ids: list[UUID]) -> None:
+def _generate_embeddings_video(
+    session: Session,
+    dataset_id: UUID,
+    sample_ids: list[UUID],
+) -> None:
     """Generate and store embeddings for samples.
 
     Args:
@@ -647,20 +660,54 @@ def _generate_embeddings(session: Session, dataset_id: UUID, sample_ids: list[UU
         return
 
     embedding_manager = EmbeddingManagerProvider.get_embedding_manager()
-    model_id = embedding_manager.load_or_get_default_model(
+    model_id = embedding_manager.load_or_get_default_model(session=session, dataset_id=dataset_id)
+    if model_id is None:
+        logger.warning("No embedding model loaded. Skipping embedding generation.")
+        return
+
+    embedding_manager.embed_videos(
         session=session,
         dataset_id=dataset_id,
+        sample_ids=sample_ids,
+        embedding_model_id=model_id,
     )
+
+    _mark_embedding_features_enabled()
+
+
+def _generate_embeddings_image(
+    session: Session,
+    dataset_id: UUID,
+    sample_ids: list[UUID],
+) -> None:
+    """Generate and store embeddings for samples.
+
+    Args:
+        session: Database session for resolver operations.
+        dataset_id: The ID of the dataset to associate with the embedding model.
+        sample_ids: List of sample IDs to generate embeddings for.
+        sample_type: The sample_type to generate embeddings for.
+    """
+    if not sample_ids:
+        return
+
+    embedding_manager = EmbeddingManagerProvider.get_embedding_manager()
+    model_id = embedding_manager.load_or_get_default_model(session=session, dataset_id=dataset_id)
     if model_id is None:
         logger.warning("No embedding model loaded. Skipping embedding generation.")
         return
 
     embedding_manager.embed_images(
         session=session,
+        dataset_id=dataset_id,
         sample_ids=sample_ids,
         embedding_model_id=model_id,
     )
 
+    _mark_embedding_features_enabled()
+
+
+def _mark_embedding_features_enabled() -> None:
     # Mark the embedding search feature as enabled.
     if "embeddingSearchEnabled" not in features.lightly_studio_active_features:
         features.lightly_studio_active_features.append("embeddingSearchEnabled")