lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/resolvers/dataset_resolver/get_or_create_child_dataset.py

@@ -0,0 +1,58 @@
+"""Function to get or create a unique child dataset with a given sample type."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.models.dataset import DatasetCreate, SampleType
+from lightly_studio.resolvers import dataset_resolver
+
+
+def get_or_create_child_dataset(
+    session: Session, dataset_id: UUID, sample_type: SampleType
+) -> UUID:
+    """Checks if a unique child dataset with the given sample type exists for the given dataset.
+
+    If it exists, returns its ID. If not, creates it and then returns its ID.
+    If multiple such datasets exist, raises an error.
+
+    The returned child is a direct child of the given dataset.
+
+    Args:
+        session: The database session.
+        dataset_id: The uuid of the dataset to attach to.
+        sample_type: The sample type of the child dataset to get or create.
+
+    Returns:
+        The uuid of the child dataset.
+
+    Raises:
+        ValueError: If multiple child datasets with the given sample type exist.
+    """
+    # Get filtered child datasets.
+    dataset = dataset_resolver.get_by_id(session=session, dataset_id=dataset_id)
+    if dataset is None:
+        raise ValueError(f"Dataset with id {dataset_id} not found.")
+    child_datasets = [ds for ds in dataset.children if ds.sample_type == sample_type]
+
+    # If we have children check if any have the given sample type.
+    if len(child_datasets) == 1:
+        return child_datasets[0].dataset_id
+    if len(child_datasets) > 1:
+        raise ValueError(
+            f"Multiple child datasets with sample type {sample_type.value} found "
+            f"for dataset id {dataset_id}."
+        )
+
+    # No child dataset with the given sample type found, create one.
+    child_dataset = dataset_resolver.create(
+        session=session,
+        dataset=DatasetCreate(
+            name=f"{dataset.name}__{sample_type.value.lower()}",
+            sample_type=sample_type,
+            parent_dataset_id=dataset_id,
+        ),
+    )
+    return child_dataset.dataset_id

lightly_studio/resolvers/dataset_resolver/get_parent_dataset_by_sample_id.py

@@ -0,0 +1,27 @@
+"""Retrieve the parent dataset for a given sample ID."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.dataset import DatasetTable
+from lightly_studio.models.sample import SampleTable
+
+
+def get_parent_dataset_by_sample_id(session: Session, sample_id: UUID) -> DatasetTable | None:
+    """Get parent dataset by sample ID.
+
+    Args:
+        session: Database session
+        sample_id: ID of the sample for which to get the parent dataset
+
+    Returns:
+        Returns parent dataset
+    """
+    child = session.exec(
+        select(DatasetTable).join(SampleTable).where(col(SampleTable.sample_id) == sample_id)
+    ).one_or_none()
+
+    return child.parent if child else None

lightly_studio/resolvers/dataset_resolver/get_parent_dataset_id.py

