lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/resolvers/sample_embedding_resolver.py

@@ -0,0 +1,132 @@
+"""Handler for database operations related to sample embeddings."""
+
+from __future__ import annotations
+
+import hashlib
+from uuid import UUID
+
+from sqlalchemy import String, cast, func
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.sample_embedding import (
+    SampleEmbeddingCreate,
+    SampleEmbeddingTable,
+)
+from lightly_studio.resolvers.sample_resolver.sample_filter import SampleFilter
+
+
+def create(session: Session, sample_embedding: SampleEmbeddingCreate) -> SampleEmbeddingTable:
+    """Create a new SampleEmbedding in the database."""
+    db_sample_embedding = SampleEmbeddingTable.model_validate(sample_embedding)
+    session.add(db_sample_embedding)
+    session.commit()
+    session.refresh(db_sample_embedding)
+    return db_sample_embedding
+
+
+def create_many(session: Session, sample_embeddings: list[SampleEmbeddingCreate]) -> None:
+    """Create many sample embeddings in a single database commit."""
+    db_sample_embeddings = [SampleEmbeddingTable.model_validate(e) for e in sample_embeddings]
+    session.bulk_save_objects(db_sample_embeddings)
+    session.commit()
+
+
+def get_by_sample_ids(
+    session: Session,
+    sample_ids: list[UUID],
+    embedding_model_id: UUID,
+) -> list[SampleEmbeddingTable]:
+    """Get sample embeddings for the specified sample IDs.
+
+    Output order matches the input order.
+
+    Args:
+        session: The database session.
+        sample_ids: List of sample IDs to get embeddings for.
+        embedding_model_id: The embedding model ID to filter by.
+
+    Returns:
+        List of sample embeddings associated with the provided IDs.
+    """
+    string_ids = [str(id_) for id_ in sample_ids]
+    results = list(
+        session.exec(
+            select(SampleEmbeddingTable)
+            .where(cast(SampleEmbeddingTable.sample_id, String).in_(string_ids))
+            .where(SampleEmbeddingTable.embedding_model_id == embedding_model_id)
+        ).all()
+    )
+    # Return embeddings in the same order as the input IDs
+    embedding_map = {embedding.sample_id: embedding for embedding in results}
+    return [embedding_map[id_] for id_ in sample_ids if id_ in embedding_map]
+
+
+def get_all_by_dataset_id(
+    session: Session,
+    dataset_id: UUID,
+    embedding_model_id: UUID,
+    filters: SampleFilter | None = None,
+) -> list[SampleEmbeddingTable]:
+    """Get all sample embeddings for samples in a specific dataset.
+
+    Args:
+        session: The database session.
+        dataset_id: The dataset ID to filter by.
+        embedding_model_id: The embedding model ID to filter by.
+        filters: Filters to apply to the samples.
+
+    Returns:
+        List of sample embeddings associated with the dataset.
+    """
+    query = (
+        select(SampleEmbeddingTable)
+        .join(SampleTable, col(SampleEmbeddingTable.sample_id) == col(SampleTable.sample_id))
+        .where(SampleTable.dataset_id == dataset_id)
+        .where(SampleEmbeddingTable.embedding_model_id == embedding_model_id)
+        .order_by(col(SampleTable.created_at).asc())
+    )
+    if filters:
+        query = filters.apply(query)
+    return list(session.exec(query).all())
+
+
+def get_hash_by_sample_ids(
+    session: Session,
+    sample_ids_ordered: list[UUID],
+    embedding_model_id: UUID,
+) -> tuple[str, list[UUID]]:
+    """Return a combined hash and the ordered sample IDs with stored embeddings.
+
+    Args:
+        session: Database session.
+        sample_ids_ordered: Sample IDs to consider, order defines deterministic hash.
+        embedding_model_id: Embedding model identifier.
+
+    Returns:
+        Tuple of (combined hash, ordered sample IDs that have stored embeddings).
+    """
+    if not sample_ids_ordered:
+        return "empty", []
+
+    rows = session.exec(
+        select(
+            SampleEmbeddingTable.sample_id,
+            func.hash(SampleEmbeddingTable.embedding).label("hash_column"),
+        )
+        .where(col(SampleEmbeddingTable.sample_id).in_(set(sample_ids_ordered)))
+        .where(SampleEmbeddingTable.embedding_model_id == embedding_model_id)
+    ).all()
+
+    # Mypy does not get that 'hash_column' is an attribute of the returned rows
+    sample_id_to_hash = {row.sample_id: row.hash_column for row in rows}  # type: ignore[attr-defined]
+    sample_ids_of_samples_with_embeddings = [
+        sample_id for sample_id in sample_ids_ordered if sample_id in sample_id_to_hash
+    ]
+    hashes_ordered = [
+        sample_id_to_hash[sample_id] for sample_id in sample_ids_of_samples_with_embeddings
+    ]
+
+    hasher = hashlib.sha256()
+    hasher.update("".join(str(h) for h in hashes_ordered).encode("utf-8"))
+    return hasher.hexdigest(), sample_ids_of_samples_with_embeddings

