lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/resolvers/caption_resolver.py

@@ -0,0 +1,129 @@
+"""Resolvers for caption."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.caption import CaptionCreate, CaptionTable
+from lightly_studio.models.dataset import SampleType
+from lightly_studio.models.sample import SampleCreate
+from lightly_studio.resolvers import dataset_resolver, sample_resolver
+
+
+class CaptionCreateHelper(CaptionCreate):
+    """Helper class to create CaptionTable with sample_id."""
+
+    sample_id: UUID
+
+
+def create_many(
+    session: Session, parent_dataset_id: UUID, captions: Sequence[CaptionCreate]
+) -> list[UUID]:
+    """Create captions for a single dataset in bulk.
+
+    It is responsibility of the caller to ensure that all parent samples belong to the same
+    dataset with ID `parent_dataset_id`. This function does not perform this check for performance
+    reasons.
+
+    Args:
+        session: Database session
+        parent_dataset_id: UUID of the parent dataset of which the caption dataset is a child
+        captions: The captions to create
+
+    Returns:
+        List of created CaptionTable sample_ids
+    """
+    if not captions:
+        return []
+
+    caption_dataset_id = dataset_resolver.get_or_create_child_dataset(
+        session=session, dataset_id=parent_dataset_id, sample_type=SampleType.CAPTION
+    )
+    sample_ids = sample_resolver.create_many(
+        session=session,
+        samples=[SampleCreate(dataset_id=caption_dataset_id) for _ in captions],
+    )
+
+    # Bulk create CaptionTable entries using the generated sample_ids.
+    db_captions = [
+        CaptionTable.model_validate(
+            CaptionCreateHelper(
+                parent_sample_id=sample.parent_sample_id,
+                text=sample.text,
+                sample_id=sample_id,
+            )
+        )
+        for sample_id, sample in zip(sample_ids, captions)
+    ]
+    session.bulk_save_objects(db_captions)
+    session.commit()
+    return sample_ids
+
+
+def get_by_ids(session: Session, sample_ids: Sequence[UUID]) -> list[CaptionTable]:
+    """Retrieve captions by IDs."""
+    results = session.exec(
+        select(CaptionTable).where(col(CaptionTable.sample_id).in_(set(sample_ids)))
+    ).all()
+    # Return samples in the same order as the input IDs
+    caption_map = {caption.sample_id: caption for caption in results}
+    return [caption_map[id_] for id_ in sample_ids if id_ in caption_map]
+
+
+def update_text(
+    session: Session,
+    sample_id: UUID,
+    text: str,
+) -> CaptionTable:
+    """Update the text of a caption.
+
+    Args:
+        session: Database session for executing the operation.
+        sample_id: UUID of the caption to update.
+        text: New text.
+
+    Returns:
+        The updated caption with the new text.
+
+    Raises:
+        ValueError: If the caption is not found.
+    """
+    captions = get_by_ids(session, [sample_id])
+    if not captions:
+        raise ValueError(f"Caption with ID {sample_id} not found.")
+
+    caption = captions[0]
+    try:
+        caption.text = text
+        session.commit()
+        session.refresh(caption)
+        return caption
+    except Exception:
+        session.rollback()
+        raise
+
+
+def delete_caption(
+    session: Session,
+    sample_id: UUID,
+) -> None:
+    """Delete a caption.
+
+    Args:
+        session: Database session for executing the operation.
+        sample_id: UUID of the caption to update.
+
+    Raises:
+        ValueError: If the caption is not found.
+    """
+    captions = get_by_ids(session=session, sample_ids=[sample_id])
+    if len(captions) == 0:
+        raise ValueError(f"Caption with ID {sample_id} not found.")
+
+    caption = captions[0]
+    session.commit()
+    session.delete(caption)
+    session.commit()

lightly_studio/resolvers/dataset_resolver/__init__.py

@@ -0,0 +1,55 @@
+"""Resolvers for database operations."""
+
+from lightly_studio.resolvers.dataset_resolver.check_dataset_type import (
+    check_dataset_type,
+)
+from lightly_studio.resolvers.dataset_resolver.create import create
+from lightly_studio.resolvers.dataset_resolver.delete import delete
+from lightly_studio.resolvers.dataset_resolver.export import (
+    export,
+    get_filtered_samples_count,
+)
+from lightly_studio.resolvers.dataset_resolver.get_all import get_all
+from lightly_studio.resolvers.dataset_resolver.get_by_id import get_by_id
+from lightly_studio.resolvers.dataset_resolver.get_by_name import get_by_name
+from lightly_studio.resolvers.dataset_resolver.get_dataset_details import (
+    get_dataset_details,
+)
+from lightly_studio.resolvers.dataset_resolver.get_hierarchy import (
+    get_hierarchy,
+)
+from lightly_studio.resolvers.dataset_resolver.get_or_create_child_dataset import (
+    get_or_create_child_dataset,
+)
+from lightly_studio.resolvers.dataset_resolver.get_parent_dataset_by_sample_id import (
+    get_parent_dataset_by_sample_id,
+)
+from lightly_studio.resolvers.dataset_resolver.get_parent_dataset_id import (
+    get_parent_dataset_id,
+)
+from lightly_studio.resolvers.dataset_resolver.get_root_dataset import (
+    get_root_dataset,
+)
+from lightly_studio.resolvers.dataset_resolver.get_root_datasets_overview import (
+    get_root_datasets_overview,
+)
+from lightly_studio.resolvers.dataset_resolver.update import update
+
+__all__ = [
+    "check_dataset_type",
+    "create",
+    "delete",
+    "export",
+    "get_all",
+    "get_by_id",
+    "get_by_name",
+    "get_dataset_details",
+    "get_filtered_samples_count",
+    "get_hierarchy",
+    "get_or_create_child_dataset",
+    "get_parent_dataset_by_sample_id",
+    "get_parent_dataset_id",
+    "get_root_dataset",
+    "get_root_datasets_overview",
+    "update",
+]

lightly_studio/resolvers/dataset_resolver/check_dataset_type.py

@@ -0,0 +1,29 @@
+"""Implementation of check_dataset_type resolver function."""
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.models.dataset import SampleType
+from lightly_studio.resolvers import dataset_resolver
+
+
+def check_dataset_type(session: Session, dataset_id: UUID, expected_type: SampleType) -> None:
+    """Check that the dataset has the expected sample type.
+
+    Raises a ValueError if the dataset does not have the expected sample type or
+    if it does not exist.
+
+    Args:
+        session: The database session.
+        dataset_id: The ID of the dataset to check.
+        expected_type: The expected sample type.
+    """
+    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.")
+    if dataset.sample_type != expected_type:
+        raise ValueError(
+            f"Dataset with id {dataset_id} is having sample type "
+            f"'{dataset.sample_type.value}', expected '{expected_type.value}'."
+        )

lightly_studio/resolvers/dataset_resolver/create.py

@@ -0,0 +1,20 @@
+"""Implementation of create dataset resolver function."""
+
+from __future__ import annotations
+
+from sqlmodel import Session
+
+from lightly_studio.models.dataset import DatasetCreate, DatasetTable
+from lightly_studio.resolvers import dataset_resolver
+
+
+def create(session: Session, dataset: DatasetCreate) -> DatasetTable:
+    """Create a new dataset in the database."""
+    existing = dataset_resolver.get_by_name(session=session, name=dataset.name)
+    if existing:
+        raise ValueError(f"Dataset with name '{dataset.name}' already exists.")
+    db_dataset = DatasetTable.model_validate(dataset)
+    session.add(db_dataset)
+    session.commit()
+    session.refresh(db_dataset)
+    return db_dataset

lightly_studio/resolvers/dataset_resolver/delete.py

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

lightly_studio/resolvers/dataset_resolver/export.py

@@ -0,0 +1,267 @@
+"""Resolver functions for exporting dataset samples based on filters."""
+
+from __future__ import annotations
+
+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 SampleType
+from lightly_studio.models.image import ImageTable
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.tag import TagTable
+from lightly_studio.resolvers.dataset_resolver.get_hierarchy import get_hierarchy
+
+
+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
+
+
+# TODO(Michal, 10/2025): Consider moving the export logic to a separate service.
+# This is a legacy code from the initial implementation of the export feature.
+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
+    """
+    # Get all child dataset IDs that could contain annotations
+    annotation_dataset_ids = _get_annotation_dataset_ids(session, dataset_id)
+    query = _build_export_query(
+        dataset_id=dataset_id,
+        annotation_dataset_ids=annotation_dataset_ids,
+        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
+    """
+    # Get all child dataset IDs that could contain annotations
+    annotation_dataset_ids = _get_annotation_dataset_ids(session, dataset_id)
+    query = _build_export_query(
+        dataset_id=dataset_id,
+        annotation_dataset_ids=annotation_dataset_ids,
+        include=include,
+        exclude=exclude,
+    )
+    count_query = select(func.count()).select_from(query.subquery())
+    return session.exec(count_query).one() or 0
+
+
+def _get_annotation_dataset_ids(session: Session, dataset_id: UUID) -> list[UUID]:
+    """Get all child dataset IDs that could contain annotations.
+
+    This includes the dataset itself and all its child datasets (recursively)
+    that have sample_type ANNOTATION.
+
+    Args:
+        session: SQLAlchemy session.
+        dataset_id: UUID of the root dataset.
+
+    Returns:
+        List of dataset IDs that could contain annotations.
+    """
+    hierarchy = get_hierarchy(session, dataset_id)
+    return [ds.dataset_id for ds in hierarchy if ds.sample_type == SampleType.ANNOTATION]
+
+
+def _build_export_query(  # noqa: C901
+    dataset_id: UUID,
+    annotation_dataset_ids: list[UUID],
+    include: ExportFilter | None = None,
+    exclude: ExportFilter | None = None,
+) -> SelectOfScalar[ImageTable]:
+    """Build the export query based on filters.
+
+    Args:
+        dataset_id: UUID of the dataset.
+        annotation_dataset_ids: List of dataset IDs that could contain annotations.
+        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(ImageTable)
+                .join(ImageTable.sample)
+                .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(ImageTable.created_at).asc())
+                .distinct()
+            )
+
+        # get samples by specific sample_ids
+        if include.sample_ids:
+            return (
+                select(ImageTable)
+                .join(ImageTable.sample)
+                .where(SampleTable.dataset_id == dataset_id)
+                .where(col(ImageTable.sample_id).in_(include.sample_ids))
+                .order_by(col(ImageTable.created_at).asc())
+                .distinct()
+            )
+
+        # get samples by specific annotation_ids
+        if include.annotation_ids:
+            # Annotations are stored in child datasets, so filter by all annotation dataset IDs
+            # Filter by checking if the annotation's sample_id belongs to a sample in
+            # annotation_dataset_ids
+            annotation_sample_subquery = select(SampleTable.sample_id).where(
+                col(SampleTable.dataset_id).in_(annotation_dataset_ids)
+            )
+            return (
+                select(ImageTable)
+                .join(ImageTable.sample)
+                .join(SampleTable.annotations)
+                .where(col(AnnotationBaseTable.sample_id).in_(annotation_sample_subquery))
+                .where(col(AnnotationBaseTable.sample_id).in_(include.annotation_ids))
+                .order_by(col(ImageTable.created_at).asc())
+                .distinct()
+            )
+
+    # exclude tags or sample_ids or annotation_ids from result
+    elif exclude:
+        if exclude.tag_ids:
+            return (
+                select(ImageTable)
+                .join(ImageTable.sample)
+                .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(ImageTable.created_at).asc())
+                .distinct()
+            )
+        if exclude.sample_ids:
+            return (
+                select(ImageTable)
+                .join(ImageTable.sample)
+                .where(SampleTable.dataset_id == dataset_id)
+                .where(col(ImageTable.sample_id).notin_(exclude.sample_ids))
+                .order_by(col(ImageTable.created_at).asc())
+                .distinct()
+            )
+        if exclude.annotation_ids:
+            return (
+                select(ImageTable)
+                .join(ImageTable.sample)
+                .where(SampleTable.dataset_id == dataset_id)
+                .where(
+                    or_(
+                        ~col(SampleTable.annotations).any(),
+                        ~col(SampleTable.annotations).any(
+                            col(AnnotationBaseTable.sample_id).in_(exclude.annotation_ids)
+                        ),
+                    )
+                )
+                .order_by(col(ImageTable.created_at).asc())
+                .distinct()
+            )
+
+    raise ValueError("Invalid include or export filter combination.")