@@ -0,0 +1,22 @@
+"""Retrieve the parent dataset ID for a given dataset ID."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlalchemy.orm import aliased
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.dataset import DatasetTable
+
+ParentDataset = aliased(DatasetTable)
+ChildDataset = aliased(DatasetTable)
+
+
+def get_parent_dataset_id(session: Session, dataset_id: UUID) -> DatasetTable | None:
+    """Retrieve the parent dataset for a given dataset ID."""
+    return session.exec(
+        select(ParentDataset)
+        .join(ChildDataset, col(ChildDataset.parent_dataset_id) == col(ParentDataset.dataset_id))
+        .where(ChildDataset.dataset_id == dataset_id)
+    ).one_or_none()

lightly_studio/resolvers/dataset_resolver/get_root_dataset.py

@@ -0,0 +1,61 @@
+"""Implementation of get_root_dataset resolver function."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.dataset import DatasetTable
+
+
+# TODO (Mihnea, 12/2025): Update the dataset_id to be required.
+#  The dataset_id is currently optional for backwards compatibility.
+def get_root_dataset(session: Session, dataset_id: UUID | None = None) -> DatasetTable:
+    """Retrieve the root dataset for a given dataset or the first root dataset.
+
+    If dataset_id is provided, traverses up the hierarchy to find the root ancestor.
+    If dataset_id is None, returns the first root dataset (backwards compatibility).
+
+    A root dataset is defined as a dataset where parent_dataset_id is None.
+    The root dataset may or may not have children.
+
+    Args:
+        session: The database session.
+        dataset_id: Optional ID of a dataset to find the root for.
+
+    Returns:
+        The root dataset.
+
+    Raises:
+        ValueError: If no root dataset is found or dataset_id doesn't exist.
+    """
+    if dataset_id is not None:
+        # Find the dataset.
+        dataset = session.get(DatasetTable, dataset_id)
+        if dataset is None:
+            raise ValueError(f"Dataset with ID {dataset_id} not found.")
+
+        # Traverse up the hierarchy until we find the root.
+        # TODO (Mihnea, 12/2025): Consider replacing the loop with a recursive CTE,
+        #  if this becomes a bottleneck.
+        while dataset.parent_dataset_id is not None:
+            parent = session.get(DatasetTable, dataset.parent_dataset_id)
+            if parent is None:
+                raise ValueError(
+                    f"Parent dataset {dataset.parent_dataset_id} not found "
+                    f"for dataset {dataset.dataset_id}."
+                )
+            dataset = parent
+
+        return dataset
+
+    # Backwards compatibility: return first root dataset
+    root_datasets = session.exec(
+        select(DatasetTable).where(col(DatasetTable.parent_dataset_id).is_(None))
+    ).all()
+
+    if len(root_datasets) == 0:
+        raise ValueError("No root dataset found. A root dataset must exist.")
+
+    return root_datasets[0]

lightly_studio/resolvers/dataset_resolver/get_root_datasets_overview.py

@@ -0,0 +1,41 @@
+"""Handler for database operations related to fetching root datasets with details."""
+
+from __future__ import annotations
+
+from sqlmodel import Session, col, func, select
+
+from lightly_studio.models.dataset import DatasetOverviewView, DatasetTable
+from lightly_studio.models.sample import SampleTable
+
+
+def get_root_datasets_overview(session: Session) -> list[DatasetOverviewView]:
+    """Get root datasets with detailed metadata including sample counts."""
+    datasets_query = (
+        select(  # type: ignore[call-overload]
+            DatasetTable.dataset_id,
+            DatasetTable.name,
+            DatasetTable.sample_type,
+            DatasetTable.created_at,
+            func.count(col(SampleTable.dataset_id)).label("sample_count"),
+        )
+        .outerjoin(SampleTable)
+        .where(col(DatasetTable.parent_dataset_id).is_(None))
+        .group_by(
+            DatasetTable.dataset_id,
+            DatasetTable.name,
+            DatasetTable.sample_type,
+            DatasetTable.created_at,
+        )
+        .order_by(col(DatasetTable.created_at).desc())
+    )
+
+    return [
+        DatasetOverviewView(
+            dataset_id=row.dataset_id,
+            name=row.name,
+            sample_type=row.sample_type,
+            created_at=row.created_at,
+            total_sample_count=row.sample_count,
+        )
+        for row in session.exec(datasets_query).all()
+    ]

lightly_studio/resolvers/dataset_resolver/update.py

@@ -0,0 +1,25 @@
+"""Implementation of update dataset resolver function."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.models.dataset import DatasetCreate, DatasetTable
+from lightly_studio.resolvers.dataset_resolver.get_by_id import get_by_id
+
+
+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.updated_at = datetime.now(timezone.utc)
+
+    session.commit()
+    session.refresh(dataset)
+    return dataset

lightly_studio/resolvers/embedding_model_resolver.py

@@ -0,0 +1,120 @@
+"""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_or_create(session: Session, embedding_model: EmbeddingModelCreate) -> EmbeddingModelTable:
+    """Retrieve an existing EmbeddingModel by hash or create a new one if it does not exist."""
+    db_model = get_by_model_hash(
+        session=session, embedding_model_hash=embedding_model.embedding_model_hash
+    )
+    if db_model is None:
+        return create(session=session, embedding_model=embedding_model)
+
+    # Validate that the existing model matches the provided data.
+    if (
+        db_model.name != embedding_model.name
+        or db_model.parameter_count_in_mb != embedding_model.parameter_count_in_mb
+        or db_model.embedding_dimension != embedding_model.embedding_dimension
+    ):
+        raise ValueError(
+            "An embedding model with the same hash but different parameters already exists."
+        )
+    return db_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/image_filter.py