lightly_studio/resolvers/sample_resolver/__init__.py

@@ -0,0 +1,17 @@
+"""Resolvers for database operations."""
+
+from lightly_studio.resolvers.sample_resolver.count_by_dataset_id import count_by_dataset_id
+from lightly_studio.resolvers.sample_resolver.create import create
+from lightly_studio.resolvers.sample_resolver.create_many import create_many
+from lightly_studio.resolvers.sample_resolver.get_by_id import get_by_id
+from lightly_studio.resolvers.sample_resolver.get_filtered_samples import get_filtered_samples
+from lightly_studio.resolvers.sample_resolver.get_many_by_id import get_many_by_id
+
+__all__ = [
+    "count_by_dataset_id",
+    "create",
+    "create_many",
+    "get_by_id",
+    "get_filtered_samples",
+    "get_many_by_id",
+]

lightly_studio/resolvers/sample_resolver/count_by_dataset_id.py

@@ -0,0 +1,16 @@
+"""Implementation of count_by_dataset_id for sample resolver."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, func, select
+
+from lightly_studio.models.sample import SampleTable
+
+
+def count_by_dataset_id(session: Session, dataset_id: UUID) -> int:
+    """Count the number of samples in a dataset."""
+    return session.exec(
+        select(func.count()).select_from(SampleTable).where(SampleTable.dataset_id == dataset_id)
+    ).one()

lightly_studio/resolvers/sample_resolver/create.py

@@ -0,0 +1,16 @@
+"""Implementation of create for sample resolver."""
+
+from __future__ import annotations
+
+from sqlmodel import Session
+
+from lightly_studio.models.sample import SampleCreate, SampleTable
+
+
+def create(session: Session, sample: SampleCreate) -> SampleTable:
+    """Create a new sample in the database."""
+    db_sample = SampleTable.model_validate(sample)
+    session.add(db_sample)
+    session.commit()
+    session.refresh(db_sample)
+    return db_sample

lightly_studio/resolvers/sample_resolver/create_many.py

@@ -0,0 +1,25 @@
+"""Implementation of create_many for sample resolver."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from uuid import UUID
+
+from sqlalchemy import ScalarResult
+from sqlmodel import Session, col, insert
+
+from lightly_studio.models.sample import SampleCreate, SampleTable
+
+
+def create_many(session: Session, samples: Sequence[SampleCreate]) -> list[UUID]:
+    """Create multiple samples in a single database commit."""
+    if not samples:
+        return []
+    # Note: We are using bulk insert for SampleTable to get sample_ids efficiently.
+    statement = (
+        insert(SampleTable)
+        .values([sample.model_dump() for sample in samples])
+        .returning(col(SampleTable.sample_id))
+    )
+    sample_ids: ScalarResult[UUID] = session.execute(statement).scalars()
+    return list(sample_ids)

