lightly-studio 0.3.1__py3-none-any.whl

lightly_studio/models/classifier.py

@@ -0,0 +1,20 @@
+"""This module defines the data model for the classifier."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from pydantic import BaseModel
+
+
+class EmbeddingClassifier(BaseModel):
+    """Base class for the Classifier model."""
+
+    """The name of the classifier."""
+    classifier_name: str
+
+    """The ID of the classifier."""
+    classifier_id: UUID
+
+    """List of classes supported by the classifier."""
+    class_list: list[str]

lightly_studio/models/dataset.py

@@ -0,0 +1,84 @@
+"""This module contains the Dataset model and related enumerations."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from datetime import datetime, timezone
+from typing import cast
+from uuid import UUID, uuid4
+
+from sqlalchemy.orm import Session as SQLAlchemySession
+from sqlmodel import Field, Session, SQLModel
+
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.resolvers import sample_resolver
+from lightly_studio.resolvers.samples_filter import SampleFilter
+
+
+class DatasetBase(SQLModel):
+    """Base class for the Dataset model."""
+
+    name: str
+    directory: str
+
+
+class DatasetCreate(DatasetBase):
+    """Dataset class when inserting."""
+
+
+class DatasetView(DatasetBase):
+    """Dataset class when retrieving."""
+
+    dataset_id: UUID
+    created_at: datetime
+    updated_at: datetime
+
+
+class DatasetTable(DatasetBase, table=True):
+    """This class defines the Dataset model."""
+
+    __tablename__ = "datasets"
+    dataset_id: UUID = Field(default_factory=uuid4, primary_key=True)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), index=True)
+    updated_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+    )
+
+    def get_samples(
+        self,
+        offset: int = 0,
+        limit: int | None = None,
+        filters: SampleFilter | None = None,
+        text_embedding: list[float] | None = None,
+        sample_ids: list[UUID] | None = None,
+    ) -> Sequence[SampleTable]:
+        """Retrieve samples for this dataset with optional filtering.
+
+        Just passes the parameters to the sample resolver.
+
+        Args:
+            offset: Offset for pagination.
+            limit: Limit for pagination.
+            filters: Optional filters to apply.
+            text_embedding: Optional text embedding for filtering.
+            sample_ids: Optional list of sample IDs to filter by.
+
+        Returns:
+            A sequence of SampleTable objects.
+        """
+        # Get the session from the instance.
+        # SQLAlchemy Session is compatible with SQLModel's Session at runtime,
+        # but we have to help mypy.
+        session = cast(Session, SQLAlchemySession.object_session(self))
+        if session is None:
+            raise RuntimeError("No database session found for this instance")
+
+        return sample_resolver.get_all_by_dataset_id(
+            session=session,
+            dataset_id=self.dataset_id,
+            offset=offset,
+            limit=limit,
+            filters=filters,
+            text_embedding=text_embedding,
+            sample_ids=sample_ids,
+        ).samples

lightly_studio/models/embedding_model.py

@@ -0,0 +1,30 @@
+"""This module defines the Embedding_Model model for the application."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from uuid import UUID, uuid4
+
+from sqlmodel import CHAR, Column, Field, SQLModel
+
+
+class EmbeddingModelBase(SQLModel):
+    """Base class for the EmbeddingModel."""
+
+    name: str
+    parameter_count_in_mb: int | None = None
+    embedding_model_hash: str = Field(default="", sa_column=Column(CHAR(128)))
+    embedding_dimension: int
+    dataset_id: UUID = Field(default=None, foreign_key="datasets.dataset_id")
+
+
+class EmbeddingModelCreate(EmbeddingModelBase):
+    """Model used for creating an embedding model."""
+
+
+class EmbeddingModelTable(EmbeddingModelBase, table=True):
+    """This class defines the EmbeddingModel model."""
+
+    __tablename__ = "embedding_models"
+    embedding_model_id: UUID = Field(default_factory=uuid4, primary_key=True)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), index=True)

lightly_studio/models/metadata.py

@@ -0,0 +1,208 @@
+"""Metadata models for storing custom metadata."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+from uuid import UUID, uuid4
+
+from pydantic import BaseModel
+from sqlalchemy.orm.attributes import flag_modified
+from sqlalchemy.types import JSON
+from sqlmodel import Field, Relationship, SQLModel
+
+from lightly_studio.metadata.complex_metadata import (
+    COMPLEX_METADATA_TYPES,
+    deserialize_complex_metadata,
+    serialize_complex_metadata,
+)
+
+if TYPE_CHECKING:
+    from lightly_studio.models.sample import SampleTable
+else:
+    SampleTable = object
+
+
+TYPE_TO_NAME_MAP = {
+    bool: "boolean",
+    int: "integer",
+    float: "float",
+    str: "string",
+    list: "list",
+    dict: "dict",
+}
+
+NAME_TO_TYPE_MAP = {
+    "string": str,
+    "integer": int,
+    "float": float,
+    "boolean": bool,
+    "list": list,
+    "dict": dict,
+}
+
+
+def get_type_name(value: Any) -> str:
+    """Get the type name for a value.
+
+    Args:
+        value: The value to get the type name for.
+
+    Returns:
+        The type name as a string.
+    """
+    if value is None:
+        return "null"
+    # Check if it's a complex metadata type.
+    for name, cls in COMPLEX_METADATA_TYPES.items():
+        if isinstance(value, cls):
+            return name
+
+    # Return mapped type name or fallback to class name.
+    return TYPE_TO_NAME_MAP.get(type(value), type(value).__name__.lower())
+
+
+def validate_type_compatibility(expected_type: str, value: Any) -> bool:
+    """Validate that a value is compatible with an expected type.
+
+    Args:
+        expected_type: The expected type name.
+        value: The value to validate.
+
+    Returns:
+        True if compatible, False otherwise.
+    """
+    if value is None:
+        return expected_type == "null"
+
+    # Check complex types.
+    if expected_type in COMPLEX_METADATA_TYPES:
+        expected_complex_cls = COMPLEX_METADATA_TYPES[expected_type]
+        assert expected_complex_cls is not None
+        return isinstance(value, expected_complex_cls)
+
+    # Check simple types
+    expected_cls = NAME_TO_TYPE_MAP.get(expected_type)
+    if expected_cls is None:
+        return False
+
+    return isinstance(value, expected_cls)
+
+
+class MetadataBase(SQLModel):
+    """Base class for CustomMetadata models."""
+
+    custom_metadata_id: UUID = Field(default_factory=uuid4, primary_key=True)
+    created_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+    )
+    updated_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+    )
+    # Dictionary storing the actual metadata values as JSON.
+    data: dict[str, Any] = Field(
+        default_factory=dict,
+        sa_type=JSON,
+        description="Custom metadata stored as JSON",
+    )
+    # Dictionary storing the metadata schema.
+    metadata_schema: dict[str, str] = Field(
+        default_factory=dict,
+        sa_type=JSON,
+        description="Schema information for metadata keys",
+    )
+
+    def ensure_schema(self, key: str, value: Any) -> None:
+        """Ensure schema exists for a key and validate value type.
+
+        This method handles schema management for metadata keys:
+         If the key doesn't exist in the schema, it creates a new entry
+         with the inferred type from the provided value
+         If the key exists, it validates that the new value matches
+         the expected type from the schema
+
+        This ensures type consistency across metadata operations and prevents
+        accidental type mismatches that could cause issues in applications.
+
+        Args:
+            key: The metadata key to validate/update schema for
+            value: The value to validate/use for type inference
+
+        Raises:
+            ValueError: If the value type doesn't match existing schema
+        """
+        if key not in self.metadata_schema:
+            # New key - create schema with actual type name.
+            self.metadata_schema[key] = get_type_name(value)
+        else:
+            # Existing key - validate type.
+            existing_type = self.metadata_schema[key]
+            if not validate_type_compatibility(existing_type, value):
+                raise ValueError(
+                    f"Value type mismatch for key '{key}'. "
+                    f"Expected {existing_type}, got {get_type_name(value)}."
+                )
+
+    def set_value(self, key: str, value: Any) -> None:
+        """Set a metadata value with schema validation and database tracking.
+
+        Args:
+            key: The metadata key to set
+            value: The value to set
+
+        Raises:
+            ValueError: If the value type doesn't match the schema
+        """
+        self.ensure_schema(key, value)
+        # Serialize complex metadata for storage.
+        self.data[key] = serialize_complex_metadata(value)
+        self.updated_at = datetime.now(timezone.utc)
+        # Mark the object as modified so SQLAlchemy knows to update it.
+        flag_modified(self, "data")
+        flag_modified(self, "metadata_schema")
+
+    def get_value(self, key: str) -> Any:
+        """Get a metadata value with automatic deserialization.
+
+        Args:
+            key: The metadata key.
+
+        Returns:
+            The deserialized value (complex metadata object if applicable)
+            or None if the key doesn't exist.
+        """
+        value = self.data.get(key)
+        if value is not None:
+            # Get expected type from schema for deserialization.
+            expected_type = self.metadata_schema.get(key)
+            if expected_type:
+                return deserialize_complex_metadata(value, expected_type)
+        return value
+
+
+class MetadataCreate(MetadataBase):
+    """Input class for Metadata model."""
+
+
+class SampleMetadataTable(MetadataBase, table=True):
+    """This class defines the SampleMetadataTable model."""
+
+    __tablename__ = "metadata"
+    sample_id: UUID = Field(foreign_key="samples.sample_id")
+
+    sample: SampleTable = Relationship(back_populates="metadata_dict")
+
+
+class SampleMetadataView(SQLModel):
+    """Sample metadata class when retrieving."""
+
+    data: dict[str, Any]
+
+
+class MetadataInfoView(BaseModel):
+    """Metadata info response model for API endpoints."""
+
+    name: str = Field(description="The metadata key name")
+    type: str = Field(description="The metadata type (e.g., 'string', 'integer', 'float')")
+    min: int | float | None = Field(None, description="Minimum value for numerical metadata")
+    max: int | float | None = Field(None, description="Maximum value for numerical metadata")

lightly_studio/models/sample.py

@@ -0,0 +1,180 @@
+"""This module defines the User model for the application."""
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, List, Literal, Optional
+from uuid import UUID, uuid4
+
+from sqlalchemy.orm import Mapped, Session
+from sqlmodel import Field, Relationship, SQLModel
+
+from lightly_studio.models.annotation.annotation_base import AnnotationView
+from lightly_studio.resolvers import metadata_resolver
+
+if TYPE_CHECKING:
+    from lightly_studio.models.annotation.annotation_base import (
+        AnnotationBaseTable,
+    )
+    from lightly_studio.models.metadata import (
+        SampleMetadataTable,
+        SampleMetadataView,
+    )
+    from lightly_studio.models.sample_embedding import SampleEmbeddingTable
+    from lightly_studio.models.tag import TagTable
+else:
+    AnnotationBaseTable = object
+    SampleEmbeddingTable = object
+    SampleMetadataTable = object
+    TagTable = object
+    SampleMetadataView = object
+
+
+class SampleBase(SQLModel):
+    """Base class for the Sample model."""
+
+    """The name of the image file."""
+    file_name: str
+
+    """The width of the image in pixels."""
+    width: int
+
+    """The height of the image in pixels."""
+    height: int
+
+    """The dataset ID to which the sample belongs."""
+    dataset_id: UUID = Field(default=None, foreign_key="datasets.dataset_id")
+
+    """The dataset image path."""
+    file_path_abs: str
+
+
+class SampleCreate(SampleBase):
+    """Sample class when inserting."""
+
+
+class SampleViewForAnnotation(SQLModel):
+    """Sample class for annotation view."""
+
+    """The name of the image file."""
+    file_path_abs: str
+    sample_id: UUID
+
+    """The width of the image in pixels."""
+    width: int
+
+    """The height of the image in pixels."""
+    height: int
+
+    created_at: datetime
+    updated_at: datetime
+
+
+class SampleTagLinkTable(SQLModel, table=True):
+    """Model to define links between Sample and Tag Many-to-Many."""
+
+    sample_id: Optional[UUID] = Field(
+        default=None, foreign_key="samples.sample_id", primary_key=True
+    )
+    tag_id: Optional[UUID] = Field(default=None, foreign_key="tags.tag_id", primary_key=True)
+
+
+class SampleTable(SampleBase, table=True):
+    """This class defines the Sample model."""
+
+    __tablename__ = "samples"
+    sample_id: UUID = Field(default_factory=uuid4, primary_key=True)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), index=True)
+    updated_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+    )
+    annotations: Mapped[List["AnnotationBaseTable"]] = Relationship(
+        back_populates="sample",
+    )
+
+    """The tag ids associated with the sample."""
+    tags: Mapped[List["TagTable"]] = Relationship(
+        back_populates="samples", link_model=SampleTagLinkTable
+    )
+    embeddings: Mapped[List["SampleEmbeddingTable"]] = Relationship(back_populates="sample")
+    metadata_dict: "SampleMetadataTable" = Relationship(back_populates="sample")
+
+    def __getitem__(self, key: str) -> Any:
+        """Provides dict-like access to sample metadata.
+
+        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.metadata_dict is None:
+            return None
+        return self.metadata_dict.get_value(key)
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        """Sets a metadata key-value pair for this sample.
+
+        Args:
+            key: The metadata key.
+            value: The metadata value.
+
+        Note:
+            If the sample has no metadata, a new Metadata Table instance
+            will be created. Changes are automatically committed to the
+            database.
+
+        Raises:
+            RuntimeError: If no database session is found.
+        """
+        # Get the session from the instance
+        session = Session.object_session(self)
+        if session is None:
+            raise RuntimeError("No database session found for this instance")
+
+        # Use metadata_resolver to handle the database operations.
+        # Added type: ignore to avoid type checking issues. SQLAlchemy and
+        # SQLModel sessions are compatible at runtime but have different type
+        # annotations.
+        metadata_resolver.set_value_for_sample(
+            session=session,  # type: ignore[arg-type]
+            sample_id=self.sample_id,
+            key=key,
+            value=value,
+        )
+
+
+TagKind = Literal[
+    "sample",
+    "annotation",
+]
+
+
+class SampleView(SQLModel):
+    """Sample class when retrieving."""
+
+    class SampleViewTag(SQLModel):
+        """Tag view inside Sample view."""
+
+        tag_id: UUID
+        name: str
+        kind: TagKind
+        created_at: datetime
+        updated_at: datetime
+
+    """The name of the image file."""
+    file_name: str
+    file_path_abs: str
+    sample_id: UUID
+    dataset_id: UUID
+    annotations: List["AnnotationView"] = Field([])
+    tags: List[SampleViewTag] = Field([])
+    metadata_dict: Optional["SampleMetadataView"] = None
+    width: int
+    height: int
+
+
+class SampleViewsWithCount(SQLModel):
+    """Response model for counted samples."""
+
+    data: List[SampleView]
+    total_count: int

lightly_studio/models/sample_embedding.py

@@ -0,0 +1,37 @@
+"""This module defines the SampleEmbedding model for the application."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from uuid import UUID, uuid4
+
+from sqlalchemy import ARRAY, Float
+from sqlmodel import Column, Field, Relationship, SQLModel
+
+if TYPE_CHECKING:
+    from lightly_studio.models.sample import SampleTable
+
+else:
+    SampleTable = object
+
+
+class SampleEmbeddingBase(SQLModel):
+    """Base class for the Embeddings used for Samples."""
+
+    sample_embedding_id: UUID = Field(default_factory=uuid4, primary_key=True)
+
+    sample_id: UUID = Field(foreign_key="samples.sample_id")
+
+    embedding_model_id: UUID = Field(foreign_key="embedding_models.embedding_model_id")
+    embedding: list[float] = Field(sa_column=Column(ARRAY(Float)))
+
+
+class SampleEmbeddingCreate(SampleEmbeddingBase):
+    """Sample class when inserting."""
+
+
+class SampleEmbeddingTable(SampleEmbeddingBase, table=True):
+    """This class defines the SampleEmbedding model."""
+
+    __tablename__ = "sample_embeddings"
+    sample: SampleTable = Relationship(back_populates="embeddings")

lightly_studio/models/settings.py

@@ -0,0 +1,60 @@
+"""This module contains settings model for user preferences."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from enum import Enum
+from uuid import UUID, uuid4
+
+from sqlmodel import Field, SQLModel
+
+
+class GridViewSampleRenderingType(str, Enum):
+    """Defines how samples are rendered in the grid view."""
+
+    COVER = "cover"
+    CONTAIN = "contain"
+
+
+class SettingBase(SQLModel):
+    """Base class for Settings model."""
+
+    grid_view_sample_rendering: GridViewSampleRenderingType = Field(
+        default=GridViewSampleRenderingType.CONTAIN,
+        description="Controls how samples are rendered in the grid view",
+    )
+
+    # Keyboard shortcuts.
+    key_hide_annotations: str = Field(
+        default="v",
+        description="Key to temporarily hide annotations while pressed",
+    )
+    key_go_back: str = Field(
+        default="Escape",
+        description="Key to navigate back from detail view to grid view",
+    )
+
+    # New setting for annotation text visibility.
+    show_annotation_text_labels: bool = Field(
+        default=True,
+        description="Controls whether to show text labels on annotations",
+    )
+
+
+class SettingView(SettingBase):
+    """View class for Settings model."""
+
+    setting_id: UUID
+    created_at: datetime
+    updated_at: datetime
+
+
+class SettingTable(SettingBase, table=True):
+    """This class defines the Setting model."""
+
+    __tablename__ = "settings"
+    setting_id: UUID = Field(default_factory=uuid4, primary_key=True)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), index=True)
+    updated_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+    )