lightly-studio 0.3.1__py3-none-any.whl

lightly_studio/resolvers/metadata_resolver/sample/get_metadata_info.py

@@ -0,0 +1,104 @@
+"""Resolver for operations for retrieving metadata info."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlalchemy import Float, func
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.metadata import (
+    MetadataInfoView,
+    SampleMetadataTable,
+)
+from lightly_studio.models.sample import SampleTable
+
+
+def get_all_metadata_keys_and_schema(
+    session: Session,
+    dataset_id: UUID,
+) -> list[MetadataInfoView]:
+    """Get all unique metadata keys and their schema for a dataset.
+
+    Args:
+        session: The database session.
+        dataset_id: The dataset's UUID.
+
+    Returns:
+        List of dicts with 'name', 'type', and optionally 'min'/'max' for numerical types.
+    """
+    # Query all metadata_schema dicts for samples in the dataset
+    rows = session.exec(
+        select(SampleMetadataTable.metadata_schema)
+        .select_from(SampleTable)
+        .join(
+            SampleMetadataTable,
+            col(SampleMetadataTable.sample_id) == col(SampleTable.sample_id),
+        )
+        .where(SampleTable.dataset_id == dataset_id)
+    ).all()
+    # Merge all schemas
+    merged: dict[str, str] = {}
+    for schema_dict in rows:
+        merged.update(schema_dict)
+
+    # Get min and max values for numerical metadata
+    result = []
+    for key, metadata_type in merged.items():
+        metadata_info = MetadataInfoView(name=key, type=metadata_type)
+
+        # Add min and max for numerical types
+        if metadata_type in ["integer", "float"]:
+            min_max_values = _get_metadata_min_max_values(session, dataset_id, key, metadata_type)
+            if min_max_values:
+                metadata_info.min = min_max_values[0]
+                metadata_info.max = min_max_values[1]
+
+        result.append(metadata_info)
+
+    return result
+
+
+def _get_metadata_min_max_values(
+    session: Session,
+    dataset_id: UUID,
+    metadata_key: str,
+    metadata_type: str,
+) -> tuple[int, int] | tuple[float, float] | None:
+    """Get min and max values for a specific numerical metadata key.
+
+    Args:
+        session: The database session.
+        dataset_id: The dataset's UUID.
+        metadata_key: The metadata key to get min/max for.
+        metadata_type: The metadata type ("integer" or "float").
+
+    Returns:
+        Tuple with 'min' and 'max' values, or None if no values found.
+    """
+    # Build JSON path for the metadata key.
+    json_path = f"$.{metadata_key}"
+
+    query = (
+        select(
+            func.min(func.cast(func.json_extract(SampleMetadataTable.data, json_path), Float)),
+            func.max(func.cast(func.json_extract(SampleMetadataTable.data, json_path), Float)),
+        )
+        .select_from(SampleTable)
+        .join(SampleMetadataTable, col(SampleMetadataTable.sample_id) == col(SampleTable.sample_id))
+        .where(
+            SampleTable.dataset_id == dataset_id,
+            func.json_extract(SampleMetadataTable.data, json_path).is_not(None),
+        )
+    )
+
+    result = session.exec(query).first()
+
+    if result and result[0] is not None and result[1] is not None:
+        # Convert to appropriate type
+        if metadata_type == "integer":
+            return int(result[0]), int(result[1])
+        if metadata_type == "float":
+            return float(result[0]), float(result[1])
+
+    return None

lightly_studio/resolvers/metadata_resolver/sample/get_value_for_sample.py

@@ -0,0 +1,27 @@
+"""Resolver for operations for retrieving metadata."""
+
+from __future__ import annotations
+
+from typing import Any
+from uuid import UUID
+
+from sqlmodel import Session
+
+from .get_by_sample_id import get_by_sample_id
+
+
+def get_value_for_sample(session: Session, sample_id: UUID, key: str) -> Any | None:
+    """Get a specific metadata value for a sample.
+
+    Args:
+        session: The database session.
+        sample_id: The sample's UUID.
+        key: The metadata key.
+
+    Returns:
+        The value for the given key, or None if not found.
+    """
+    metadata = get_by_sample_id(session=session, sample_id=sample_id)
+    if metadata is None:
+        return None
+    return metadata.data.get(key)

lightly_studio/resolvers/metadata_resolver/sample/set_value_for_sample.py

@@ -0,0 +1,53 @@
+"""Resolver for operations for setting metadata."""
+
+from __future__ import annotations
+
+from typing import Any
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.models.metadata import (
+    SampleMetadataTable,
+)
+from lightly_studio.resolvers.metadata_resolver.sample.get_by_sample_id import (
+    get_by_sample_id,
+)
+
+
+def set_value_for_sample(
+    session: Session,
+    sample_id: UUID,
+    key: str,
+    value: Any,
+) -> SampleMetadataTable:
+    """Set a specific metadata value for a sample.
+
+    Args:
+        session: The database session.
+        sample_id: The sample's UUID.
+        key: The metadata key.
+        value: The value to set.
+
+    Returns:
+        The updated CustomMetadataTable instance.
+
+    Raises:
+        ValueError: If the value type doesn't match the schema.
+    """
+    metadata = get_by_sample_id(session=session, sample_id=sample_id)
+    if metadata is None:
+        # Create new metadata row if it does not exist
+        metadata = SampleMetadataTable(
+            sample_id=sample_id,
+            data={},
+            metadata_schema={},
+        )
+        session.add(metadata)
+
+    metadata.set_value(key, value)
+
+    # Commit changes and refresh the object
+    session.commit()
+    session.refresh(metadata)
+    return metadata

lightly_studio/resolvers/sample_embedding_resolver.py

@@ -0,0 +1,86 @@
+"""Handler for database operations related to sample embeddings."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlalchemy import String, cast
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.sample_embedding import (
+    SampleEmbeddingCreate,
+    SampleEmbeddingTable,
+)
+
+
+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,
+) -> 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.
+
+    Returns:
+        List of sample embeddings associated with the dataset.
+    """
+    query = (
+        select(SampleEmbeddingTable)
+        .join(SampleTable)
+        .where(SampleEmbeddingTable.sample_id == SampleTable.sample_id)
+        .where(SampleTable.dataset_id == dataset_id)
+        .where(SampleEmbeddingTable.embedding_model_id == embedding_model_id)
+        .order_by(col(SampleTable.created_at).asc())
+    )
+    return list(session.exec(query).all())

lightly_studio/resolvers/sample_resolver.py

@@ -0,0 +1,249 @@
+"""Handler for database operations related to samples."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from datetime import datetime, timezone
+from uuid import UUID
+
+from pydantic import BaseModel
+from sqlmodel import Session, col, func, select
+from sqlmodel.sql.expression import Select
+
+from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable
+from lightly_studio.models.annotation_label import AnnotationLabelTable
+from lightly_studio.models.embedding_model import EmbeddingModelTable
+from lightly_studio.models.sample import SampleCreate, SampleTable
+from lightly_studio.models.sample_embedding import SampleEmbeddingTable
+from lightly_studio.models.tag import TagTable
+from lightly_studio.resolvers.samples_filter import SampleFilter
+
+
+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
+
+
+def create_many(session: Session, samples: list[SampleCreate]) -> list[SampleTable]:
+    """Create multiple samples in a single database commit."""
+    db_samples = [SampleTable.model_validate(sample) for sample in samples]
+    session.bulk_save_objects(db_samples)
+    session.commit()
+    return db_samples
+
+
+def get_by_id(session: Session, dataset_id: UUID, sample_id: UUID) -> SampleTable | None:
+    """Retrieve a single sample by ID."""
+    return session.exec(
+        select(SampleTable).where(
+            SampleTable.sample_id == sample_id, SampleTable.dataset_id == dataset_id
+        )
+    ).one_or_none()
+
+
+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]
+
+
+class GetAllSamplesByDatasetIdResult(BaseModel):
+    """Result of getting all samples."""
+
+    samples: Sequence[SampleTable]
+    total_count: int
+
+
+def get_all_by_dataset_id(  # noqa: PLR0913
+    session: Session,
+    dataset_id: UUID,
+    offset: int = 0,
+    limit: int | None = None,
+    filters: SampleFilter | None = None,
+    text_embedding: list[float] | None = None,
+    sample_ids: list[UUID] | None = None,
+) -> GetAllSamplesByDatasetIdResult:
+    """Retrieve samples for a specific dataset with optional filtering."""
+    samples_query = select(SampleTable).where(SampleTable.dataset_id == dataset_id)
+    total_count_query = (
+        select(func.count()).select_from(SampleTable).where(SampleTable.dataset_id == dataset_id)
+    )
+
+    if filters:
+        samples_query = filters.apply(samples_query)
+        total_count_query = filters.apply(total_count_query)
+
+    # TODO(Michal, 06/2025): Consider adding sample_ids to the filters.
+    if sample_ids:
+        samples_query = samples_query.where(col(SampleTable.sample_id).in_(sample_ids))
+        total_count_query = total_count_query.where(col(SampleTable.sample_id).in_(sample_ids))
+
+    if text_embedding:
+        # Fetch the first embedding_model_id for the given dataset_id
+        embedding_model_id = session.exec(
+            select(EmbeddingModelTable.embedding_model_id)
+            .where(EmbeddingModelTable.dataset_id == dataset_id)
+            .limit(1)
+        ).first()
+        if embedding_model_id:
+            # Join with SampleEmbedding table to access embeddings
+            samples_query = (
+                samples_query.join(
+                    SampleEmbeddingTable,
+                    col(SampleTable.sample_id) == col(SampleEmbeddingTable.sample_id),
+                )
+                .where(SampleEmbeddingTable.embedding_model_id == embedding_model_id)
+                .order_by(
+                    func.list_cosine_distance(
+                        SampleEmbeddingTable.embedding,
+                        text_embedding,
+                    )
+                )
+            )
+            total_count_query = total_count_query.join(
+                SampleEmbeddingTable,
+                col(SampleTable.sample_id) == col(SampleEmbeddingTable.sample_id),
+            ).where(SampleEmbeddingTable.embedding_model_id == embedding_model_id)
+    else:
+        samples_query = samples_query.order_by(
+            col(SampleTable.created_at).asc(), col(SampleTable.sample_id).asc()
+        )
+
+    # paginate query when offset or limit are set/positive
+    if offset > 0:
+        samples_query = samples_query.offset(offset)
+    if limit is not None:
+        samples_query = samples_query.limit(limit)
+
+    return GetAllSamplesByDatasetIdResult(
+        samples=session.exec(samples_query).all(),
+        total_count=session.exec(total_count_query).one(),
+    )
+
+
+def get_dimension_bounds(
+    session: Session,
+    dataset_id: UUID,
+    annotation_label_ids: list[UUID] | None = None,
+    tag_ids: list[UUID] | None = None,
+) -> dict[str, int]:
+    """Get min and max dimensions of samples in a dataset."""
+    # Prepare the base query for dimensions
+    query: Select[tuple[int | None, int | None, int | None, int | None]] = select(
+        func.min(SampleTable.width).label("min_width"),
+        func.max(SampleTable.width).label("max_width"),
+        func.min(SampleTable.height).label("min_height"),
+        func.max(SampleTable.height).label("max_height"),
+    )
+
+    if annotation_label_ids:
+        # Subquery to filter samples matching all annotation labels
+        label_filter = (
+            select(SampleTable.sample_id)
+            .join(
+                AnnotationBaseTable,
+                col(SampleTable.sample_id) == col(AnnotationBaseTable.sample_id),
+            )
+            .join(
+                AnnotationLabelTable,
+                col(AnnotationBaseTable.annotation_label_id)
+                == col(AnnotationLabelTable.annotation_label_id),
+            )
+            .where(
+                SampleTable.dataset_id == dataset_id,
+                col(AnnotationLabelTable.annotation_label_id).in_(annotation_label_ids),
+            )
+            .group_by(col(SampleTable.sample_id))
+            .having(
+                func.count(col(AnnotationLabelTable.annotation_label_id).distinct())
+                == len(annotation_label_ids)
+            )
+        )
+        # Filter the dimension query based on the subquery
+        query = query.where(col(SampleTable.sample_id).in_(label_filter))
+    else:
+        # If no labels specified, filter dimensions
+        # for all samples in the dataset
+        query = query.where(SampleTable.dataset_id == dataset_id)
+
+    if tag_ids:
+        query = (
+            query.join(SampleTable.tags)
+            .where(SampleTable.tags.any(col(TagTable.tag_id).in_(tag_ids)))
+            .distinct()
+        )
+
+    # Note: We use SQLAlchemy's session.execute instead of SQLModel's
+    # ession.exec to be able to fetch the columns with names with the
+    # `mappings()` method.
+    result = session.execute(query).mappings().one()
+    return {key: value for key, value in result.items() if value is not None}
+
+
+def update(session: Session, sample_id: UUID, sample_data: SampleCreate) -> SampleTable | None:
+    """Update an existing sample."""
+    sample = get_by_id(session=session, dataset_id=sample_data.dataset_id, sample_id=sample_id)
+    if not sample:
+        return None
+
+    sample.file_name = sample_data.file_name
+    sample.width = sample_data.width
+    sample.height = sample_data.height
+    sample.updated_at = datetime.now(timezone.utc)
+
+    session.commit()
+    session.refresh(sample)
+    return sample
+
+
+def delete(session: Session, dataset_id: UUID, sample_id: UUID) -> bool:
+    """Delete a sample."""
+    sample = get_by_id(session=session, dataset_id=dataset_id, sample_id=sample_id)
+    if not sample:
+        return False
+
+    session.delete(sample)
+    session.commit()
+    return True
+
+
+def get_samples_excluding(
+    session: Session,
+    dataset_id: UUID,
+    excluded_sample_ids: list[UUID],
+    limit: int | None = None,
+) -> Sequence[SampleTable]:
+    """Get random samples excluding specified sample IDs.
+
+    Args:
+        session: The database session.
+        dataset_id: The dataset ID to filter by.
+        excluded_sample_ids: List of sample IDs to exclude from the result.
+        limit: Maximum number of samples to return.
+                If None, returns all matches.
+
+    Returns:
+        List of samples not associated with the excluded IDs.
+    """
+    query = (
+        select(SampleTable)
+        .where(SampleTable.dataset_id == dataset_id)
+        .where(col(SampleTable.sample_id).not_in(excluded_sample_ids))
+        .order_by(func.random())
+    )
+
+    if limit is not None:
+        query = query.limit(limit)
+
+    return session.exec(query).all()