@@ -0,0 +1,50 @@
+"""Utility functions for building database queries."""
+# TODO(Michal, 11/2025): Move to image_resolver once DatasetTable.get_samples() is removed.
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+from lightly_studio.models.image import ImageTable
+from lightly_studio.resolvers.sample_resolver.sample_filter import SampleFilter
+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 ImageFilter(BaseModel):
+    """Encapsulates filter parameters for querying samples."""
+
+    sample_filter: Optional[SampleFilter] = None
+    width: Optional[FilterDimensions] = None
+    height: Optional[FilterDimensions] = None
+
+    def apply(self, query: QueryType) -> QueryType:
+        """Apply the filters to the given query."""
+        # Apply sample filters to the query.
+        if self.sample_filter is not None:
+            query = self.sample_filter.apply(query)
+
+        # Apply dimension-based filters to the query.
+        query = self._apply_dimension_filters(query)
+
+        # Return the modified query.
+        return query  # noqa: RET504
+
+    def _apply_dimension_filters(self, query: QueryType) -> QueryType:
+        if self.width:
+            if self.width.min is not None:
+                query = query.where(ImageTable.width >= self.width.min)
+            if self.width.max is not None:
+                query = query.where(ImageTable.width <= self.width.max)
+        if self.height:
+            if self.height.min is not None:
+                query = query.where(ImageTable.height >= self.height.min)
+            if self.height.max is not None:
+                query = query.where(ImageTable.height <= self.height.max)
+        return query

lightly_studio/resolvers/image_resolver/__init__.py

@@ -0,0 +1,21 @@
+"""Resolvers for database operations."""
+
+from lightly_studio.resolvers.image_resolver.create_many import create_many
+from lightly_studio.resolvers.image_resolver.delete import delete
+from lightly_studio.resolvers.image_resolver.filter_new_paths import filter_new_paths
+from lightly_studio.resolvers.image_resolver.get_all_by_dataset_id import get_all_by_dataset_id
+from lightly_studio.resolvers.image_resolver.get_by_id import get_by_id
+from lightly_studio.resolvers.image_resolver.get_dimension_bounds import get_dimension_bounds
+from lightly_studio.resolvers.image_resolver.get_many_by_id import get_many_by_id
+from lightly_studio.resolvers.image_resolver.get_samples_excluding import get_samples_excluding
+
+__all__ = [
+    "create_many",
+    "delete",
+    "filter_new_paths",
+    "get_all_by_dataset_id",
+    "get_by_id",
+    "get_dimension_bounds",
+    "get_many_by_id",
+    "get_samples_excluding",
+]

lightly_studio/resolvers/image_resolver/create_many.py

@@ -0,0 +1,52 @@
+"""Implementation of create functions for images."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.models.dataset import SampleType
+from lightly_studio.models.image import ImageCreate, ImageTable
+from lightly_studio.models.sample import SampleCreate
+from lightly_studio.resolvers import dataset_resolver, sample_resolver
+
+
+class ImageCreateHelper(ImageCreate):
+    """Helper class to create ImageTable with sample_id."""
+
+    sample_id: UUID
+    dataset_id: UUID
+
+
+def create_many(session: Session, dataset_id: UUID, samples: list[ImageCreate]) -> list[UUID]:
+    """Create multiple samples in a single database commit.
+
+    Returns the list of created sample IDs that matches the order of input samples.
+    """
+    dataset_resolver.check_dataset_type(
+        session=session,
+        dataset_id=dataset_id,
+        expected_type=SampleType.IMAGE,
+    )
+    sample_ids = sample_resolver.create_many(
+        session=session,
+        samples=[SampleCreate(dataset_id=dataset_id) for _ in samples],
+    )
+    # Bulk create ImageTable entries using the generated sample_ids.
+    db_images = [
+        ImageTable.model_validate(
+            ImageCreateHelper(
+                file_name=sample.file_name,
+                width=sample.width,
+                height=sample.height,
+                dataset_id=dataset_id,
+                file_path_abs=sample.file_path_abs,
+                sample_id=sample_id,
+            )
+        )
+        for sample_id, sample in zip(sample_ids, samples)
+    ]
+    session.bulk_save_objects(db_images)
+    session.commit()
+    return sample_ids

lightly_studio/resolvers/image_resolver/delete.py

@@ -0,0 +1,20 @@
+"""Implementation of delete function for images."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.resolvers import image_resolver
+
+
+def delete(session: Session, sample_id: UUID) -> bool:
+    """Delete a sample."""
+    sample = image_resolver.get_by_id(session=session, sample_id=sample_id)
+    if not sample:
+        return False
+
+    session.delete(sample)
+    session.commit()
+    return True

