lightly-studio 0.3.1__py3-none-any.whl

lightly_studio/resolvers/dataset_resolver.py

@@ -0,0 +1,278 @@
+"""Handler for database operations related to datasets."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from uuid import UUID
+
+from pydantic import BaseModel, Field, model_validator
+from sqlmodel import Session, and_, col, func, or_, select
+from sqlmodel.sql.expression import SelectOfScalar
+
+from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable
+from lightly_studio.models.dataset import DatasetCreate, DatasetTable
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.tag import TagTable
+
+
+class ExportFilter(BaseModel):
+    """Export Filter to be used for including or excluding."""
+
+    tag_ids: list[UUID] | None = Field(default=None, min_length=1, description="List of tag UUIDs")
+    sample_ids: list[UUID] | None = Field(
+        default=None, min_length=1, description="List of sample UUIDs"
+    )
+    annotation_ids: list[UUID] | None = Field(
+        default=None, min_length=1, description="List of annotation UUIDs"
+    )
+
+    @model_validator(mode="after")
+    def check_exactly_one(self) -> ExportFilter:  # noqa: N804
+        """Ensure that exactly one of the fields is set."""
+        count = (
+            (self.tag_ids is not None)
+            + (self.sample_ids is not None)
+            + (self.annotation_ids is not None)
+        )
+        if count != 1:
+            raise ValueError("Either tag_ids, sample_ids, or annotation_ids must be set.")
+        return self
+
+
+def create(session: Session, dataset: DatasetCreate) -> DatasetTable:
+    """Create a new dataset in the database."""
+    db_dataset = DatasetTable.model_validate(dataset)
+    session.add(db_dataset)
+    session.commit()
+    session.refresh(db_dataset)
+    return db_dataset
+
+
+# TODO(Michal, 06/2025): Use Paginated struct instead of offset and limit
+def get_all(session: Session, offset: int = 0, limit: int = 100) -> list[DatasetTable]:
+    """Retrieve all datasets with pagination."""
+    datasets = session.exec(
+        select(DatasetTable)
+        .order_by(col(DatasetTable.created_at).asc())
+        .offset(offset)
+        .limit(limit)
+    ).all()
+    return list(datasets) if datasets else []
+
+
+def get_by_id(session: Session, dataset_id: UUID) -> DatasetTable | None:
+    """Retrieve a single dataset by ID."""
+    return session.exec(
+        select(DatasetTable).where(DatasetTable.dataset_id == dataset_id)
+    ).one_or_none()
+
+
+def update(session: Session, dataset_id: UUID, dataset_data: DatasetCreate) -> DatasetTable:
+    """Update an existing dataset."""
+    dataset = get_by_id(session=session, dataset_id=dataset_id)
+    if not dataset:
+        raise ValueError(f"Dataset ID was not found '{dataset_id}'.")
+
+    dataset.name = dataset_data.name
+    dataset.directory = dataset_data.directory
+    dataset.updated_at = datetime.now(timezone.utc)
+
+    session.commit()
+    session.refresh(dataset)
+    return dataset
+
+
+def delete(session: Session, dataset_id: UUID) -> bool:
+    """Delete a dataset."""
+    dataset = get_by_id(session=session, dataset_id=dataset_id)
+    if not dataset:
+        return False
+
+    session.delete(dataset)
+    session.commit()
+    return True
+
+
+def _build_export_query(  # noqa: C901
+    dataset_id: UUID,
+    include: ExportFilter | None = None,
+    exclude: ExportFilter | None = None,
+) -> SelectOfScalar[SampleTable]:
+    """Build the export query based on filters.
+
+    Args:
+        session: SQLAlchemy session.
+        dataset_id: UUID of the dataset.
+        include: Filter to include samples.
+        exclude: Filter to exclude samples.
+
+    Returns:
+        SQLModel select query
+    """
+    if not include and not exclude:
+        raise ValueError("Include or exclude filter is required.")
+    if include and exclude:
+        raise ValueError("Cannot include and exclude at the same time.")
+
+    # include tags or sample_ids or annotation_ids from result
+    if include:
+        if include.tag_ids:
+            return (
+                select(SampleTable)
+                .where(SampleTable.dataset_id == dataset_id)
+                .where(
+                    or_(
+                        # Samples with matching sample tags
+                        col(SampleTable.tags).any(
+                            and_(
+                                TagTable.kind == "sample",
+                                col(TagTable.tag_id).in_(include.tag_ids),
+                            )
+                        ),
+                        # Samples with matching annotation tags
+                        col(SampleTable.annotations).any(
+                            col(AnnotationBaseTable.tags).any(
+                                and_(
+                                    TagTable.kind == "annotation",
+                                    col(TagTable.tag_id).in_(include.tag_ids),
+                                )
+                            )
+                        ),
+                    )
+                )
+                .order_by(col(SampleTable.created_at).asc())
+                .distinct()
+            )
+
+        # get samples by specific sample_ids
+        if include.sample_ids:
+            return (
+                select(SampleTable)
+                .where(SampleTable.dataset_id == dataset_id)
+                .where(col(SampleTable.sample_id).in_(include.sample_ids))
+                .order_by(col(SampleTable.created_at).asc())
+                .distinct()
+            )
+
+        # get samples by specific annotation_ids
+        if include.annotation_ids:
+            return (
+                select(SampleTable)
+                .join(SampleTable.annotations)
+                .where(AnnotationBaseTable.dataset_id == dataset_id)
+                .where(col(AnnotationBaseTable.annotation_id).in_(include.annotation_ids))
+                .order_by(col(SampleTable.created_at).asc())
+                .distinct()
+            )
+
+    # exclude tags or sample_ids or annotation_ids from result
+    elif exclude:
+        if exclude.tag_ids:
+            return (
+                select(SampleTable)
+                .where(SampleTable.dataset_id == dataset_id)
+                .where(
+                    and_(
+                        ~col(SampleTable.tags).any(
+                            and_(
+                                TagTable.kind == "sample",
+                                col(TagTable.tag_id).in_(exclude.tag_ids),
+                            )
+                        ),
+                        or_(
+                            ~col(SampleTable.annotations).any(),
+                            ~col(SampleTable.annotations).any(
+                                col(AnnotationBaseTable.tags).any(
+                                    and_(
+                                        TagTable.kind == "annotation",
+                                        col(TagTable.tag_id).in_(exclude.tag_ids),
+                                    )
+                                )
+                            ),
+                        ),
+                    )
+                )
+                .order_by(col(SampleTable.created_at).asc())
+                .distinct()
+            )
+        if exclude.sample_ids:
+            return (
+                select(SampleTable)
+                .where(SampleTable.dataset_id == dataset_id)
+                .where(col(SampleTable.sample_id).notin_(exclude.sample_ids))
+                .order_by(col(SampleTable.created_at).asc())
+                .distinct()
+            )
+        if exclude.annotation_ids:
+            return (
+                select(SampleTable)
+                .where(SampleTable.dataset_id == dataset_id)
+                .where(
+                    or_(
+                        ~col(SampleTable.annotations).any(),
+                        ~col(SampleTable.annotations).any(
+                            col(AnnotationBaseTable.annotation_id).in_(exclude.annotation_ids)
+                        ),
+                    )
+                )
+                .order_by(col(SampleTable.created_at).asc())
+                .distinct()
+            )
+
+    raise ValueError("Invalid include or export filter combination.")
+
+
+# TODO: this fn should be moved to a "business logic" layer outside of the
+# resolvers and abstracted in to reusable components.
+# https://linear.app/lightly/issue/LIG-6196/figure-out-architecture-follow-up-changes-in-python-package
+# TODO: this should be abstracted to allow different export formats.
+def export(
+    session: Session,
+    dataset_id: UUID,
+    include: ExportFilter | None = None,
+    exclude: ExportFilter | None = None,
+) -> list[str]:
+    """Retrieve samples for exporting from a dataset.
+
+    Only one of include or exclude should be set and not both.
+    Furthermore, the include and exclude filter can only have
+    one type (tag_ids, sample_ids or annotations_ids) set.
+
+    Args:
+        session: SQLAlchemy session.
+        dataset_id: UUID of the dataset.
+        include: Filter to include samples.
+        exclude: Filter to exclude samples.
+
+    Returns:
+        List of file paths
+    """
+    query = _build_export_query(dataset_id=dataset_id, include=include, exclude=exclude)
+    result = session.exec(query).all()
+    return [sample.file_path_abs for sample in result]
+
+
+def get_filtered_samples_count(
+    session: Session,
+    dataset_id: UUID,
+    include: ExportFilter | None = None,
+    exclude: ExportFilter | None = None,
+) -> int:
+    """Get statistics about the export query.
+
+    Only one of include or exclude should be set and not both.
+    Furthermore, the include and exclude filter can only have
+    one type (tag_ids, sample_ids or annotations_ids) set.
+
+    Args:
+        session: SQLAlchemy session.
+        dataset_id: UUID of the dataset.
+        include: Filter to include samples.
+        exclude: Filter to exclude samples.
+
+    Returns:
+        Count of files to be exported
+    """
+    query = _build_export_query(dataset_id=dataset_id, include=include, exclude=exclude)
+    count_query = select(func.count()).select_from(query.subquery())
+    return session.exec(count_query).one() or 0

lightly_studio/resolvers/embedding_model_resolver.py

@@ -0,0 +1,100 @@
+"""Handler for database operations related to embedding models."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.embedding_model import (
+    EmbeddingModelCreate,
+    EmbeddingModelTable,
+)
+
+
+def create(session: Session, embedding_model: EmbeddingModelCreate) -> EmbeddingModelTable:
+    """Create a new EmbeddingModel in the database."""
+    db_embedding_model = EmbeddingModelTable.model_validate(embedding_model)
+    session.add(db_embedding_model)
+    session.commit()
+    session.refresh(db_embedding_model)
+    return db_embedding_model
+
+
+def get_all_by_dataset_id(session: Session, dataset_id: UUID) -> list[EmbeddingModelTable]:
+    """Retrieve all embedding models."""
+    embedding_models = session.exec(
+        select(EmbeddingModelTable)
+        .where(EmbeddingModelTable.dataset_id == dataset_id)
+        .order_by(col(EmbeddingModelTable.created_at).asc())
+    ).all()
+    return list(embedding_models)
+
+
+def get_by_id(session: Session, embedding_model_id: UUID) -> EmbeddingModelTable | None:
+    """Retrieve a single embedding model by ID."""
+    return session.exec(
+        select(EmbeddingModelTable).where(
+            EmbeddingModelTable.embedding_model_id == embedding_model_id
+        )
+    ).one_or_none()
+
+
+def get_by_model_hash(session: Session, embedding_model_hash: str) -> EmbeddingModelTable | None:
+    """Retrieve a single embedding model by hash."""
+    return session.exec(
+        select(EmbeddingModelTable).where(
+            EmbeddingModelTable.embedding_model_hash == embedding_model_hash
+        )
+    ).one_or_none()
+
+
+def get_by_name(
+    session: Session, dataset_id: UUID, embedding_model_name: str | None
+) -> EmbeddingModelTable:
+    """Helper function to resolve the embedding model name to its ID.
+
+    Args:
+        session: The database session.
+        dataset_id: The ID of the dataset.
+        embedding_model_name: The name of the embedding model.
+            If None, expects the dataset to have exactly one embedding model and
+            returns it. Otherwise raises a ValueError.
+            If set, expects the dataset to have an embedding model with the given name.
+            Otherwise raises a ValueError.
+
+    Returns:
+        The embedding model with the given name.
+    """
+    embedding_models = get_all_by_dataset_id(
+        session=session,
+        dataset_id=dataset_id,
+    )
+
+    if embedding_model_name is None:
+        if len(embedding_models) != 1:
+            raise ValueError(
+                f"Expected exactly one embedding model, "
+                f"but found {len(embedding_models)} with names "
+                f"{[model.name for model in embedding_models]}."
+            )
+        return embedding_models[0]
+
+    embedding_model_with_name = next(
+        (model for model in embedding_models if model.name == embedding_model_name), None
+    )
+    if embedding_model_with_name is None:
+        raise ValueError(f"Embedding model with name `{embedding_model_name}` not found.")
+
+    return embedding_model_with_name
+
+
+def delete(session: Session, embedding_model_id: UUID) -> bool:
+    """Delete an embedding model."""
+    embedding_model = get_by_id(session=session, embedding_model_id=embedding_model_id)
+    if not embedding_model:
+        return False
+
+    session.delete(embedding_model)
+    session.commit()
+    return True