lightly_studio/resolvers/samples_filter.py

@@ -0,0 +1,81 @@
+"""Utility functions for building database queries."""
+
+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.metadata_filter import (
+    MetadataFilter,
+    apply_metadata_filters,
+)
+from lightly_studio.type_definitions import QueryType
+
+
+class FilterDimensions(BaseModel):
+    """Encapsulates dimension-based filter parameters for querying samples."""
+
+    min: Optional[int] = None
+    max: Optional[int] = None
+
+
+class SampleFilter(BaseModel):
+    """Encapsulates filter parameters for querying samples."""
+
+    width: Optional[FilterDimensions] = None
+    height: Optional[FilterDimensions] = None
+    annotation_label_ids: Optional[List[UUID]] = None
+    tag_ids: Optional[List[UUID]] = None
+    metadata_filters: Optional[List[MetadataFilter]] = None
+
+    def apply(self, query: QueryType) -> QueryType:
+        """Apply the filters to the given query."""
+        # Apply dimension-based filters to the query.
+        if self.width:
+            if self.width.min is not None:
+                query = query.where(SampleTable.width >= self.width.min)
+            if self.width.max is not None:
+                query = query.where(SampleTable.width <= self.width.max)
+        if self.height:
+            if self.height.min is not None:
+                query = query.where(SampleTable.height >= self.height.min)
+            if self.height.max is not None:
+                query = query.where(SampleTable.height <= self.height.max)
+
+        # Apply annotation label filters to the query.
+        if self.annotation_label_ids:
+            sample_ids_subquery = (
+                select(AnnotationBaseTable.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)
+                .select_from(SampleTable)
+                .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 = apply_metadata_filters(
+                query,
+                self.metadata_filters,
+                metadata_model=SampleMetadataTable,
+                metadata_join_condition=SampleMetadataTable.sample_id == SampleTable.sample_id,
+            )
+        return query

lightly_studio/resolvers/settings_resolver.py

@@ -0,0 +1,58 @@
+"""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
+
+    # Update show annotation text labels
+    current_settings.show_annotation_text_labels = settings.show_annotation_text_labels
+
+    session.commit()
+    session.refresh(current_settings)
+
+    return SettingView.model_validate(current_settings)