Embed video frames during extraction when embed_frames is enabled.

lightly_studio/src/lightly_studio/core/video/add_videos.py

@@ -30,6 +30,7 @@
     SingleObjectDetectionTrack,
     VideoObjectDetectionTrack,
 )
+from PIL import Image
 from sqlmodel import Session
 from tqdm import tqdm
 
@@ -40,6 +41,7 @@
     FileOutcomeReport,
     MissingInputFileError,
 )
+from lightly_studio.dataset.embedding_manager import EmbeddingManagerProvider
 from lightly_studio.models.annotation.annotation_base import (
     AnnotationCreate,
 )
@@ -59,7 +61,7 @@
 
 DEFAULT_VIDEO_CHANNEL = 0
 # Number of samples to process in a single batch
-SAMPLE_BATCH_SIZE = 128
+SAMPLE_BATCH_SIZE = 32
 
 # Video file extensions
 # These are commonly supported by PyAV/FFmpeg.
@@ -81,6 +83,8 @@ class FrameExtractionContext:
     session: Session
     collection_id: UUID
     video_sample_id: UUID
+    embed_frames: bool = False
+    embedding_model_id: UUID | None = None
 
 
 @dataclass
@@ -93,6 +97,8 @@ class VideoLoadContext:
     video_channel: int
     num_decode_threads: int | None
     target_fps: float | None
