lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/dataset/edge_embedding_generator.py

@@ -0,0 +1,155 @@
+"""EdgeCLIP embedding generator."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Tuple
+from uuid import UUID
+
+import cv2
+import fsspec
+import numpy as np
+from lightly_edge_sdk import (
+    InferenceDeviceType,
+    LightlyEdge,
+    LightlyEdgeConfig,
+    LightlyEdgeDetectorConfig,
+)
+from numpy.typing import NDArray
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from lightly_studio.dataset.embedding_generator import ImageEmbeddingGenerator
+from lightly_studio.models.embedding_model import EmbeddingModelCreate
+
+MAX_BATCH_SIZE: int = 1
+
+
+class _ImageFileDatasetEdge(Dataset[Tuple[bytes, int, int]]):
+    """Dataset wrapping image file paths for processing."""
+
+    def __init__(
+        self,
+        filepaths: Sequence[str],
+    ) -> None:
+        self.filepaths = filepaths
+
+    def __len__(self) -> int:
+        return len(self.filepaths)
+
+    def __getitem__(self, idx: int) -> tuple[bytes, int, int]:
+        # Load the image.
+        with fsspec.open(self.filepaths[idx], "rb") as file:
+            image_bytes = file.read()
+            # Decode image from bytes using OpenCV
+            nparr = np.frombuffer(image_bytes, np.uint8)
+            bgr_image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+            rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
+            rgb_bytes = rgb_image.tobytes()
+            height, width, _ = rgb_image.shape
+            return rgb_bytes, width, height
+
+
+class EdgeSDKEmbeddingGenerator(ImageEmbeddingGenerator):
+    """Embedding generator using Edge SDK runtime."""
+
+    def __init__(self, model_path: str) -> None:
+        """Initialize the LightlyEdge object.
+
+        Args:
+            model_path: Path to the model tar file.
+        """
+        # Initialize the LightlyEdge SDK.
+        config = _create_edge_config()
+        self.lightly_edge = LightlyEdge(
+            path=model_path,
+            config=config,
+        )
+        model_config = self.lightly_edge.get_image_model_config()
+        self._model_hash = model_config.model_hash
+        self._embedding_size = model_config.embedding_size
+        self._model_name = model_config.model_name
+
+    def get_embedding_model_input(self, dataset_id: UUID) -> EmbeddingModelCreate:
+        """Generate an EmbeddingModelInput instance.
+
+        Args:
+            dataset_id: The ID of the dataset.
+
+        Returns:
+            An EmbeddingModelInput instance with the model details.
+        """
+        return EmbeddingModelCreate(
+            name=self._model_name,
+            embedding_model_hash=self._model_hash,
+            embedding_dimension=self._embedding_size,
+            dataset_id=dataset_id,
+        )
+
+    def embed_text(self, text: str) -> list[float]:
+        """Embed a text with EdgeCLIP.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            A list of floats representing the generated embedding.
+        """
+        embeddings = self.lightly_edge.embed_texts([text])
+        if len(embeddings):
+            return embeddings[0]
+        return []
+
+    def embed_images(self, filepaths: list[str]) -> NDArray[np.float32]:
+        """Embed images with EdgeSDK.
+
+        Args:
+            filepaths: A list of file paths to the images to embed.
+
+        Returns:
+            A numpy array representing the generated embeddings.
+        """
+        total_images = len(filepaths)
+        if not total_images:
+            return np.empty((0, self._embedding_size), dtype=np.float32)
+
+        dataset = _ImageFileDatasetEdge(filepaths)
+
+        # To avoid issues with db locking and multiprocessing we set the
+        # number of workers to 0 (no multiprocessing). The DataLoader is still
+        # very useful for batching and async prefetching of images.
+        loader = DataLoader(
+            dataset,
+            batch_size=MAX_BATCH_SIZE,
+            num_workers=0,  # must be 0 to avoid multiprocessing issues
+            pin_memory=True,
+        )
+
+        embeddings = np.empty((total_images, self._embedding_size), dtype=np.float32)
+        with tqdm(total=total_images, desc="Generating embeddings", unit=" images") as progress_bar:
+            for i, (rgb_bytes, width, height) in enumerate(loader):
+                embedding = self.lightly_edge.embed_frame_rgb_bytes(
+                    rgb_bytes=rgb_bytes[0],
+                    width=width[0].item(),
+                    height=height[0].item(),
+                )
+                embeddings[i] = embedding
+                progress_bar.update(1)
+
+        return embeddings
+
+
+def _create_edge_config() -> LightlyEdgeConfig:
+    """Create configuration for LightlyEdge.
+
+    Returns:
+        Configured LightlyEdgeConfig instance.
+    """
+    config = LightlyEdgeConfig.default()
+    config.inference_device_type = InferenceDeviceType.Auto
+    config.detector_config = LightlyEdgeDetectorConfig(
+        object_detector_enable=False,
+        classifiers_enable=False,
+        max_classifications=0,
+    )
+    return config

lightly_studio/dataset/embedding_generator.py

@@ -0,0 +1,129 @@
+"""EmbeddingGenerator implementations."""
+
+from __future__ import annotations
+
+import random
+from typing import Protocol, runtime_checkable
+from uuid import UUID
+
+import numpy as np
+from numpy.typing import NDArray
+
+from lightly_studio.models.embedding_model import EmbeddingModelCreate
+
+
+@runtime_checkable
+class EmbeddingGenerator(Protocol):
+    """Protocol defining the interface for embedding models.
+
+    This protocol defines the interface that all embedding models must
+    implement. Concrete implementations will use different techniques
+    for creating embeddings.
+    """
+
+    def get_embedding_model_input(self, dataset_id: UUID) -> EmbeddingModelCreate:
+        """Generate an EmbeddingModelCreate instance.
+
+        Args:
+            dataset_id: The ID of the dataset.
+
+        Returns:
+            An EmbeddingModelCreate instance with the model details.
+        """
+
+    def embed_text(self, text: str) -> list[float]:
+        """Generate an embedding for a text sample.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            A list of floats representing the generated embedding.
+        """
+        ...
+
+
+@runtime_checkable
+class ImageEmbeddingGenerator(EmbeddingGenerator, Protocol):
+    """Protocol defining the interface for image embedding models.
+
+    This protocol defines the interface that all image embedding models must
+    implement. Concrete implementations will use different techniques
+    for creating embeddings.
+    """
+
+    def embed_images(self, filepaths: list[str]) -> NDArray[np.float32]:
+        """Generate embeddings for multiple image samples.
+
+        TODO(Michal, 04/2025): Use DatasetLoader as input instead.
+
+        Args:
+            filepaths: A list of file paths to the images to embed.
+
+        Returns:
+            A numpy array representing the generated embeddings
+            in the same order as the input file paths.
+        """
+        ...
+
+
+@runtime_checkable
+class VideoEmbeddingGenerator(EmbeddingGenerator, Protocol):
+    """Protocol defining the interface for video embedding models.
+
+    This protocol defines the interface that all video embedding models must
+    implement. Concrete implementations will use different techniques
+    for creating embeddings.
+    """
+
+    def embed_videos(self, filepaths: list[str]) -> NDArray[np.float32]:
+        """Generate embeddings for multiple video samples.
+
+        Args:
+            filepaths: A list of file paths to the videos to embed.
+
+        Returns:
+            A numpy array representing the generated embeddings
+            in the same order as the input file paths.
+        """
+        ...
+
+
+class RandomEmbeddingGenerator(ImageEmbeddingGenerator, VideoEmbeddingGenerator):
+    """Model that produces random embeddings with a fixed dimension."""
+
+    def __init__(self, dimension: int = 3):
+        """Initialize the random embedding model.
+
+        Args:
+            dimension: The dimension of the embedding vectors to generate.
+        """
+        self._dimension = dimension
+
+    def get_embedding_model_input(self, dataset_id: UUID) -> EmbeddingModelCreate:
+        """Generate an EmbeddingModelCreate instance.
+
+        Args:
+            dataset_id: The ID of the dataset.
+
+        Returns:
+            An EmbeddingModelCreate instance with the model details.
+        """
+        return EmbeddingModelCreate(
+            name="Random",
+            embedding_model_hash="random_model",
+            embedding_dimension=self._dimension,
+            dataset_id=dataset_id,
+        )
+
+    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]) -> NDArray[np.float32]:
+        """Generate random embeddings for multiple image samples."""
+        return np.random.rand(len(filepaths), self._dimension).astype(np.float32)
+
+    def embed_videos(self, filepaths: list[str]) -> NDArray[np.float32]:
+        """Generate random embeddings for multiple image samples."""
+        return np.random.rand(len(filepaths), self._dimension).astype(np.float32)

lightly_studio/dataset/embedding_manager.py

@@ -0,0 +1,349 @@
+"""Embedding manager for dataset processing."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.dataset import env
+from lightly_studio.dataset.embedding_generator import (
+    EmbeddingGenerator,
+    ImageEmbeddingGenerator,
+    VideoEmbeddingGenerator,
+)
+from lightly_studio.models.dataset import SampleType
+from lightly_studio.models.embedding_model import EmbeddingModelTable
+from lightly_studio.models.sample_embedding import SampleEmbeddingCreate
+from lightly_studio.resolvers import (
+    dataset_resolver,
+    embedding_model_resolver,
+    image_resolver,
+    sample_embedding_resolver,
+    video_resolver,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingManagerProvider:
+    """Provider for the EmbeddingManager singleton instance."""
+
+    _instance: EmbeddingManager | None = None
+
+    @classmethod
+    def get_embedding_manager(cls) -> EmbeddingManager:
+        """Get the singleton instance of EmbeddingManager.
+
+        Returns:
+            The singleton instance of EmbeddingManager.
+
+        Raises:
+            ValueError: If no instance exists and no session is provided.
+        """
+        if cls._instance is None:
+            cls._instance = EmbeddingManager()
+        return cls._instance
+
+
+@dataclass
+class TextEmbedQuery:
+    """Parameters for text embedding generation."""
+
+    text: str
+    embedding_model_id: UUID | None = None
+
+
+class EmbeddingManager:
+    """Manages embedding models and handles embedding generation and storage."""
+
+    def __init__(self) -> None:
+        """Initialize the embedding manager."""
+        self._models: dict[UUID, EmbeddingGenerator] = {}
+        self._dataset_id_to_default_model_id: dict[UUID, UUID] = {}
+
+    def register_embedding_model(
+        self,
+        session: Session,
+        dataset_id: UUID,
+        embedding_generator: EmbeddingGenerator,
+        set_as_default: bool = False,
+    ) -> EmbeddingModelTable:
+        """Register an embedding model in the database.
+
+        The model is stored in an internal dictionary for later use.
+        The model is set as default if requested or if it's the first model.
+
+        Args:
+            session: Database session for resolver operations.
+            dataset_id: The ID of the dataset to associate with the model.
+                And to register as default, if requested.
+            embedding_generator: The model implementation used for embeddings.
+            set_as_default: Whether to set this model as the default.
+
+        Returns:
+            The created EmbeddingModel.
+        """
+        # Get or create embedding model record in the database.
+        db_model = embedding_model_resolver.get_or_create(
+            session=session,
+            embedding_model=embedding_generator.get_embedding_model_input(dataset_id=dataset_id),
+        )
+        model_id = db_model.embedding_model_id
+
+        # Store the model in our dictionary
+        self._models[model_id] = embedding_generator
+
+        # Set as default if requested or if it's the first model
+        if set_as_default or dataset_id not in self._dataset_id_to_default_model_id:
+            self._dataset_id_to_default_model_id[dataset_id] = model_id
+
+        return db_model
+
+    def embed_text(self, dataset_id: UUID, text_query: TextEmbedQuery) -> list[float]:
+        """Generate an embedding for a text sample.
+
+        Args:
+            dataset_id: The ID of the dataset to determine the registered default model.
+                It is used if embedding_model_id is not valid.
+            text_query: Text embedding query containing text and model ID.
+
+        Returns:
+            A list of floats representing the generated embedding.
+        """
+        model_id = self._get_default_or_validate(
+            dataset_id=dataset_id, embedding_model_id=text_query.embedding_model_id
+        )
+
+        model = self._models[model_id]
+
+        return model.embed_text(text_query.text)
+
+    def embed_images(
+        self,
+        session: Session,
+        dataset_id: UUID,
+        sample_ids: list[UUID],
+        embedding_model_id: UUID | None = None,
+    ) -> None:
+        """Generate and store embeddings for image samples.
+
+        Args:
+            session: Database session for resolver operations.
+            dataset_id: The ID of the dataset to determine the registered default model.
+                It is used if embedding_model_id is not valid.
+            sample_ids: List of sample IDs to generate embeddings for.
+            embedding_model_id: ID of the model to use. Uses default if None.
+
+        Raises:
+            ValueError: If no embedding model is registered, provided model
+            ID doesn't exist or if the embedding model does not support images.
+        """
+        model_id = self._get_default_or_validate(
+            dataset_id=dataset_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.")
+
+        # Query image filenames from the database.
+        sample_id_to_filepath = {
+            sample.sample_id: sample.file_path_abs
+            for sample in image_resolver.get_many_by_id(
+                session=session,
+                sample_ids=sample_ids,
+            )
+        }
+
+        # Extract filepaths in the same order as sample_ids.
+        filepaths = [sample_id_to_filepath[sample_id] for sample_id in sample_ids]
+
+        # Generate embeddings for the samples.
+        embeddings = model.embed_images(filepaths=filepaths)
+
+        # Convert to SampleEmbeddingCreate objects.
+        sample_embeddings = [
+            SampleEmbeddingCreate(
+                sample_id=sample_id,
+                embedding_model_id=model_id,
+                embedding=embedding,
+            )
+            for sample_id, embedding in zip(sample_ids, embeddings)
+        ]
+
+        # Store the embeddings in the database.
+        sample_embedding_resolver.create_many(session=session, sample_embeddings=sample_embeddings)
+
+    def embed_videos(
+        self,
+        session: Session,
+        dataset_id: UUID,
+        sample_ids: list[UUID],
+        embedding_model_id: UUID | None = None,
+    ) -> None:
+        """Generate and store embeddings for video samples.
+
+        Args:
+            session: Database session for resolver operations.
+            dataset_id: The ID of the dataset to determine the registered default model.
+                It is used if embedding_model_id is not valid.
+            sample_ids: List of sample IDs to generate embeddings for.
+            embedding_model_id: ID of the model to use. Uses default if None.
+
+        Raises:
+            ValueError: If no embedding model is registered, provided model
+            ID doesn't exist or if the embedding model does not support videos.
+        """
+        model_id = self._get_default_or_validate(
+            dataset_id=dataset_id, embedding_model_id=embedding_model_id
+        )
+
+        model = self._models[model_id]
+        if not isinstance(model, VideoEmbeddingGenerator):
+            raise ValueError("Embedding model not compatible with videos.")
+
+        # Get the samples
+        filepaths = []
+        for sample_id in sample_ids:
+            sample = video_resolver.get_by_id(session=session, sample_id=sample_id)
+            if sample is not None:
+                filepaths.append(sample.file_path_abs)
+
+        if len(filepaths) != len(sample_ids):
+            raise ValueError("Could not fetch all video paths for the provided IDs.")
+
+        # Generate embeddings for the samples.
+        embeddings = model.embed_videos(filepaths=filepaths)
+
+        # Convert to SampleEmbeddingCreate objects.
+        sample_embeddings = [
+            SampleEmbeddingCreate(
+                sample_id=sample_id,
+                embedding_model_id=model_id,
+                embedding=embedding,
+            )
+            for sample_id, embedding in zip(sample_ids, embeddings)
+        ]
+
+        # Store the embeddings in the database.
+        sample_embedding_resolver.create_many(session=session, sample_embeddings=sample_embeddings)
+
+    def load_or_get_default_model(
+        self,
+        session: Session,
+        dataset_id: UUID,
+    ) -> UUID | None:
+        """Ensure a default embedding model exists and return its ID.
+
+        Args:
+            session: Database session for resolver operations.
+            dataset_id: Dataset identifier the model should belong to.
+
+        Returns:
+            UUID of the default embedding model or None if the model cannot be loaded.
+        """
+        # Return the existing default model ID if available.
+
+        if dataset_id in self._dataset_id_to_default_model_id:
+            return self._dataset_id_to_default_model_id[dataset_id]
+
+        # Load the embedding generator based on sample_type from the env var.
+        dataset = dataset_resolver.get_by_id(session=session, dataset_id=dataset_id)
+        if dataset is None:
+            raise ValueError("Provided dataset_id could not be found.")
+
+        embedding_generator = _load_embedding_generator_from_env(sample_type=dataset.sample_type)
+        if embedding_generator is None:
+            return None
+
+        # Register the embedding model and set it as default.
+        embedding_model = self.register_embedding_model(
+            session=session,
+            dataset_id=dataset_id,
+            embedding_generator=embedding_generator,
+            set_as_default=True,
+        )
+
+        return embedding_model.embedding_model_id
+
+    def _get_default_or_validate(self, dataset_id: UUID, embedding_model_id: UUID | None) -> UUID:
+        """Get a valid model_id or raise error of non available.
+
+        If embedding_model_id is not provided, returns the default model for dataset_id.
+        If embedding_model_id is provided, validates that the model has been loaded and returns it.
+        """
+        default_model_id = self._dataset_id_to_default_model_id.get(dataset_id, None)
+        if embedding_model_id is None and default_model_id is None:
+            raise ValueError(
+                "No embedding_model_id provided and no default embedding model registered."
+            )
+
+        if embedding_model_id is None and default_model_id is not None:
+            return default_model_id
+
+        if embedding_model_id not in self._models:
+            raise ValueError(f"No embedding model found with ID {embedding_model_id}")
+        return embedding_model_id
+
+
+def _load_embedding_generator_from_env(sample_type: SampleType) -> EmbeddingGenerator | None:
+    """Load the embedding generator based on environment variable configuration."""
+    if sample_type == SampleType.IMAGE:
+        return _load_image_embedding_generator_from_env()
+    if sample_type == SampleType.VIDEO:
+        return _load_video_embedding_generator()
+    return None
+
+
+# TODO(Michal, 09/2025): Write tests for this function.
+def _load_image_embedding_generator_from_env() -> ImageEmbeddingGenerator | None:
+    if env.LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE == "EDGE":
+        try:
+            from lightly_studio.dataset.edge_embedding_generator import (
+                EdgeSDKEmbeddingGenerator,
+            )
+
+            logger.info("Using LightlyEdge embedding generator for images.")
+            return EdgeSDKEmbeddingGenerator(model_path=env.LIGHTLY_STUDIO_EDGE_MODEL_FILE_PATH)
+        except ImportError:
+            logger.warning("Embedding functionality is disabled.")
+    elif env.LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE == "MOBILE_CLIP":
+        try:
+            from lightly_studio.dataset.mobileclip_embedding_generator import (
+                MobileCLIPEmbeddingGenerator,
+            )
+
+            logger.info("Using MobileCLIP embedding generator for images.")
+            return MobileCLIPEmbeddingGenerator()
+        except ImportError:
+            logger.warning("Embedding functionality is disabled.")
+    elif env.LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE == "PE":
+        try:
+            from lightly_studio.dataset.perception_encoder_embedding_generator import (
+                PerceptionEncoderEmbeddingGenerator,
+            )
+
+            logger.info("Using PerceptionEncoder embedding generator for images.")
+            return PerceptionEncoderEmbeddingGenerator()
+        except ImportError:
+            logger.warning("Embedding functionality is disabled.")
+    else:
+        logger.warning(f"Unsupported model type: '{env.LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE}'")
+        logger.warning("Embedding functionality is disabled.")
+    return None
+
+
+def _load_video_embedding_generator() -> VideoEmbeddingGenerator | None:
+    try:
+        from lightly_studio.dataset.perception_encoder_embedding_generator import (
+            PerceptionEncoderEmbeddingGenerator,
+        )
+
+        logger.info("Using PerceptionEncoder embedding generator for videos.")
+        return PerceptionEncoderEmbeddingGenerator()
+    except ImportError:
+        logger.warning("Embedding functionality is disabled.")
+        return None

lightly_studio/dataset/env.py

@@ -0,0 +1,20 @@
+"""Initialize environment variables for the dataset module."""
+
+from typing import Optional
+
+from environs import Env
+
+env = Env()
+env.read_env()
+LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE: str = env.str(
+    "LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE", "MOBILE_CLIP"
+)
+LIGHTLY_STUDIO_EDGE_MODEL_FILE_PATH: str = env.str("EDGE_MODEL_PATH", "./lightly_model.tar")
+LIGHTLY_STUDIO_PROTOCOL: str = env.str("LIGHTLY_STUDIO_PROTOCOL", "http")
+LIGHTLY_STUDIO_PORT: int = env.int("LIGHTLY_STUDIO_PORT", 8001)
+LIGHTLY_STUDIO_HOST: str = env.str("LIGHTLY_STUDIO_HOST", "localhost")
+LIGHTLY_STUDIO_DEBUG: bool = env.bool("LIGHTLY_STUDIO_DEBUG", False)
+
+APP_URL = f"{LIGHTLY_STUDIO_PROTOCOL}://{LIGHTLY_STUDIO_HOST}:{LIGHTLY_STUDIO_PORT}"
+
+LIGHTLY_STUDIO_LICENSE_KEY: Optional[str] = env.str("LIGHTLY_STUDIO_LICENSE_KEY", default=None)

lightly_studio/dataset/file_utils.py

@@ -0,0 +1,49 @@
+"""File manipulation utilities."""
+
+from __future__ import annotations
+
+import logging
+import shutil
+from pathlib import Path
+
+import requests
+import xxhash
+
+logger = logging.getLogger(__name__)
+
+
+def download_file_if_does_not_exist(url: str, local_filename: Path) -> None:
+    """Download a file from a URL if it does not already exist locally."""
+    if local_filename.exists():
+        return
+
+    try:
+        logger.info(f"Downloading {url} to {local_filename}")
+        with requests.get(url, stream=True, timeout=30) as r:
+            # Raise an error for bad status codes
+            r.raise_for_status()
+            with open(local_filename, "wb") as f:
+                shutil.copyfileobj(r.raw, f)
+    except Exception:
+        # If download fails, remove any partial file to allow retry.
+        if local_filename.exists():
+            local_filename.unlink()
+        raise
+
+
+def get_file_xxhash(file_path: Path) -> str:
+    """Calculate the xxhash of a file.
+
+    XXHash is a fast non-cryptographic hash function.
+
+    Args:
+        file_path: Path to the file.
+
+    Returns:
+        The xxhash of the file as a string.
+    """
+    hasher = xxhash.xxh64()
+    with file_path.open("rb") as f:
+        while chunk := f.read(8192):
+            hasher.update(chunk)
+    return hasher.hexdigest()