From d2f988d9c406ab150a0730b46565b9b17670439a Mon Sep 17 00:00:00 2001 From: JonasWurst Date: Wed, 22 Jul 2026 07:48:02 +0200 Subject: [PATCH] Extract EmbeddingResult to a shared, path-agnostic module Rename ImageEmbeddingResult -> EmbeddingResult and move it out of image_embedding.py into its own dependency-free module (dataset/embedding_result.py). The type carries embeddings plus kept_indices and is not image-specific; giving it a neutral home lets the image, image-crop, and (future) video embedding paths reuse it without importing a path-specific module. Pure rename/move: no behavior change. All references and docstrings updated. Only the image embedding path is touched here; the crop path still returns NDArray on main and adopts EmbeddingResult in the stacked crop PR. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../dataset/embedding_generator.py | 14 +++----- .../dataset/embedding_result.py | 32 +++++++++++++++++++ .../lightly_studio/dataset/image_embedding.py | 32 ++++--------------- .../dataset/mobileclip_embedding_generator.py | 9 +++--- .../perception_encoder_embedding_generator.py | 9 +++--- .../example_custom_embedding_model.py | 7 ++-- .../tests/dataset/test_embedding_manager.py | 12 +++---- 7 files changed, 59 insertions(+), 56 deletions(-) create mode 100644 lightly_studio/src/lightly_studio/dataset/embedding_result.py diff --git a/lightly_studio/src/lightly_studio/dataset/embedding_generator.py b/lightly_studio/src/lightly_studio/dataset/embedding_generator.py index d88f8977a..baf991a04 100644 --- a/lightly_studio/src/lightly_studio/dataset/embedding_generator.py +++ b/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 diff --git a/lightly_studio/src/lightly_studio/dataset/embedding_result.py b/lightly_studio/src/lightly_studio/dataset/embedding_result.py new file mode 100644 index 000000000..3dcefe63a --- /dev/null +++ b/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] diff --git a/lightly_studio/src/lightly_studio/dataset/image_embedding.py b/lightly_studio/src/lightly_studio/dataset/image_embedding.py index 493005722..71a51e767 100644 --- a/lightly_studio/src/lightly_studio/dataset/image_embedding.py +++ b/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( 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 d9c4a859e..2bd5e3efe 100644 --- a/lightly_studio/src/lightly_studio/dataset/mobileclip_embedding_generator.py +++ b/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,9 +85,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 images with MobileCLIP. Args: @@ -94,7 +93,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. """ return image_embedding.embed_image_files_batched( 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 31e7e412c..2078ecf70 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 @@ -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,9 +163,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 images with Perception Encoder. Args: @@ -172,7 +171,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. """ return image_embedding.embed_image_files_batched( diff --git a/lightly_studio/src/lightly_studio/examples/example_custom_embedding_model.py b/lightly_studio/src/lightly_studio/examples/example_custom_embedding_model.py index 3f6cc7463..ba5c0dc5a 100644 --- a/lightly_studio/src/lightly_studio/examples/example_custom_embedding_model.py +++ b/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, diff --git a/lightly_studio/tests/dataset/test_embedding_manager.py b/lightly_studio/tests/dataset/test_embedding_manager.py index 3e0f079e9..75d47acab 100644 --- a/lightly_studio/tests/dataset/test_embedding_manager.py +++ b/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))), )