lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/resolvers/image_resolver/get_dimension_bounds.py

@@ -0,0 +1,75 @@
+"""Implementation of get_dimension_bounds function for images."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+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.image import ImageTable
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.tag import TagTable
+
+
+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(ImageTable.width).label("min_width"),
+        func.max(ImageTable.width).label("max_width"),
+        func.min(ImageTable.height).label("min_height"),
+        func.max(ImageTable.height).label("max_height"),
+    )
+    query = query.join(ImageTable.sample)
+
+    if annotation_label_ids:
+        # Subquery to filter samples matching all annotation labels
+        label_filter = (
+            select(ImageTable.sample_id)
+            .join(ImageTable.sample)
+            .join(
+                AnnotationBaseTable,
+                col(ImageTable.sample_id) == col(AnnotationBaseTable.parent_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(ImageTable.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(ImageTable.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}

lightly_studio/resolvers/image_resolver/get_many_by_id.py

@@ -0,0 +1,22 @@
+"""Implementation of get_many_by_id function for images."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.image import ImageTable
+
+
+def get_many_by_id(session: Session, sample_ids: list[UUID]) -> list[ImageTable]:
+    """Retrieve multiple samples by their IDs.
+
+    Output order matches the input order.
+    """
+    results = session.exec(
+        select(ImageTable).where(col(ImageTable.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/image_resolver/get_samples_excluding.py

@@ -0,0 +1,43 @@
+"""Implementation of get_samples_excluding function for images."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from uuid import UUID
+
+from sqlmodel import Session, col, func, select
+
+from lightly_studio.models.image import ImageTable
+from lightly_studio.models.sample import SampleTable
+
+
+def get_samples_excluding(
+    session: Session,
+    dataset_id: UUID,
+    excluded_sample_ids: list[UUID],
+    limit: int | None = None,
+) -> Sequence[ImageTable]:
+    """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(ImageTable)
+        .join(ImageTable.sample)
+        .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/metadata_resolver/__init__.py

@@ -0,0 +1,15 @@
+"""Metadata resolver module."""
+
+from lightly_studio.resolvers.metadata_resolver.sample import (
+    bulk_update_metadata,
+    get_by_sample_id,
+    get_value_for_sample,
+    set_value_for_sample,
+)
+
+__all__ = [
+    "bulk_update_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 ==
+                                    ImageTable.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_update_metadata import (
+    bulk_update_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_update_metadata",
+    "get_by_sample_id",
+    "get_value_for_sample",
+    "set_value_for_sample",
+]

lightly_studio/resolvers/metadata_resolver/sample/bulk_update_metadata.py

@@ -0,0 +1,46 @@
+"""Resolver for operations for setting metadata."""
+
+from __future__ import annotations
+
+from typing import Any
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.metadata import SampleMetadataTable
+
+
+def bulk_update_metadata(
+    session: Session,
+    sample_metadata: list[tuple[UUID, dict[str, Any]]],
+) -> None:
+    """Bulk insert or update metadata for multiple samples.
+
+    If a sample does not have metadata, a new metadata row is created.
+    If a sample already has metadata, the new key-value pairs are merged with the existing metadata.
+
+    Args:
+        session: The database session.
+        sample_metadata: List of (sample_id, metadata_dict) tuples.
+    """
+    # TODO(Mihnea, 10/2025): Consider using SQLAlchemy's bulk operations
+    #  (Session.bulk_insert/update_mappings) if performance becomes a bottleneck.
+    if not sample_metadata:
+        return
+
+    # Get all existing metadata rows for the given sample IDs.
+    sample_ids = [s[0] for s in sample_metadata]
+    existing_metadata = session.exec(
+        select(SampleMetadataTable).where(col(SampleMetadataTable.sample_id).in_(sample_ids))
+    ).all()
+    sample_id_to_existing_metadata = {meta.sample_id: meta for meta in existing_metadata}
+
+    for sample_id, new_metadata in sample_metadata:
+        metadata = sample_id_to_existing_metadata.get(
+            sample_id, SampleMetadataTable(sample_id=sample_id)
+        )
+        for key, value in new_metadata.items():
+            metadata.set_value(key, value)
+        session.add(metadata)
+
+    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()

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