lightly_studio/resolvers/sample_resolver/get_by_id.py

@@ -0,0 +1,14 @@
+"""Implementation of get_by_id for sample resolver."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, select
+
+from lightly_studio.models.sample import SampleTable
+
+
+def get_by_id(session: Session, sample_id: UUID) -> SampleTable | None:
+    """Retrieve a single sample by ID."""
+    return session.exec(select(SampleTable).where(SampleTable.sample_id == sample_id)).one_or_none()

lightly_studio/resolvers/sample_resolver/get_filtered_samples.py

@@ -0,0 +1,56 @@
+"""Implementation of get_filtered_samples resolver function."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from pydantic import BaseModel
+from sqlmodel import Session, col, func, select
+
+from lightly_studio.api.routes.api.validators import Paginated
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.resolvers.sample_resolver.sample_filter import SampleFilter
+
+
+class SamplesWithCount(BaseModel):
+    """Result of getting all samples."""
+
+    samples: Sequence[SampleTable]
+    total_count: int
+    next_cursor: int | None = None
+
+
+def get_filtered_samples(
+    session: Session,
+    filters: SampleFilter | None = None,
+    pagination: Paginated | None = None,
+) -> SamplesWithCount:
+    """Retrieve samples for a specific dataset with optional filtering."""
+    samples_query = select(SampleTable)
+    total_count_query = select(func.count()).select_from(SampleTable)
+
+    if filters is not None:
+        samples_query = filters.apply(samples_query)
+        total_count_query = filters.apply(total_count_query)
+
+    # Apply default ordering
+    samples_query = samples_query.order_by(
+        col(SampleTable.created_at).asc(),
+        col(SampleTable.sample_id).asc(),
+    )
+
+    # Apply pagination if provided
+    if pagination is not None:
+        samples_query = samples_query.offset(pagination.offset).limit(pagination.limit)
+
+    total_count = session.exec(total_count_query).one()
+
+    next_cursor = None
+    if pagination is not None and pagination.offset + pagination.limit < total_count:
+        next_cursor = pagination.offset + pagination.limit
+
+    return SamplesWithCount(
+        samples=session.exec(samples_query).all(),
+        total_count=total_count,
+        next_cursor=next_cursor,
+    )

lightly_studio/resolvers/sample_resolver/get_many_by_id.py

@@ -0,0 +1,22 @@
+"""Implementation of get_many_by_id for sample resolver."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.sample import SampleTable
+
+
+def get_many_by_id(session: Session, sample_ids: list[UUID]) -> list[SampleTable]:
+    """Retrieve multiple samples by their IDs.
+
+    Output order matches the input order.
+    """
+    results = session.exec(
+        select(SampleTable).where(col(SampleTable.sample_id).in_(sample_ids))
+    ).all()
+    # Return samples in the same order as the input IDs
+    sample_map = {sample.sample_id: sample for sample in results}
+    return [sample_map[id_] for id_ in sample_ids if id_ in sample_map]

lightly_studio/resolvers/sample_resolver/sample_filter.py

@@ -0,0 +1,74 @@
+"""SampleFilter class."""
+
+from typing import List, Optional
+from uuid import UUID
+
+from pydantic import BaseModel
+from sqlmodel import col, select
+
+from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable
+from lightly_studio.models.annotation_label import AnnotationLabelTable
+from lightly_studio.models.metadata import SampleMetadataTable
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.tag import TagTable
+from lightly_studio.resolvers.metadata_resolver import metadata_filter
+from lightly_studio.resolvers.metadata_resolver.metadata_filter import MetadataFilter
+from lightly_studio.type_definitions import QueryType
+
+
+class SampleFilter(BaseModel):
+    """Encapsulates filter parameters for querying samples."""
+
+    dataset_id: Optional[UUID] = None
+    annotation_label_ids: Optional[List[UUID]] = None
+    tag_ids: Optional[List[UUID]] = None
+    metadata_filters: Optional[List[MetadataFilter]] = None
+    sample_ids: Optional[List[UUID]] = None
+    has_captions: Optional[bool] = None
+
+    def apply(self, query: QueryType) -> QueryType:
+        """Apply the filters to the given query."""
+        if self.dataset_id:
+            query = query.where(col(SampleTable.dataset_id) == self.dataset_id)
+
+        if self.sample_ids:
+            query = query.where(col(SampleTable.sample_id).in_(self.sample_ids))
+
+        # Apply annotation label filters to the query.
+        if self.annotation_label_ids:
+            sample_ids_subquery = (
+                select(AnnotationBaseTable.parent_sample_id)
+                .select_from(AnnotationBaseTable)
+                .join(AnnotationBaseTable.annotation_label)
+                .where(col(AnnotationLabelTable.annotation_label_id).in_(self.annotation_label_ids))
+                .distinct()
+            )
+            query = query.where(col(SampleTable.sample_id).in_(sample_ids_subquery))
+
+        # Apply tag filters to the query.
+        if self.tag_ids:
+            sample_ids_subquery = (
+                select(SampleTable.sample_id)
+                .join(SampleTable.tags)
+                .where(col(TagTable.tag_id).in_(self.tag_ids))
+                .distinct()
+            )
+            query = query.where(col(SampleTable.sample_id).in_(sample_ids_subquery))
+
+        # Apply metadata filters to the query.
+        if self.metadata_filters:
+            query = metadata_filter.apply_metadata_filters(
+                query,
+                self.metadata_filters,
+                metadata_model=SampleMetadataTable,
+                metadata_join_condition=SampleMetadataTable.sample_id == SampleTable.sample_id,
+            )
+
+        # Apply caption presence filter to the query.
+        if self.has_captions is not None:
+            if self.has_captions:
+                query = query.where(col(SampleTable.captions).any())
+            else:
+                query = query.where(~col(SampleTable.captions).any())
+
+        return query

lightly_studio/resolvers/settings_resolver.py

@@ -0,0 +1,62 @@
+"""This module contains the resolvers for user settings."""
+
+from sqlmodel import Session, select
+
+from lightly_studio.models.settings import SettingTable, SettingView
+
+
+def get_settings(session: Session) -> SettingView:
+    """Get current settings.
+
+    Args:
+        session: Database session.
+
+    Returns:
+        The current settings.
+    """
+    statement = select(SettingTable)
+    result = session.exec(statement).first()
+
+    # If no settings exist, create default settings
+    if result is None:
+        result = SettingTable()
+        session.add(result)
+        session.commit()
+        session.refresh(result)
+
+    return SettingView.model_validate(result)
+
+
+def set_settings(session: Session, settings: SettingView) -> SettingView:
+    """Update settings.
+
+    Args:
+        session: Database session.
+        settings: New settings to apply.
+
+    Returns:
+        Updated settings.
+    """
+    current_settings = session.exec(select(SettingTable)).first()
+    if current_settings is None:
+        current_settings = SettingTable()
+        session.add(current_settings)
+
+    # Update grid view sample rendering
+    current_settings.grid_view_sample_rendering = settings.grid_view_sample_rendering
+
+    # Update keyboard shortcut mapping
+    current_settings.key_hide_annotations = settings.key_hide_annotations
+    current_settings.key_go_back = settings.key_go_back
+    current_settings.key_toggle_edit_mode = settings.key_toggle_edit_mode
+
+    # Update show annotation text labels
+    current_settings.show_annotation_text_labels = settings.show_annotation_text_labels
+
+    # Update sample filename visibility
+    current_settings.show_sample_filenames = settings.show_sample_filenames
+
+    session.commit()
+    session.refresh(current_settings)
+
+    return SettingView.model_validate(current_settings)