Extract EmbeddingResult to a shared, path-agnostic module

lightly_studio/src/lightly_studio/dataset/embedding_generator.py

@@ -11,7 +11,7 @@
 from numpy.typing import NDArray
 from PIL import Image
 
-from lightly_studio.dataset.image_embedding import ImageEmbeddingResult
+from lightly_studio.dataset.embedding_result import EmbeddingResult
 from lightly_studio.models.embedding_model import EmbeddingModelCreate
 
 
@@ -66,9 +66,7 @@ class ImageEmbeddingGenerator(EmbeddingGenerator, Protocol):
     for creating embeddings.
     """
 
-    def embed_images(
-        self, filepaths: list[str], show_progress: bool = True
-    ) -> ImageEmbeddingResult:
+    def embed_images(self, filepaths: list[str], show_progress: bool = True) -> EmbeddingResult:
         """Generate embeddings for multiple image samples.
 
         TODO(Michal, 04/2025): Use DatasetLoader as input instead.
@@ -78,7 +76,7 @@ def embed_images(
             show_progress: Whether to show a progress bar during embedding.
 
         Returns:
-            An ``ImageEmbeddingResult`` with embeddings for the readable files, in the same
+            An ``EmbeddingResult`` with embeddings for the readable files, in the same
             order as the corresponding input file paths.
         """
         ...
@@ -167,13 +165,11 @@ def embed_text(self, _text: str) -> list[float]:
         """Generate a random embedding for a text sample."""
         return [random.random() for _ in range(self._dimension)]
 
-    def embed_images(
-        self, filepaths: list[str], show_progress: bool = True
-    ) -> ImageEmbeddingResult:
+    def embed_images(self, filepaths: list[str], show_progress: bool = True) -> EmbeddingResult:
         """Generate random embeddings for multiple image samples."""
         _ = show_progress  # Not used for random embeddings.
         embeddings = np.random.rand(len(filepaths), self._dimension).astype(np.float32)
-        return ImageEmbeddingResult(embeddings=embeddings, kept_indices=list(range(len(filepaths))))
+        return EmbeddingResult(embeddings=embeddings, kept_indices=list(range(len(filepaths))))
 
     def embed_image_crops(
         self, image_crops: list[ImageCrop], show_progress: bool = True

lightly_studio/src/lightly_studio/dataset/embedding_result.py

@@ -0,0 +1,32 @@
+"""Result type shared by the batched embedding paths.
+
+Holds :class:`EmbeddingResult`, the model-agnostic return type used across the
+image, image-crop, and video embedding paths. It lives in its own module so any
+embedding path can depend on it without pulling in a path-specific module.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+@dataclass(frozen=True)
+class EmbeddingResult:
+    """Embeddings for the inputs that could be read, plus which inputs they cover.
+
+    Broken inputs (unreadable/undecodable files) are skipped rather than aborting the
+    whole run, so ``embeddings`` may hold fewer rows than the input list. ``kept_indices``
+    maps each row back to its position in the input list, in input order, letting callers
+    realign any parallel per-input data (e.g. sample IDs) with the embeddings.
+
+    Attributes:
+        embeddings: Float32 array of shape ``(len(kept_indices), embedding_dimension)``.
+        kept_indices: Indices into the input list of the inputs that were embedded,
+            in input order.
+    """
+
+    embeddings: NDArray[np.float32]
+    kept_indices: list[int]

lightly_studio/src/lightly_studio/dataset/image_embedding.py

@@ -21,6 +21,7 @@
 from tqdm import tqdm
 
 from lightly_studio.core.file_outcome_report import FileOutcome, FileOutcomeReport
+from lightly_studio.dataset.embedding_result import EmbeddingResult
 from lightly_studio.utils import batching, parallelize
 
 _ItemT = TypeVar("_ItemT")
@@ -53,30 +54,11 @@ class _EmbeddingProgress:
     unit: str
 
 
-@dataclass(frozen=True)
-class ImageEmbeddingResult:
-    """Embeddings for the image files that could be read, plus which inputs they cover.
-
-    Broken files (unreadable/undecodable) are skipped rather than aborting the whole
-    run, so ``embeddings`` may hold fewer rows than the input file list. ``kept_indices``
-    maps each row back to its position in the input list, in input order, letting callers
-    realign any parallel per-file data (e.g. sample IDs) with the embeddings.
-
-    Attributes:
-        embeddings: Float32 array of shape ``(len(kept_indices), embedding_dimension)``.
-        kept_indices: Indices into the input file list of the files that were embedded,
-            in input order.
-    """
-
-    embeddings: NDArray[np.float32]
-    kept_indices: list[int]
-
-
 def embed_image_files_batched(
     filepaths: list[str],
     context: EmbeddingContext,
     show_progress: bool,
-) -> ImageEmbeddingResult:
+) -> EmbeddingResult:
     """Embed image files in batches, preserving input order and skipping broken files.
 
     Args:
@@ -85,7 +67,7 @@ def embed_image_files_batched(
         show_progress: Whether to show a tqdm progress bar.
 
     Returns:
-        An ``ImageEmbeddingResult`` whose embeddings cover only the readable files, with
+        An ``EmbeddingResult`` whose embeddings cover only the readable files, with
         ``kept_indices`` mapping each row back to its input position.
 
     Raises:
@@ -152,7 +134,7 @@ def _embed_items_batched(
     context: EmbeddingContext,
     show_progress: bool,
     progress: _EmbeddingProgress,
-) -> ImageEmbeddingResult:
+) -> EmbeddingResult:
     """Preprocess items on a thread pool and embed them in batches, preserving order.
 
     ``preprocess_item`` (per-item PIL decode/resize/normalize, plus remote reads for file
@@ -164,13 +146,13 @@ def _embed_items_batched(
     bounded.
 
     Returns:
-        An ``ImageEmbeddingResult`` whose embeddings cover the kept items and whose
+        An ``EmbeddingResult`` whose embeddings cover the kept items and whose
         ``kept_indices`` map each row back to its index into ``items``, both in input order.
     """
     total_items = len(items)
     if not total_items:
         empty = np.empty((0, context.embedding_dimension), dtype=np.float32)
-        return ImageEmbeddingResult(embeddings=empty, kept_indices=[])
+        return EmbeddingResult(embeddings=empty, kept_indices=[])
 
     if context.max_batch_size <= 0:
         raise ValueError("max_batch_size must be positive.")
@@ -193,7 +175,7 @@ def _embed_items_batched(
         show_progress=show_progress,
         progress=progress,
     )
-    return ImageEmbeddingResult(embeddings=embeddings, kept_indices=kept_indices)
+    return EmbeddingResult(embeddings=embeddings, kept_indices=kept_indices)
 
 
 def _keep_non_none(

lightly_studio/src/lightly_studio/dataset/mobileclip_embedding_generator.py

@@ -16,7 +16,8 @@
 
 from . import file_utils, image_crop_embedding, image_embedding
 from .embedding_generator import ImageCrop, ImageEmbeddingGenerator
-from .image_embedding import EmbeddingContext, ImageEmbeddingResult
+from .embedding_result import EmbeddingResult
+from .image_embedding import EmbeddingContext
 
 MODEL_NAME = "mobileclip_s0"
 MOBILECLIP_DOWNLOAD_URL = (
@@ -84,17 +85,15 @@ def embed_text(self, text: str) -> list[float]:
             embedding_list: list[float] = embedding.cpu().numpy().flatten().tolist()
         return embedding_list
 
-    def embed_images(
-        self, filepaths: list[str], show_progress: bool = True
-    ) -> ImageEmbeddingResult:
+    def embed_images(self, filepaths: list[str], show_progress: bool = True) -> EmbeddingResult:
         """Embed images with MobileCLIP.
 
         Args:
             filepaths: A list of file paths to the images to embed.
             show_progress: Whether to show a progress bar during embedding.
 
         Returns:
-            An ``ImageEmbeddingResult`` with embeddings for the readable files, in the same
+            An ``EmbeddingResult`` with embeddings for the readable files, in the same
             order as the corresponding input file paths.
         """
         return image_embedding.embed_image_files_batched(

lightly_studio/src/lightly_studio/dataset/perception_encoder_embedding_generator.py

@@ -21,7 +21,8 @@
 
 from . import file_utils, image_crop_embedding, image_embedding
 from .embedding_generator import ImageCrop, ImageEmbeddingGenerator, VideoEmbeddingGenerator
-from .image_embedding import EmbeddingContext, ImageEmbeddingResult
+from .embedding_result import EmbeddingResult
+from .image_embedding import EmbeddingContext
 
 MODEL_NAME = "PE-Core-T16-384"
 DEFAULT_VIDEO_CHANNEL = 0
@@ -162,17 +163,15 @@ def embed_text(self, text: str) -> list[float]:
             embedding_list: list[float] = embedding.cpu().numpy().flatten().tolist()
         return embedding_list
 
-    def embed_images(
-        self, filepaths: list[str], show_progress: bool = True
-    ) -> ImageEmbeddingResult:
+    def embed_images(self, filepaths: list[str], show_progress: bool = True) -> EmbeddingResult:
         """Embed images with Perception Encoder.
 
         Args:
             filepaths: A list of file paths to the images to embed.
             show_progress: Whether to show a progress bar during embedding.
 
         Returns:
-            An ``ImageEmbeddingResult`` with embeddings for the readable files, in the same
+            An ``EmbeddingResult`` with embeddings for the readable files, in the same
             order as the corresponding input file paths.
         """
         return image_embedding.embed_image_files_batched(

lightly_studio/src/lightly_studio/examples/example_custom_embedding_model.py

@@ -23,8 +23,9 @@
 import lightly_studio as ls
 from lightly_studio.database import db_manager
 from lightly_studio.dataset import file_utils, image_crop_embedding, image_embedding
+from lightly_studio.dataset.embedding_result import EmbeddingResult
 from lightly_studio.dataset.env import LIGHTLY_STUDIO_MODEL_CACHE_DIR
-from lightly_studio.dataset.image_embedding import EmbeddingContext, ImageEmbeddingResult
+from lightly_studio.dataset.image_embedding import EmbeddingContext
 from lightly_studio.models.embedding_model import EmbeddingModelCreate
 from lightly_studio.vendor import mobileclip
 
@@ -86,9 +87,7 @@ def embed_text(self, text: str) -> list[float]:
             embedding_list: list[float] = embedding.cpu().numpy().flatten().tolist()
         return embedding_list
 
-    def embed_images(
-        self, filepaths: list[str], show_progress: bool = True
-    ) -> ImageEmbeddingResult:
+    def embed_images(self, filepaths: list[str], show_progress: bool = True) -> EmbeddingResult:
         """Embed a batch of images, returning one row per readable input path."""
         return image_embedding.embed_image_files_batched(
             filepaths=filepaths,

lightly_studio/tests/dataset/test_embedding_manager.py

@@ -21,7 +21,7 @@
     EmbeddingManager,
     TextEmbedQuery,
 )
-from lightly_studio.dataset.image_embedding import ImageEmbeddingResult
+from lightly_studio.dataset.embedding_result import EmbeddingResult
 from lightly_studio.models.annotation.annotation_base import AnnotationType
 from lightly_studio.models.collection import CollectionTable, SampleType
 from lightly_studio.models.embedding_model import EmbeddingModelCreate, EmbeddingModelTable
@@ -100,9 +100,7 @@ def get_embedding_model_input(self, collection_id: UUID) -> EmbeddingModelCreate
         def embed_text(self, text: str) -> list[float]:
             raise NotImplementedError()
 
-        def embed_images(
-            self, filepaths: list[str], show_progress: bool = True
-        ) -> ImageEmbeddingResult:
+        def embed_images(self, filepaths: list[str], show_progress: bool = True) -> EmbeddingResult:
             raise NotImplementedError()
 
         def embed_image_crops(
@@ -627,11 +625,9 @@ def embed_text(self, text: str) -> list[float]:
             _ = text
             return [0.1, 0.2, 0.3]
 
-        def embed_images(
-            self, filepaths: list[str], show_progress: bool = True
-        ) -> ImageEmbeddingResult:
+        def embed_images(self, filepaths: list[str], show_progress: bool = True) -> EmbeddingResult:
             _ = show_progress
-            return ImageEmbeddingResult(
+            return EmbeddingResult(
                 embeddings=np.zeros((len(filepaths), 3), dtype=np.float32),
                 kept_indices=list(range(len(filepaths))),
             )