lightly-studio 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

lightly_studio/core/sample.py

@@ -2,17 +2,22 @@
 
 from __future__ import annotations
 
-from typing import Any, Generic, Protocol, TypeVar
+from collections.abc import Iterable
+from typing import Any, Generic, Protocol, TypeVar, cast
 
-from sqlalchemy.orm import object_session
+from sqlalchemy.orm import Mapped, object_session
+from sqlmodel import Session, col
 
 from lightly_studio.models.sample import SampleTable
+from lightly_studio.resolvers import metadata_resolver, tag_resolver
 
 T = TypeVar("T")
 
 
-class _HasInner(Protocol):
-    _inner: Any
+class _DBFieldOwner(Protocol):
+    inner: Any
+
+    def get_object_session(self) -> Session: ...
 
 
 class DBField(Generic[T]):
@@ -20,53 +25,42 @@ class DBField(Generic[T]):
 
     Provides interface to a SQLAlchemy model field. Setting the field
     immediately commits to the database. The owner class must implement
-    the _inner attribute.
+    the inner attribute and the get_object_session() method.
     """
 
     __slots__ = ("_sqla_descriptor",)
     """Store the SQLAlchemy descriptor for accessing the field."""
 
-    def __init__(self, sqla_descriptor: T) -> None:
-        """Initialize the DBField with a SQLAlchemy descriptor.
-
-        Note: Mypy thinks that the descriptor has type T. In reality, during
-        runtime, it will be InstrumentedAttribute[T].
-        """
+    def __init__(self, sqla_descriptor: Mapped[T]) -> None:
+        """Initialize the DBField with a SQLAlchemy descriptor."""
         self._sqla_descriptor = sqla_descriptor
 
-    def __get__(self, obj: _HasInner | None, owner: type | None = None) -> T:
+    def __get__(self, obj: _DBFieldOwner | None, owner: type | None = None) -> T:
         """Get the value of the field from the database."""
         assert obj is not None, "DBField must be accessed via an instance, not the class"
         # Delegate to SQLAlchemy's descriptor.
-        # Note: Mypy incorrectly thinks that the descriptor has type T. It complains
-        # about the lack of a __get__ method.
-        value: T = self._sqla_descriptor.__get__(obj._inner, type(obj._inner))  # type: ignore[attr-defined]  # noqa: SLF001
+        value: T = self._sqla_descriptor.__get__(obj.inner, type(obj.inner))
         return value
 
-    def __set__(self, obj: _HasInner, value: T) -> None:
+    def __set__(self, obj: _DBFieldOwner, value: T) -> None:
         """Set the value of the field in the database. Commits the session."""
         # Delegate to SQLAlchemy's descriptor.
-        # Note: Mypy incorrectly thinks that the descriptor has type T. It complains
-        # about the lack of a __set__ method.
-        self._sqla_descriptor.__set__(obj._inner, value)  # type: ignore[attr-defined]  # noqa: SLF001
-        sess = object_session(obj._inner)  # noqa: SLF001
-        if sess is None:
-            raise RuntimeError("No active session found for the DBField object")
-        sess.commit()
+        self._sqla_descriptor.__set__(obj.inner, value)
+        obj.get_object_session().commit()
 
 
 class Sample:
     """Interface to a dataset sample."""
 
-    file_name = DBField(SampleTable.file_name)
-    width = DBField(SampleTable.width)
-    height = DBField(SampleTable.height)
-    dataset_id = DBField(SampleTable.dataset_id)
-    file_path_abs = DBField(SampleTable.file_path_abs)
+    file_name = DBField(col(SampleTable.file_name))
+    width = DBField(col(SampleTable.width))
+    height = DBField(col(SampleTable.height))
+    dataset_id = DBField(col(SampleTable.dataset_id))
+    file_path_abs = DBField(col(SampleTable.file_path_abs))
 
-    sample_id = DBField(SampleTable.sample_id)
-    created_at = DBField(SampleTable.created_at)
-    updated_at = DBField(SampleTable.updated_at)
+    sample_id = DBField(col(SampleTable.sample_id))
+    created_at = DBField(col(SampleTable.created_at))
+    updated_at = DBField(col(SampleTable.updated_at))
 
     def __init__(self, inner: SampleTable) -> None:
         """Initialize the Sample.