+    embed_frames: bool
+    embedding_model_id: UUID | None
 
 
 def load_into_collection_from_paths(  # noqa: PLR0913
@@ -103,6 +109,7 @@ def load_into_collection_from_paths(  # noqa: PLR0913
     num_decode_threads: int | None = None,
     show_progress: bool = True,
     target_fps: float | None = None,
+    embed_frames: bool = False,
 ) -> tuple[list[UUID], list[UUID]]:
     """Load video samples from file paths into the dataset using PyAV.
 
@@ -118,6 +125,8 @@ def load_into_collection_from_paths(  # noqa: PLR0913
         target_fps: Optional target frame rate for subsampling. When set below the source
             frame rate, only selected frames are kept. frame_number values remain
             original. Must be greater than 0.
+        embed_frames: If True, generate image embeddings for extracted video frames during
+            decoding. Requires an image-compatible embedding model.
 
     Returns:
         A tuple containing:
@@ -143,6 +152,16 @@ def load_into_collection_from_paths(  # noqa: PLR0913
     video_frames_collection_id = collection_resolver.get_or_create_child_collection(
         session=session, collection_id=collection_id, sample_type=SampleType.VIDEO_FRAME
     )
+    embedding_model_id: UUID | None = None
+    if embed_frames:
+        embedding_manager = EmbeddingManagerProvider.get_embedding_manager()
+        embedding_model_id = embedding_manager.load_or_get_default_model(
+            session=session,
+            collection_id=video_frames_collection_id,
+        )
+        if embedding_model_id is None:
+            logger.warning("No embedding model loaded. Skipping frame embedding generation.")
+    effective_embed_frames = embed_frames and embedding_model_id is not None
 
     load_context = VideoLoadContext(
         session=session,
@@ -151,6 +170,8 @@ def load_into_collection_from_paths(  # noqa: PLR0913
         video_channel=video_channel,
         num_decode_threads=num_decode_threads,
         target_fps=target_fps,
+        embed_frames=effective_embed_frames,
+        embedding_model_id=embedding_model_id,
     )
 
     for video_path in tqdm(
@@ -246,6 +267,8 @@ def _load_single_video(
                 session=context.session,
                 collection_id=context.video_frames_collection_id,
                 video_sample_id=video_sample_ids[0],
+                embed_frames=context.embed_frames,
+                embedding_model_id=context.embedding_model_id,
             )
             frame_sample_ids = _create_video_frame_samples(
                 context=extraction_context,
@@ -272,6 +295,7 @@ def load_video_annotations_from_labelformat(  # noqa: PLR0913
     input_labels: ObjectDetectionTrackInput | InstanceSegmentationTrackInput,
     input_labels_paths_root: Path | str,
     limit: int | None = None,
+    embed_frames: bool = False,
 ) -> tuple[list[UUID], list[UUID]]:
     """Load video frame annotations from a labelformat input into the dataset.
 
@@ -292,6 +316,8 @@ def load_video_annotations_from_labelformat(  # noqa: PLR0913
         input_labels_paths_root: The root path for the paths in input_labels.
         limit: Maximum number of samples to load. By default, all samples are loaded.
             Annotations of videos beyond the limit are skipped.
+        embed_frames: If True, generate image embeddings for extracted video frames during
+            decoding. Requires an image-compatible embedding model.
 
     Returns:
         A tuple containing:
@@ -308,6 +334,7 @@ def load_video_annotations_from_labelformat(  # noqa: PLR0913
         session=session,
         collection_id=collection_id,
         video_paths=video_paths_labelformat,
+        embed_frames=embed_frames,
     )
 
     # In YouTube-VIS, the file extension is typically missing. Hence we fallback to the path
@@ -431,7 +458,8 @@ def _create_video_frame_samples(
 ) -> list[UUID]:
     """Create video frame samples for a video by parsing all frames.
 
-    This function decodes all frames to extract metadata.
+    This function decodes all frames to extract metadata. When frame embedding is enabled,
+    embeddings are generated from the decoded frames in the same pass.
 
     Args:
         context: Frame extraction context (session, dataset and parent video).
@@ -447,6 +475,7 @@ def _create_video_frame_samples(
     """
     created_sample_ids: list[UUID] = []
     samples_to_create: list[VideoFrameCreate] = []
+    pil_frames: list[Image.Image] = []
     video_stream = video_container.streams.video[video_channel]
     _configure_stream_threading(video_stream=video_stream, num_decode_threads=num_decode_threads)
 
@@ -477,25 +506,53 @@ def _create_video_frame_samples(
             rotation_deg=_get_frame_rotation_deg(frame=frame),
         )
         samples_to_create.append(sample)
+        if context.embed_frames:
+            pil_frames.append(frame.to_image().convert("RGB"))  # type: ignore[no-untyped-call]
 
-        # Process batch when it reaches SAMPLE_BATCH_SIZE
         if len(samples_to_create) >= SAMPLE_BATCH_SIZE:
-            created_samples_batch = video_frame_resolver.create_many(
-                session=context.session,
-                samples=samples_to_create,
-                collection_id=context.collection_id,
+            created_sample_ids.extend(
+                _flush_frame_batch(
+                    context=context,
+                    samples_to_create=samples_to_create,
+                    pil_frames=pil_frames,
+                )
             )
-            created_sample_ids.extend(created_samples_batch)
             samples_to_create = []
+            pil_frames = []
 
-    # Handle remaining samples for this video
     if samples_to_create:
-        created_samples_batch = video_frame_resolver.create_many(
+        created_sample_ids.extend(
+            _flush_frame_batch(
+                context=context,
+                samples_to_create=samples_to_create,
+                pil_frames=pil_frames,
+            )
+        )
+
+    return created_sample_ids
+
+
+def _flush_frame_batch(
+    context: FrameExtractionContext,
+    samples_to_create: list[VideoFrameCreate],
+    pil_frames: list[Image.Image],
+) -> list[UUID]:
+    """Persist a batch of frame samples and optionally embed them."""
+    created_sample_ids = video_frame_resolver.create_many(
+        session=context.session,
+        samples=samples_to_create,
+        collection_id=context.collection_id,
+    )
+
+    if context.embed_frames and context.embedding_model_id is not None and pil_frames:
+        embedding_manager = EmbeddingManagerProvider.get_embedding_manager()
+        embedding_manager.embed_and_store_pil_images(
             session=context.session,
-            samples=samples_to_create,
-            collection_id=context.collection_id,
+            embedding_model_id=context.embedding_model_id,
+            sample_ids=created_sample_ids,
+            images=pil_frames,
+            show_progress=False,
         )
-        created_sample_ids.extend(created_samples_batch)
 
     return created_sample_ids
 

lightly_studio/src/lightly_studio/dataset/embedding_generator.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
+from PIL import Image
 
 from lightly_studio.models.embedding_model import EmbeddingModelCreate
 
@@ -94,6 +95,21 @@ def embed_image_crops(
         """
         ...
 
+    def embed_pil_images(
+        self, images: list[Image.Image], show_progress: bool = True
+    ) -> NDArray[np.float32]:
+        """Generate embeddings for in-memory PIL images.
+
+        Args:
+            images: PIL images to embed.
+            show_progress: Whether to show a progress bar during embedding.
+
+        Returns:
+            A numpy array representing the generated embeddings in the same order
+            as the input images.
+        """
+        ...
+
 
 @runtime_checkable
 class VideoEmbeddingGenerator(EmbeddingGenerator, Protocol):
@@ -160,6 +176,13 @@ def embed_image_crops(
         _ = show_progress  # Not used for random embeddings.
         return np.random.rand(len(image_crops), self._dimension).astype(np.float32)
 
+    def embed_pil_images(
+        self, images: list[Image.Image], show_progress: bool = True
+    ) -> NDArray[np.float32]:
+        """Generate random embeddings for in-memory PIL images."""
+        _ = show_progress  # Not used for random embeddings.
+        return np.random.rand(len(images), self._dimension).astype(np.float32)
+
     def embed_videos(self, filepaths: list[str]) -> NDArray[np.float32]:
         """Generate random embeddings for multiple video samples."""
         return np.random.rand(len(filepaths), self._dimension).astype(np.float32)

lightly_studio/src/lightly_studio/dataset/embedding_manager.py

@@ -8,6 +8,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
+from PIL import Image
 from sqlmodel import Session
 from tqdm import tqdm
 
@@ -43,6 +44,7 @@
     SampleType.IMAGE: SampleType.IMAGE,
     SampleType.ANNOTATION: SampleType.IMAGE,
     SampleType.VIDEO: SampleType.VIDEO,
+    SampleType.VIDEO_FRAME: SampleType.IMAGE,
 }
 
 
@@ -221,10 +223,7 @@ def embed_images(
         model_id = self._get_default_or_validate(
             collection_id=collection_id, embedding_model_id=embedding_model_id
         )
-
-        model = self._models[model_id]
-        if not isinstance(model, ImageEmbeddingGenerator):
-            raise ValueError("Embedding model not compatible with images.")
+        model = self._get_image_model(model_id)
 
         # Query image filenames from the database.
         sample_id_to_filepath = {
@@ -270,9 +269,7 @@ def embed_annotations(
         model_id = self._get_default_or_validate(
             collection_id=annotation_collection_id, embedding_model_id=embedding_model_id
         )
-        model = self._models[model_id]
-        if not isinstance(model, ImageEmbeddingGenerator):
-            raise ValueError("Embedding model not compatible with images.")
+        model = self._get_image_model(model_id)
 
         annotation_sample_ids = annotation_resolver.get_unembedded_annotation_ids(
             session=session,
@@ -335,10 +332,7 @@ def compute_image_embedding(
         model_id = self._get_default_or_validate(
             collection_id=collection_id, embedding_model_id=embedding_model_id
         )
-
-        model = self._models[model_id]
-        if not isinstance(model, ImageEmbeddingGenerator):
-            raise ValueError("Embedding model not compatible with images.")
+        model = self._get_image_model(model_id)
 
         # Generate embedding for the image without progress bar.
         embeddings = model.embed_images(filepaths=[filepath], show_progress=False)
@@ -395,6 +389,43 @@ def embed_videos(
             embeddings=embeddings,
         )
 
+    def embed_and_store_pil_images(
+        self,
+        session: Session,
+        embedding_model_id: UUID,
+        sample_ids: list[UUID],
+        images: list[Image.Image],
+        show_progress: bool = True,
+    ) -> None:
+        """Generate and store embeddings for in-memory PIL images.
+
+        Args:
+            session: Database session for resolver operations.
+            embedding_model_id: ID of a registered image-compatible embedding model.
+            sample_ids: Sample IDs the embeddings are stored for.
+            images: PIL images to embed, in the same order as sample_ids.
+            show_progress: Whether to show a progress bar during embedding and storage.
+
+        Raises:
+            ValueError: If the model is missing, does not support image embedding, or
+                the number of images does not match the number of sample IDs.
+        """
+        if len(sample_ids) != len(images):
+            raise ValueError(
+                f"Expected the same number of sample IDs and images, got "
+                f"{len(sample_ids)} sample IDs and {len(images)} images."
+            )
+
+        model = self._get_image_model(embedding_model_id)
+        embeddings = model.embed_pil_images(images=images, show_progress=show_progress)
+        _store_embeddings(
+            session=session,
+            model_id=embedding_model_id,
+            sample_ids=sample_ids,
+            embeddings=embeddings,
+            show_progress=show_progress,
+        )
+
     def load_or_get_default_model(
         self,
         session: Session,
@@ -469,6 +500,19 @@ def _get_default_or_validate(
             raise ValueError(f"No embedding model found with ID {embedding_model_id}")
         return embedding_model_id
 
+    def _get_image_model(self, model_id: UUID) -> ImageEmbeddingGenerator:
+        """Return the registered image-compatible generator for model_id.
+
+        Raises:
+            ValueError: If no model is registered for the ID or it does not support images.
+        """
+        model = self._models.get(model_id)
+        if model is None:
+            raise ValueError(f"No embedding model found with ID {model_id}")
+        if not isinstance(model, ImageEmbeddingGenerator):
+            raise ValueError("Embedding model not compatible with images.")
+        return model
+
 
 def _store_embeddings(
     session: Session,

lightly_studio/src/lightly_studio/dataset/mobileclip_embedding_generator.py

@@ -8,6 +8,7 @@
 import numpy as np
 import torch
 from numpy.typing import NDArray
+from PIL import Image
 
 from lightly_studio.dataset.env import LIGHTLY_STUDIO_MODEL_CACHE_DIR
 from lightly_studio.models.embedding_model import EmbeddingModelCreate
@@ -119,6 +120,25 @@ def embed_image_crops(
             show_progress=show_progress,
         )
 
+    def embed_pil_images(
+        self, images: list[Image.Image], show_progress: bool = True
+    ) -> NDArray[np.float32]:
+        """Embed in-memory PIL images with MobileCLIP.
+
+        Args:
+            images: PIL images to embed.
+            show_progress: Whether to show a progress bar during embedding.
+
+        Returns:
+            A numpy array representing the generated embeddings in the same order
+            as the input images.
+        """
+        return image_embedding.embed_pil_images_batched(
+            images=images,
+            context=self._embedding_context(),
+            show_progress=show_progress,
+        )
+
     def _embedding_context(self) -> EmbeddingContext:
         """Build the model-specific configuration for batched image embedding."""
         return EmbeddingContext(