lightly_studio/resolvers/dataset_resolver/get_all.py

@@ -0,0 +1,19 @@
+"""Implementation of get all datasets resolver function."""
+
+from __future__ import annotations
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.dataset import DatasetTable
+
+
+# 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 []

lightly_studio/resolvers/dataset_resolver/get_by_id.py

@@ -0,0 +1,16 @@
+"""Implementation of get dataset by ID resolver function."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, select
+
+from lightly_studio.models.dataset import DatasetTable
+
+
+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()

lightly_studio/resolvers/dataset_resolver/get_by_name.py

@@ -0,0 +1,12 @@
+"""Implementation of get dataset by name resolver function."""
+
+from __future__ import annotations
+
+from sqlmodel import Session, select
+
+from lightly_studio.models.dataset import DatasetTable
+
+
+def get_by_name(session: Session, name: str) -> DatasetTable | None:
+    """Retrieve a single dataset by name."""
+    return session.exec(select(DatasetTable).where(DatasetTable.name == name)).one_or_none()

lightly_studio/resolvers/dataset_resolver/get_dataset_details.py

@@ -0,0 +1,27 @@
+"""Handler for database operations related to datasets."""
+
+from __future__ import annotations
+
+from sqlmodel import Session, func, select
+
+from lightly_studio.models.dataset import DatasetTable, DatasetViewWithCount
+from lightly_studio.models.sample import SampleTable
+
+
+def get_dataset_details(session: Session, dataset: DatasetTable) -> DatasetViewWithCount:
+    """Convert a DatasetTable to DatasetViewWithCount with computed sample count."""
+    sample_count = (
+        session.exec(
+            select(func.count("*")).where(SampleTable.dataset_id == dataset.dataset_id)
+        ).one()
+        or 0
+    )
+    return DatasetViewWithCount(
+        dataset_id=dataset.dataset_id,
+        parent_dataset_id=dataset.parent_dataset_id,
+        sample_type=dataset.sample_type,
+        name=dataset.name,
+        created_at=dataset.created_at,
+        updated_at=dataset.updated_at,
+        total_sample_count=sample_count,
+    )

lightly_studio/resolvers/dataset_resolver/get_hierarchy.py

@@ -0,0 +1,31 @@
+"""Implementation of get_child_datasets resolver function."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.models.dataset import DatasetTable
+from lightly_studio.resolvers import dataset_resolver
+
+
+def get_hierarchy(session: Session, root_dataset_id: UUID) -> list[DatasetTable]:
+    """Retrieve all child datasets of the given root dataset, including the root itself.
+
+    The datasets are returned in the depth-first order, starting with the root dataset.
+    The relative order of children of any given node is the order in DatasetTable.children.
+    """
+    root_dataset = dataset_resolver.get_by_id(session=session, dataset_id=root_dataset_id)
+    if root_dataset is None:
+        raise ValueError(f"Dataset with id {root_dataset_id} not found.")
+
+    # Use a stack to perform depth-first traversal of the dataset hierarchy.
+    to_process = [root_dataset]
+    all_datasets = []
+    while to_process:
+        current_dataset = to_process.pop()
+        all_datasets.append(current_dataset)
+        to_process.extend(reversed(current_dataset.children))
+
+    return all_datasets