@@ -74,4 +68,137 @@ class Sample:
         Args:
             inner: The SampleTable SQLAlchemy model instance.
         """
-        self._inner = inner
+        self.inner = inner
+        self._metadata = SampleMetadata(self)
+
+    def get_object_session(self) -> Session:
+        """Get the database session for this sample.
+
+        Returns:
+            The SQLModel session.
+
+        Raises:
+            RuntimeError: If no active session is found.
+        """
+        session = object_session(self.inner)
+        if session is None:
+            raise RuntimeError("No active session found for the sample")
+        # Cast from SQLAlchemy Session to SQLModel Session for mypy.
+        return cast(Session, session)
+
+    def add_tag(self, name: str) -> None:
+        """Add a tag to this sample.
+
+        If the tag doesn't exist, it will be created first.
+
+        Args:
+            name: The name of the tag to add.
+        """
+        session = self.get_object_session()
+
+        # Get or create the tag for this dataset.
+        tag = tag_resolver.get_or_create_sample_tag_by_name(
+            session=session, dataset_id=self.dataset_id, tag_name=name
+        )
+
+        # Add the tag to the sample if not already associated.
+        if tag not in self.inner.tags:
+            tag_resolver.add_tag_to_sample(session=session, tag_id=tag.tag_id, sample=self.inner)
+
+    def remove_tag(self, name: str) -> None:
+        """Remove a tag from this sample.
+
+        Args:
+            name: The name of the tag to remove.
+        """
+        session = self.get_object_session()
+
+        # Find the tag by name for this dataset.
+        existing_tag = tag_resolver.get_by_name(
+            session=session, tag_name=name, dataset_id=self.dataset_id
+        )
+
+        # Remove the tag from the sample if it exists and is associated
+        if existing_tag is not None and existing_tag in self.inner.tags:
+            tag_resolver.remove_tag_from_sample(
+                session=session, tag_id=existing_tag.tag_id, sample=self.inner
+            )
+
+    @property
+    def tags(self) -> set[str]:
+        """Get the tag names associated with this sample.
+
+        Returns:
+            A set of tag names as strings.
+        """
+        return {tag.name for tag in self.inner.tags}
+
+    @tags.setter
+    def tags(self, tags: Iterable[str]) -> None:
+        """Set the tags for this sample, replacing any existing tags.
+
+        Args:
+            tags: Iterable of tag names to associate with this sample.
+        """
+        # Get current tag names
+        current_tags = self.tags
+        new_tags = set(tags)
+
+        # Remove tags that are no longer needed
+        tags_to_remove = current_tags - new_tags
+        for tag_name in tags_to_remove:
+            self.remove_tag(tag_name)
+
+        # Add new tags
+        tags_to_add = new_tags - current_tags
+        for tag_name in tags_to_add:
+            self.add_tag(tag_name)
+
+    @property
+    def metadata(self) -> SampleMetadata:
+        """Get dictionary-like access to sample metadata.
+
+        Returns:
+            A dictionary-like object for accessing metadata.
+        """
+        return self._metadata
+
+
+class SampleMetadata:
+    """Dictionary-like interface for sample metadata."""
+
+    def __init__(self, sample: Sample) -> None:
+        """Initialize SampleMetadata.
+
+        Args:
+            sample: The Sample instance this metadata belongs to.
+        """
+        self._sample = sample
+
+    def __getitem__(self, key: str) -> Any:
+        """Get a metadata value by key.
+
+        Args:
+            key: The metadata key to access.
+
+        Returns:
+            The metadata value for the given key, or None if the key doesn't exist.
+        """
+        if self._sample.inner.metadata_dict is None:
+            return None
+        return self._sample.inner.metadata_dict.get_value(key)
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        """Set a metadata key-value pair.
+
+        Args:
+            key: The metadata key.
+            value: The metadata value.
+        """
+        session = self._sample.get_object_session()
+        metadata_resolver.set_value_for_sample(
+            session=session,
+            sample_id=self._sample.sample_id,
+            key=key,
+            value=value,
+        )

lightly_studio/core/start_gui.py

@@ -2,12 +2,47 @@
 
 from __future__ import annotations
 
+from lightly_studio import db_manager
 from lightly_studio.api.server import Server
 from lightly_studio.dataset import env
+from lightly_studio.resolvers import dataset_resolver, sample_resolver
+
+
+def _validate_has_samples() -> None:
+    """Validate that there are samples in the database before starting GUI.
+
+    Raises:
+        ValueError: If no datasets are found or if no samples exist in any dataset.
+    """
+    session = db_manager.persistent_session()
+
+    # Check if any datasets exist
+    datasets = dataset_resolver.get_all(session=session, offset=0, limit=1)
+
+    if not datasets:
+        raise ValueError(
+            "No datasets found. Please load a dataset using Dataset class methods "
+            "(e.g., add_samples_from_path(), add_samples_from_yolo(), etc.) "
+            "before starting the GUI."
+        )
+
+    # Check if there are any samples in the first dataset
+    first_dataset = datasets[0]
+    sample_count = sample_resolver.count_by_dataset_id(
+        session=session, dataset_id=first_dataset.dataset_id
+    )
+
+    if sample_count == 0:
+        raise ValueError(
+            "No images have been indexed for the first dataset. "
+            "Please ensure your dataset contains valid images and try loading again."
+        )
 
 
 def start_gui() -> None:
     """Launch the web interface for the loaded dataset."""
+    _validate_has_samples()
+
     server = Server(host=env.LIGHTLY_STUDIO_HOST, port=env.LIGHTLY_STUDIO_PORT)
 
     print(f"Open the LightlyStudio GUI under: {env.APP_URL}")

lightly_studio/dataset/edge_embedding_generator.py

@@ -3,11 +3,12 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from pathlib import Path
 from typing import Tuple
 from uuid import UUID
 
 import cv2
+import fsspec
+import numpy as np
 from lightly_edge_sdk import (
     InferenceDeviceType,
     LightlyEdge,
@@ -29,7 +30,7 @@ class _ImageFileDatasetEdge(Dataset[Tuple[bytes, int, int]]):
 
     def __init__(
         self,
-        filepaths: Sequence[Path],
+        filepaths: Sequence[str],
     ) -> None:
         self.filepaths = filepaths
 
@@ -38,11 +39,15 @@ class _ImageFileDatasetEdge(Dataset[Tuple[bytes, int, int]]):
 
     def __getitem__(self, idx: int) -> tuple[bytes, int, int]:
         # Load the image.
-        bgr_image = cv2.imread(str(self.filepaths[idx]))
-        rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
-        rgb_bytes = rgb_image.tobytes()
-        height, width, _ = rgb_image.shape
-        return rgb_bytes, width, height
+        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(EmbeddingGenerator):
@@ -95,7 +100,7 @@ class EdgeSDKEmbeddingGenerator(EmbeddingGenerator):
             return embeddings[0]
         return []
 
-    def embed_images(self, filepaths: list[Path]) -> list[list[float]]:
+    def embed_images(self, filepaths: list[str]) -> list[list[float]]:
         """Embed images with EdgeSDK.
 
         Args:

lightly_studio/dataset/embedding_generator.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import random
-from pathlib import Path
 from typing import Protocol, runtime_checkable
 from uuid import UUID
 
@@ -40,7 +39,7 @@ class EmbeddingGenerator(Protocol):
         """
         ...
 
-    def embed_images(self, filepaths: list[Path]) -> list[list[float]]:
+    def embed_images(self, filepaths: list[str]) -> list[list[float]]:
         """Generate embeddings for multiple image samples.
 
         TODO(Michal, 04/2025): Use DatasetLoader as input instead.
@@ -86,6 +85,6 @@ class RandomEmbeddingGenerator(EmbeddingGenerator):
         """Generate a random embedding for a text sample."""
         return [random.random() for _ in range(self._dimension)]
 
-    def embed_images(self, filepaths: list[Path]) -> list[list[float]]:
+    def embed_images(self, filepaths: list[str]) -> list[list[float]]:
         """Generate random embeddings for multiple image samples."""
         return [[random.random() for _ in range(self._dimension)] for _ in range(len(filepaths))]

lightly_studio/dataset/embedding_manager.py

@@ -3,11 +3,11 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from pathlib import Path
 from uuid import UUID
 
 from sqlmodel import Session
 
+from lightly_studio.dataset import env
 from lightly_studio.dataset.embedding_generator import EmbeddingGenerator
 from lightly_studio.models.embedding_model import EmbeddingModelTable
 from lightly_studio.models.sample_embedding import SampleEmbeddingCreate