lightly_studio/resolvers/image_resolver/filter_new_paths.py

@@ -0,0 +1,23 @@
+"""Implementation of filter_new_paths function for images."""
+
+from __future__ import annotations
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.image import ImageTable
+
+
+def filter_new_paths(session: Session, file_paths_abs: list[str]) -> tuple[list[str], list[str]]:
+    """Return a) file_path_abs that do not already exist in the database and b) those that do."""
+    existing_file_paths_abs = set(
+        session.exec(
+            select(col(ImageTable.file_path_abs)).where(
+                col(ImageTable.file_path_abs).in_(file_paths_abs)
+            )
+        ).all()
+    )
+    file_paths_abs_set = set(file_paths_abs)
+    return (
+        list(file_paths_abs_set - existing_file_paths_abs),  # paths that are not in the DB
+        list(file_paths_abs_set & existing_file_paths_abs),  # paths that are already in the DB
+    )

lightly_studio/resolvers/image_resolver/get_all_by_dataset_id.py

@@ -0,0 +1,117 @@
+"""Implementation of get_all_by_dataset_id function for images."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from uuid import UUID
+
+from pydantic import BaseModel
+from sqlalchemy.orm import joinedload, selectinload
+from sqlmodel import Session, col, func, select
+
+from lightly_studio.api.routes.api.validators import Paginated
+from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable
+from lightly_studio.models.embedding_model import EmbeddingModelTable
+from lightly_studio.models.image import ImageTable
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.sample_embedding import SampleEmbeddingTable
+from lightly_studio.resolvers.image_filter import ImageFilter
+
+
+class GetAllSamplesByDatasetIdResult(BaseModel):
+    """Result of getting all samples."""
+
+    samples: Sequence[ImageTable]
+    total_count: int
+    next_cursor: int | None = None
+
+
+def get_all_by_dataset_id(  # noqa: PLR0913
+    session: Session,
+    dataset_id: UUID,
+    pagination: Paginated | None = None,
+    filters: ImageFilter | 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(ImageTable)
+        .options(
+            selectinload(ImageTable.sample).options(
+                joinedload(SampleTable.tags),
+                # Ignore type checker error below as it's a false positive caused by TYPE_CHECKING.
+                joinedload(SampleTable.metadata_dict),  # type: ignore[arg-type]
+                selectinload(SampleTable.captions),
+                selectinload(SampleTable.annotations).options(
+                    joinedload(AnnotationBaseTable.annotation_label),
+                    joinedload(AnnotationBaseTable.object_detection_details),
+                    joinedload(AnnotationBaseTable.instance_segmentation_details),
+                    joinedload(AnnotationBaseTable.semantic_segmentation_details),
+                    selectinload(AnnotationBaseTable.tags),
+                ),
+            ),
+        )
+        .join(ImageTable.sample)
+        .where(SampleTable.dataset_id == dataset_id)
+    )
+    total_count_query = (
+        select(func.count())
+        .select_from(ImageTable)
+        .join(ImageTable.sample)
+        .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(ImageTable.sample_id).in_(sample_ids))
+        total_count_query = total_count_query.where(col(ImageTable.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(ImageTable.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(ImageTable.sample_id) == col(SampleEmbeddingTable.sample_id),
+            ).where(SampleEmbeddingTable.embedding_model_id == embedding_model_id)
+    else:
+        samples_query = samples_query.order_by(col(ImageTable.file_path_abs).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 and pagination.offset + pagination.limit < total_count:
+        next_cursor = pagination.offset + pagination.limit
+
+    return GetAllSamplesByDatasetIdResult(
+        samples=session.exec(samples_query).all(),
+        total_count=total_count,
+        next_cursor=next_cursor,
+    )

lightly_studio/resolvers/image_resolver/get_by_id.py

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