lightly_studio/resolvers/metadata_resolver/__init__.py

@@ -0,0 +1,15 @@
+"""Metadata resolver module."""
+
+from lightly_studio.resolvers.metadata_resolver.sample import (
+    bulk_set_metadata,
+    get_by_sample_id,
+    get_value_for_sample,
+    set_value_for_sample,
+)
+
+__all__ = [
+    "bulk_set_metadata",
+    "get_by_sample_id",
+    "get_value_for_sample",
+    "set_value_for_sample",
+]

lightly_studio/resolvers/metadata_resolver/metadata_filter.py

@@ -0,0 +1,163 @@
+"""Generic metadata filtering utilities."""
+
+import json
+import re
+from typing import Any, Dict, List, Literal, Protocol, Type, TypeVar
+
+from pydantic import BaseModel
+from sqlalchemy import text
+
+from lightly_studio.type_definitions import QueryType
+
+# Type variables for generic constraints
+T = TypeVar("T", bound=BaseModel)
+M = TypeVar("M", bound="HasMetadata")
+
+# Valid operators for metadata filtering
+MetadataOperator = Literal[">", "<", "==", ">=", "<=", "!="]
+
+# Default metadata column name
+METADATA_COLUMN = "metadata.data"
+
+
+class HasMetadata(Protocol):
+    """Protocol for models that have metadata."""
+
+    data: Dict[str, Any]
+    metadata_schema: Dict[str, str]
+
+
+class MetadataFilter(BaseModel):
+    """Encapsulates a single metadata filter condition."""
+
+    key: str
+    op: MetadataOperator
+    value: Any
+
+    def model_post_init(self, __context: Any) -> None:
+        """Post-initialization hook to serialize string values."""
+        # Pre-serialize string values for JSON comparison
+        if isinstance(self.value, str):
+            # Avoid double-serialization
+            try:
+                json.loads(self.value)
+                # Already serialized, don't serialize again
+            except (json.JSONDecodeError, TypeError):
+                # Not serialized, serialize it
+                self.value = json.dumps(self.value)
+
+
+class Metadata:
+    """Helper class for creating metadata filters with operator syntax."""
+
+    def __init__(self, key: str) -> None:
+        """Initialize metadata filter with key."""
+        self.key = key
+
+    def __gt__(self, value: Any) -> MetadataFilter:
+        """Create greater than filter."""
+        return MetadataFilter(key=self.key, op=">", value=value)
+
+    def __lt__(self, value: Any) -> MetadataFilter:
+        """Create less than filter."""
+        return MetadataFilter(key=self.key, op="<", value=value)
+
+    def __ge__(self, value: Any) -> MetadataFilter:
+        """Create greater than or equal filter."""
+        return MetadataFilter(key=self.key, op=">=", value=value)
+
+    def __le__(self, value: Any) -> MetadataFilter:
+        """Create less than or equal filter."""
+        return MetadataFilter(key=self.key, op="<=", value=value)
+
+    def __eq__(self, value: Any) -> MetadataFilter:  # type: ignore
+        """Create equality filter."""
+        return MetadataFilter(key=self.key, op="==", value=value)
+
+    def __ne__(self, value: Any) -> MetadataFilter:  # type: ignore
+        """Create not equal filter."""
+        return MetadataFilter(key=self.key, op="!=", value=value)
+
+
+def _sanitize_param_name(field: str) -> str:
+    """Sanitize field name for use as SQL parameter name.
+
+    Args:
+        field: The field name (may contain dots for nested paths).
+
+    Returns:
+        A sanitized parameter name safe for SQL binding.
+    """
+    # Replace dots and other problematic characters with underscores
+    return re.sub(r"[^a-zA-Z0-9_]", "_", field)
+
+
+def apply_metadata_filters(
+    query: QueryType,
+    metadata_filters: List[MetadataFilter],
+    *,
+    metadata_model: Type[M],
+    metadata_join_condition: Any,
+) -> QueryType:
+    """Apply metadata filters to a query.
+
+    Args:
+        query: The base query to filter.
+        metadata_filters: The list of metadata filters to apply.
+        metadata_model: The metadata table/model class.
+        metadata_join_condition: The join condition between the main table
+        and metadata table.
+
+    Returns:
+        The filtered query.
+
+    Raises:
+        ValueError: If any field name contains invalid characters.
+
+    Example:
+        ```python
+        # Simple filters (AND by default)
+        query = apply_metadata_filters(
+            query,
+            metadata_filters=[
+                Metadata("temperature") > 25,
+                Metadata("location") == "city",
+            ],
+            metadata_model=SampleMetadataTable,
+            metadata_join_condition=SampleMetadataTable.sample_id ==
+                                    SampleTable.sample_id,
+        )
+        ```
+    """
+    if not metadata_filters:
+        return query
+
+    # Apply the filters using JSON extraction
+    query = query.join(
+        metadata_model,
+        metadata_join_condition,
+    )
+
+    for i, meta_filter in enumerate(metadata_filters):
+        field = meta_filter.key
+        value = meta_filter.value
+        op = meta_filter.op
+
+        json_path = "$." + field
+        # Add unique identifier to parameter name to avoid conflicts
+        param_name = f"{_sanitize_param_name(field)}_{i}"
+
+        # Build the condition based on value type
+        if isinstance(value, (int, float)):
+            # For numeric values, use json_extract with CAST
+            condition = (
+                f"CAST(json_extract({METADATA_COLUMN}, '{json_path}')  AS FLOAT) {op} :{param_name}"
+            )
+        else:
+            # For string values, use json_extract with parameter binding
+            condition = f"json_extract({METADATA_COLUMN}, '{json_path}') {op} :{param_name}"
+
+        # Apply the condition (same for both types)
+        query = query.where(text(condition).bindparams(**{param_name: value}))
+
+    return query

lightly_studio/resolvers/metadata_resolver/sample/__init__.py

@@ -0,0 +1,21 @@
+"""Resolvers for metadata operations."""
+
+from .bulk_set_metadata import (
+    bulk_set_metadata,
+)
+from .get_by_sample_id import (
+    get_by_sample_id,
+)
+from .get_value_for_sample import (
+    get_value_for_sample,
+)
+from .set_value_for_sample import (
+    set_value_for_sample,
+)
+
+__all__ = [
+    "bulk_set_metadata",
+    "get_by_sample_id",
+    "get_value_for_sample",
+    "set_value_for_sample",
+]

lightly_studio/resolvers/metadata_resolver/sample/bulk_set_metadata.py

@@ -0,0 +1,48 @@
+"""Resolver for operations for setting metadata."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Any
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.metadata.complex_metadata import serialize_complex_metadata
+from lightly_studio.models.metadata import (
+    SampleMetadataTable,
+    get_type_name,
+)
+
+
+def bulk_set_metadata(
+    session: Session,
+    sample_metadata: list[tuple[UUID, dict[str, Any]]],
+) -> None:
+    """Bulk insert metadata for multiple samples.
+
+    Args:
+        session: The database session.
+        sample_metadata: List of (sample_id, metadata_dict) tuples.
+    """
+    now = datetime.now(timezone.utc)
+    objects = []
+    for sample_id, metadata in sample_metadata:
+        metadata_schema = {}
+        serialized_data = {}
+        for key, value in metadata.items():
+            # Serialize complex metadata objects.
+            serialized_data[key] = serialize_complex_metadata(value)
+            # Get the type name
+            metadata_schema[key] = get_type_name(value)
+
+        obj = SampleMetadataTable(
+            sample_id=sample_id,
+            data=serialized_data,
+            metadata_schema=metadata_schema,
+            created_at=now,
+            updated_at=now,
+        )
+        objects.append(obj)
+    session.bulk_save_objects(objects)
+    session.commit()

lightly_studio/resolvers/metadata_resolver/sample/get_by_sample_id.py

@@ -0,0 +1,24 @@
+"""Resolver for operations for retrieving metadata."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, select
+
+from lightly_studio.models.metadata import SampleMetadataTable
+
+
+def get_by_sample_id(session: Session, sample_id: UUID) -> SampleMetadataTable | None:
+    """Retrieve the metadata object for a given sample.
+
+    Args:
+        session: The database session.
+        sample_id: The sample's UUID.
+
+    Returns:
+        The CustomMetadataTable instance or None if not found.
+    """
+    return session.exec(
+        select(SampleMetadataTable).where(SampleMetadataTable.sample_id == sample_id)
+    ).one_or_none()