@@ -75,12 +75,12 @@ class EmbeddingManager:
         Returns:
             The created EmbeddingModel.
         """
-        # Create embedding model record in the database.
-        created_model = embedding_model_resolver.create(
+        # 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 = created_model.embedding_model_id
+        model_id = db_model.embedding_model_id
 
         # Store the model in our dictionary
         self._models[model_id] = embedding_generator
@@ -89,7 +89,7 @@ class EmbeddingManager:
         if set_as_default or self._default_model_id is None:
             self._default_model_id = model_id
 
-        return created_model
+        return db_model
 
     def embed_text(self, text_query: TextEmbedQuery) -> list[float]:
         """Generate an embedding for a text sample.
@@ -136,7 +136,7 @@ class EmbeddingManager:
 
         # Query image filenames from the database.
         sample_id_to_filepath = {
-            sample.sample_id: Path(sample.file_path_abs)
+            sample.sample_id: sample.file_path_abs
             for sample in sample_resolver.get_many_by_id(
                 session=session,
                 sample_ids=sample_ids,
@@ -161,3 +161,71 @@ class EmbeddingManager:
 
         # 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.
+        # TODO(Michal, 09/2025): We do not check if the model belongs to the dataset.
+        # The design of EmbeddingManager needs to change to support multiple datasets.
+        if self._default_model_id is not None:
+            return self._default_model_id
+
+        # Load the embedding generator based on configuration.
+        embedding_generator = _load_embedding_generator_from_env()
+        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
+
+
+# TODO(Michal, 09/2025): Write tests for this function.
+def _load_embedding_generator_from_env() -> EmbeddingGenerator | None:
+    """Load the embedding generator based on environment variable configuration."""
+    if env.LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE == "EDGE":
+        try:
+            from lightly_studio.dataset.edge_embedding_generator import (
+                EdgeSDKEmbeddingGenerator,
+            )
+
+            print("Using LightlyEdge embedding generator.")
+            return EdgeSDKEmbeddingGenerator(model_path=env.LIGHTLY_STUDIO_EDGE_MODEL_FILE_PATH)
+        except ImportError:
+            print("Embedding functionality is disabled.")
+            return None
+    elif env.LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE == "MOBILE_CLIP":
+        try:
+            from lightly_studio.dataset.mobileclip_embedding_generator import (
+                MobileCLIPEmbeddingGenerator,
+            )
+
+            print("Using MobileCLIP embedding generator.")
+            return MobileCLIPEmbeddingGenerator()
+        except ImportError:
+            print("Embedding functionality is disabled.")
+            return None
+
+    print(
+        f"Unsupported model type: '{env.LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE}'",
+    )
+    print("Embedding functionality is disabled.")
+    return None

lightly_studio/dataset/env.py

@@ -1,5 +1,7 @@
 """Initialize environment variables for the dataset module."""
 
+from typing import Optional
+
 from environs import Env
 
 env = Env()
@@ -14,3 +16,5 @@ LIGHTLY_STUDIO_HOST: str = env.str("LIGHTLY_STUDIO_HOST", "localhost")
 LIGHTLY_STUDIO_DEBUG: str = 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

@@ -13,8 +13,19 @@ 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
-    with requests.get(url, stream=True) as r, open(local_filename, "wb") as f:
-        shutil.copyfileobj(r.raw, f)
+
+    try:
+        print(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: