lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/models/dataset.py

@@ -0,0 +1,70 @@
+"""This module contains the Dataset model and related enumerations."""
+
+from datetime import datetime, timezone
+from enum import Enum
+from typing import List, Optional
+from uuid import UUID, uuid4
+
+from sqlmodel import Field, Relationship, SQLModel
+
+
+class SampleType(str, Enum):
+    """The type of samples in the dataset."""
+
+    VIDEO = "video"
+    VIDEO_FRAME = "video_frame"
+    IMAGE = "image"
+    ANNOTATION = "annotation"
+    CAPTION = "caption"
+
+
+class DatasetBase(SQLModel):
+    """Base class for the Dataset model."""
+
+    name: str = Field(unique=True, index=True)
+    parent_dataset_id: Optional[UUID] = Field(default=None, foreign_key="dataset.dataset_id")
+    sample_type: SampleType
+
+
+class DatasetCreate(DatasetBase):
+    """Dataset class when inserting."""
+
+
+class DatasetView(DatasetBase):
+    """Dataset class when retrieving."""
+
+    dataset_id: UUID
+    created_at: datetime
+    updated_at: datetime
+    children: List["DatasetView"] = []
+
+
+class DatasetViewWithCount(DatasetView):
+    """Dataset view with total sample count."""
+
+    total_sample_count: int
+
+
+class DatasetOverviewView(SQLModel):
+    """Dataset view for dashboard display."""
+
+    dataset_id: UUID
+    name: str
+    sample_type: SampleType
+    created_at: datetime
+    total_sample_count: int
+
+
+class DatasetTable(DatasetBase, table=True):
+    """This class defines the Dataset model."""
+
+    __tablename__ = "dataset"
+    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))
+
+    parent: Optional["DatasetTable"] = Relationship(
+        back_populates="children",
+        sa_relationship_kwargs={"remote_side": "DatasetTable.dataset_id"},
+    )
+    children: List["DatasetTable"] = Relationship(back_populates="parent")

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="dataset.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_model"
+    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/image.py

@@ -0,0 +1,96 @@
+"""This module defines the User model for the application."""
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, List, Literal, Optional
+from uuid import UUID
+
+from pydantic import BaseModel, ConfigDict
+from pydantic import Field as PydanticField
+from sqlalchemy.orm import Mapped
+from sqlmodel import Field, Relationship, SQLModel
+
+from lightly_studio.models.annotation.annotation_base import AnnotationView
+from lightly_studio.models.caption import CaptionView
+from lightly_studio.models.sample import SampleTable, SampleView
+
+if TYPE_CHECKING:
+    from lightly_studio.models.metadata import (
+        SampleMetadataView,
+    )
+else:
+    SampleMetadataView = object
+
+
+class ImageBase(SQLModel):
+    """Base class for the Image 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 image path."""
+    file_path_abs: str = Field(default=None)
+
+
+class ImageCreate(ImageBase):
+    """Image class when inserting."""
+
+
+class ImageTable(ImageBase, table=True):
+    """This class defines the Image model."""
+
+    __tablename__ = "image"
+    sample_id: UUID = Field(foreign_key="sample.sample_id", 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),
+    )
+    sample: Mapped["SampleTable"] = Relationship()
+
+
+TagKind = Literal[
+    "sample",
+    "annotation",
+]
+
+
+class ImageView(BaseModel):
+    """Image class when retrieving."""
+
+    class ImageViewTag(BaseModel):
+        """Tag view inside Image view."""
+
+        tag_id: UUID
+        name: str
+        kind: TagKind
+        created_at: datetime
+        updated_at: datetime
+
+    file_name: str
+    file_path_abs: str
+    sample_id: UUID
+    annotations: List["AnnotationView"]
+    width: int
+    height: int
+
+    sample: SampleView
+
+    # TODO(Michal, 10/2025): Add SampleView to ImageView, don't expose these fields directly.
+    tags: List[ImageViewTag]
+    metadata_dict: Optional["SampleMetadataView"] = None
+    captions: List[CaptionView] = []
+
+
+class ImageViewsWithCount(BaseModel):
+    """Response model for counted images."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    samples: List[ImageView] = PydanticField(..., alias="data")
+    total_count: int
+    next_cursor: Optional[int] = PydanticField(None, alias="nextCursor")

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="sample.sample_id", unique=True)
+
+    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/range.py

@@ -0,0 +1,17 @@
+"""Models representing numeric ranges."""
+
+from pydantic import BaseModel
+
+
+class FloatRange(BaseModel):
+    """Defines a range of floating-point values."""
+
+    min: float
+    max: float
+
+
+class IntRange(BaseModel):
+    """Defines a range of integer-point values."""
+
+    min: int
+    max: int

lightly_studio/models/sample.py

@@ -0,0 +1,154 @@
+"""This module defines the Sample model for the application."""
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, List, Optional
+from uuid import UUID, uuid4
+
+from pydantic import BaseModel, ConfigDict
+from pydantic import Field as PydanticField
+from sqlalchemy.orm import Mapped, Session
+from sqlmodel import Field, Relationship, SQLModel
+
+from lightly_studio.resolvers import metadata_resolver
+
+if TYPE_CHECKING:
+    from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable, AnnotationView
+    from lightly_studio.models.caption import CaptionTable, CaptionView
+    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
+    TagTable = object
+    SampleEmbeddingTable = object
+    SampleMetadataTable = object
+    SampleMetadataView = object
+    CaptionTable = object
+    CaptionView = object
+    AnnotationView = object
+
+
+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="sample.sample_id", primary_key=True
+    )
+    tag_id: Optional[UUID] = Field(default=None, foreign_key="tag.tag_id", primary_key=True)
+
+
+class SampleBase(SQLModel):
+    """Base class for the Sample model."""
+
+    """The dataset ID to which the sample belongs."""
+    dataset_id: UUID = Field(default=None, foreign_key="dataset.dataset_id")
+
+
+class SampleCreate(SampleBase):
+    """Sample class when inserting."""
+
+
+class SampleTable(SampleBase, table=True):
+    """This class defines the Sample model."""
+
+    __tablename__ = "sample"
+    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),
+    )
+
+    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")
+    annotations: Mapped[List["AnnotationBaseTable"]] = Relationship(
+        back_populates="parent_sample",
+        sa_relationship_kwargs={
+            "lazy": "select",
+            "foreign_keys": "[AnnotationBaseTable.parent_sample_id]",
+        },
+    )
+    captions: Mapped[List["CaptionTable"]] = Relationship(
+        back_populates="parent_sample",
+        sa_relationship_kwargs={
+            "foreign_keys": "[CaptionTable.parent_sample_id]",
+        },
+    )
+
+    # TODO(Michal, 9/2025): Remove this function in favour of Sample.metadata.
+    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)
+
+    # TODO(Michal, 9/2025): Remove this function in favour of Sample.metadata.
+    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,
+        )
+
+
+class SampleView(SampleBase):
+    """This class defines the Sample view model."""
+
+    dataset_id: UUID
+
+    sample_id: UUID
+    created_at: datetime
+    updated_at: datetime
+
+    tags: List["TagTable"] = []
+    metadata_dict: Optional["SampleMetadataView"] = None
+    captions: List[CaptionView] = []
+    annotations: List["AnnotationView"] = []
+
+
+class SampleViewsWithCount(BaseModel):
+    """Result of getting all sample views."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    samples: List[SampleView] = PydanticField(..., alias="data")
+    total_count: int
+    next_cursor: Optional[int] = PydanticField(None, alias="nextCursor")

lightly_studio/models/sample_embedding.py

@@ -0,0 +1,36 @@
+"""This module defines the SampleEmbedding model for the application."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+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_id: UUID = Field(foreign_key="sample.sample_id", primary_key=True)
+    embedding_model_id: UUID = Field(
+        foreign_key="embedding_model.embedding_model_id", primary_key=True
+    )
+    embedding: list[float] = Field(sa_column=Column(ARRAY(Float)))
+
+
+class SampleEmbeddingCreate(SampleEmbeddingBase):
+    """Sample embedding class when inserting."""
+
+
+class SampleEmbeddingTable(SampleEmbeddingBase, table=True):
+    """This class defines the SampleEmbedding model."""
+
+    __tablename__ = "sample_embedding"
+    sample: SampleTable = Relationship(back_populates="embeddings")

lightly_studio/models/settings.py

@@ -0,0 +1,69 @@
+"""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",
+    )
+    key_toggle_edit_mode: str = Field(
+        default="e",
+        description="Key to toggle annotation edit mode",
+    )
+
+    # New setting for annotation text visibility.
+    show_annotation_text_labels: bool = Field(
+        default=True,
+        description="Controls whether to show text labels on annotations",
+    )
+
+    show_sample_filenames: bool = Field(
+        default=False,
+        description="Controls whether to show sample filenames in the samples grid view",
+    )
+
+
+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__ = "setting"
+    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),
+    )