lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/resolvers/annotation_resolver/get_all_with_payload.py

@@ -0,0 +1,179 @@
+"""Get all annotations with payload resolver."""
+
+from __future__ import annotations
+
+from typing import Any
+from uuid import UUID
+
+from sqlalchemy.orm import aliased, joinedload, load_only
+from sqlmodel import Session, col, func, select
+from sqlmodel.sql.expression import Select
+
+from lightly_studio.api.routes.api.validators import Paginated
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+    AnnotationWithPayloadAndCountView,
+    AnnotationWithPayloadView,
+    ImageAnnotationView,
+    SampleAnnotationView,
+    VideoAnnotationView,
+    VideoFrameAnnotationView,
+)
+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.video import VideoFrameTable, VideoTable
+from lightly_studio.resolvers import dataset_resolver
+from lightly_studio.resolvers.annotations.annotations_filter import (
+    AnnotationsFilter,
+)
+
+
+def get_all_with_payload(
+    session: Session,
+    dataset_id: UUID,
+    pagination: Paginated | None = None,
+    filters: AnnotationsFilter | None = None,
+) -> AnnotationWithPayloadAndCountView:
+    """Get all annotations with payload from the database.
+
+    Args:
+        session: Database session
+        pagination: Optional pagination parameters
+        filters: Optional filters to apply to the query
+        dataset_id: ID of the dataset to get annotations for
+
+    Returns:
+        List of annotations matching the filters with payload
+    """
+    parent_dataset = dataset_resolver.get_parent_dataset_id(session=session, dataset_id=dataset_id)
+
+    if parent_dataset is None:
+        raise ValueError(f"Dataset with id {dataset_id} does not have a parent dataset.")
+
+    sample_type = parent_dataset.sample_type
+
+    base_query = _build_base_query(sample_type=sample_type)
+
+    if filters:
+        base_query = filters.apply(base_query)
+
+    annotations_query = base_query.order_by(
+        *_extra_order_by(sample_type=sample_type),
+        col(AnnotationBaseTable.created_at).asc(),
+        col(AnnotationBaseTable.sample_id).asc(),
+    )
+
+    total_count_query = select(func.count()).select_from(base_query.subquery())
+    total_count = session.exec(total_count_query).one()
+
+    if pagination is not None:
+        annotations_query = annotations_query.offset(pagination.offset).limit(pagination.limit)
+
+    next_cursor = None
+    if pagination and pagination.offset + pagination.limit < total_count:
+        next_cursor = pagination.offset + pagination.limit
+
+    rows = session.exec(annotations_query).all()
+
+    return AnnotationWithPayloadAndCountView(
+        total_count=total_count,
+        next_cursor=next_cursor,
+        annotations=[
+            AnnotationWithPayloadView(
+                parent_sample_type=sample_type,
+                annotation=annotation,
+                parent_sample_data=_serialize_annotation_payload(payload),
+            )
+            for annotation, payload in rows
+        ],
+    )
+
+
+def _build_base_query(
+    sample_type: SampleType,
+) -> Select[tuple[AnnotationBaseTable, Any]]:
+    if sample_type == SampleType.IMAGE:
+        # this alias is needed to avoid name clashes in joins
+        SampleFromImage = aliased(SampleTable)  # noqa: N806
+
+        return (
+            select(AnnotationBaseTable, ImageTable)
+            .join(
+                ImageTable,
+                col(ImageTable.sample_id) == col(AnnotationBaseTable.parent_sample_id),
+            )
+            .join(SampleFromImage, col(SampleFromImage.sample_id) == col(ImageTable.sample_id))
+            .options(
+                load_only(
+                    ImageTable.file_path_abs,  # type: ignore[arg-type]
+                    ImageTable.sample_id,  # type: ignore[arg-type]
+                    ImageTable.height,  # type: ignore[arg-type]
+                    ImageTable.width,  # type: ignore[arg-type]
+                ),
+                joinedload(ImageTable.sample).load_only(
+                    SampleTable.dataset_id,  # type: ignore[arg-type]
+                ),
+            )
+        )
+
+    if sample_type in (SampleType.VIDEO_FRAME, SampleType.VIDEO):
+        return (
+            select(AnnotationBaseTable, VideoFrameTable)
+            .join(
+                VideoFrameTable,
+                col(VideoFrameTable.sample_id) == col(AnnotationBaseTable.parent_sample_id),
+            )
+            .join(VideoFrameTable.video)
+            .options(
+                load_only(VideoFrameTable.sample_id),  # type: ignore[arg-type]
+                joinedload(VideoFrameTable.video).load_only(
+                    VideoTable.height,  # type: ignore[arg-type]
+                    VideoTable.width,  # type: ignore[arg-type]
+                    VideoTable.file_path_abs,  # type: ignore[arg-type]
+                ),
+            )
+        )
+
+    raise NotImplementedError(f"Unsupported sample type: {sample_type}")
+
+
+def _extra_order_by(sample_type: SampleType) -> list[Any]:
+    """Return extra order by clauses for the query."""
+    if sample_type == SampleType.IMAGE:
+        return [
+            col(ImageTable.file_path_abs).asc(),
+        ]
+
+    if sample_type in (SampleType.VIDEO_FRAME, SampleType.VIDEO):
+        return [
+            col(VideoTable.file_path_abs).asc(),
+        ]
+
+    return []
+
+
+def _serialize_annotation_payload(
+    payload: ImageTable | VideoFrameTable,
+) -> ImageAnnotationView | VideoFrameAnnotationView:
+    """Serialize annotation based on sample type."""
+    if isinstance(payload, ImageTable):
+        return ImageAnnotationView(
+            height=payload.height,
+            width=payload.width,
+            file_path_abs=payload.file_path_abs,
+            sample_id=payload.sample_id,
+            sample=SampleAnnotationView(dataset_id=payload.sample.dataset_id),
+        )
+
+    if isinstance(payload, VideoFrameTable):
+        return VideoFrameAnnotationView(
+            sample_id=payload.sample_id,
+            video=VideoAnnotationView(
+                width=payload.video.width,
+                height=payload.video.height,
+                file_path_abs=payload.video.file_path_abs,
+            ),
+        )
+
+    raise NotImplementedError("Unsupported sample type")

lightly_studio/resolvers/annotation_resolver/get_by_id.py

@@ -0,0 +1,34 @@
+"""Handler for database operations related to annotations."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+)
+
+
+def get_by_id(session: Session, annotation_id: UUID) -> AnnotationBaseTable | None:
+    """Retrieve a single annotation by ID."""
+    return session.exec(
+        select(AnnotationBaseTable).where(AnnotationBaseTable.sample_id == annotation_id)
+    ).one_or_none()
+
+
+def get_by_ids(session: Session, annotation_ids: Sequence[UUID]) -> Sequence[AnnotationBaseTable]:
+    """Retrieve multiple annotations by their IDs.
+
+    Args:
+        session: The database session to use for the query.
+        annotation_ids: A list of annotation IDs to retrieve.
+
+    Returns:
+        A list of annotations matching the provided IDs.
+    """
+    return session.exec(
+        select(AnnotationBaseTable).where(col(AnnotationBaseTable.sample_id).in_(annotation_ids))
+    ).all()

lightly_studio/resolvers/annotation_resolver/get_by_id_with_payload.py

@@ -0,0 +1,130 @@
+"""Get annotation by id with payload resolver."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlalchemy.orm import aliased, joinedload, load_only
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+    AnnotationDetailsWithPayloadView,
+    ImageAnnotationDetailsView,
+    VideoFrameAnnotationDetailsView,
+)
+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.video import VideoFrameTable, VideoTable
+from lightly_studio.resolvers import dataset_resolver
+
+
+def get_by_id_with_payload(
+    session: Session,
+    sample_id: UUID,
+) -> AnnotationDetailsWithPayloadView | None:
+    """Get annotation by its id with payload from the database.
+
+    Args:
+        session: Database session
+        sample_id: ID of the sample to get annotations for
+
+    Returns:
+        Returns annotations with payload
+    """
+    parent_dataset = dataset_resolver.get_parent_dataset_by_sample_id(
+        session=session, sample_id=sample_id
+    )
+
+    if parent_dataset is None:
+        raise ValueError(f"Sample with id {sample_id} does not have a parent dataset.")
+
+    parent_sample_type = parent_dataset.sample_type
+
+    if parent_sample_type in SampleType.VIDEO_FRAME:
+        return _get_video_frame_annotation_by_id(
+            session=session, sample_id=sample_id, parent_sample_type=parent_sample_type
+        )
+    if parent_sample_type == SampleType.IMAGE:
+        return _get_image_annotation_by_id(session=session, sample_id=sample_id)
+
+    raise NotImplementedError("Unsupported sample type")
+
+
+SampleFromImage = aliased(SampleTable)
+
+
+def _get_image_annotation_by_id(
+    session: Session, sample_id: UUID
+) -> AnnotationDetailsWithPayloadView | None:
+    base_query = (
+        select(AnnotationBaseTable, ImageTable)
+        .join(
+            ImageTable,
+            col(ImageTable.sample_id) == col(AnnotationBaseTable.parent_sample_id),
+        )
+        .join(SampleFromImage, col(SampleFromImage.sample_id) == col(ImageTable.sample_id))
+        .options(
+            load_only(
+                ImageTable.file_path_abs,  # type: ignore[arg-type]
+                ImageTable.file_name,  # type: ignore[arg-type]
+                ImageTable.sample_id,  # type: ignore[arg-type]
+                ImageTable.height,  # type: ignore[arg-type]
+                ImageTable.width,  # type: ignore[arg-type]
+            ),
+        )
+        .where(col(AnnotationBaseTable.sample_id) == sample_id)
+    )
+
+    row = session.exec(base_query).one_or_none()
+
+    if row is None:
+        return None
+
+    (annotation, payload) = row
+
+    return AnnotationDetailsWithPayloadView(
+        parent_sample_type=SampleType.IMAGE,
+        annotation=annotation,
+        parent_sample_data=ImageAnnotationDetailsView.from_image_table(payload),
+    )
+
+
+def _get_video_frame_annotation_by_id(
+    session: Session, sample_id: UUID, parent_sample_type: SampleType
+) -> AnnotationDetailsWithPayloadView | None:
+    base_query = (
+        select(AnnotationBaseTable, VideoFrameTable)
+        .join(
+            VideoFrameTable,
+            col(VideoFrameTable.sample_id) == col(AnnotationBaseTable.parent_sample_id),
+        )
+        .join(VideoFrameTable.video)
+        .options(
+            load_only(
+                VideoFrameTable.sample_id,  # type: ignore[arg-type]
+                VideoFrameTable.frame_number,  # type: ignore[arg-type]
+                VideoFrameTable.frame_timestamp_s,  # type: ignore[arg-type]
+            ),
+            joinedload(VideoFrameTable.video).load_only(
+                VideoTable.height,  # type: ignore[arg-type]
+                VideoTable.width,  # type: ignore[arg-type]
+                VideoTable.file_path_abs,  # type: ignore[arg-type]
+            ),
+        )
+        .where(col(AnnotationBaseTable.sample_id) == sample_id)
+    )
+
+    row = session.exec(base_query).one_or_none()
+
+    if row is None:
+        return None
+
+    (annotation, payload) = row
+
+    return AnnotationDetailsWithPayloadView(
+        parent_sample_type=parent_sample_type,
+        annotation=annotation,
+        parent_sample_data=VideoFrameAnnotationDetailsView.from_video_frame_table(payload),
+    )

lightly_studio/resolvers/annotation_resolver/update_annotation_label.py

@@ -0,0 +1,142 @@
+"""Module for handling the update of annotation labels in the database."""
+
+from __future__ import annotations
+
+from typing import TypeVar
+from uuid import UUID
+
+from sqlmodel import Session, SQLModel
+
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+)
+from lightly_studio.models.annotation.instance_segmentation import (
+    InstanceSegmentationAnnotationTable,
+)
+from lightly_studio.models.annotation.links import AnnotationTagLinkTable
+from lightly_studio.models.annotation.object_detection import ObjectDetectionAnnotationTable
+from lightly_studio.models.annotation.semantic_segmentation import (
+    SemanticSegmentationAnnotationTable,
+)
+from lightly_studio.resolvers import (
+    annotation_resolver,
+)
+
+T = TypeVar("T", bound=SQLModel)
+
+
+def update_annotation_label(
+    session: Session, annotation_id: UUID, annotation_label_id: UUID
+) -> AnnotationBaseTable:
+    """Update the label of an annotation.
+
+    Args:
+        session: Database session for executing the operation.
+        annotation_id: UUID of the annotation to update.
+        annotation_label_id: UUID of the new label to assign to the annotation.
+
+    Returns:
+        The updated annotation with the new label assigned.
+
+    Raises:
+        ValueError: If the annotation is not found.
+    """
+    annotation = annotation_resolver.get_by_id(session, annotation_id)
+    if not annotation:
+        raise ValueError(f"Annotation with ID {annotation_id} not found.")
+
+    # DuckDB has no "looking ahead" functionality for referenced tables.
+    # We need to work around this by deleting the existing and re-inserting
+    # Check https://duckdb.org/docs/stable/sql/indexes.html#over-eager-constraint-checking-in-foreign-keys
+
+    # DuckDB has no "looking ahead" functionality for referenced tables and neither does it support cascading updates.  # noqa: E501
+    # Herefore we need to delete and re-insert the affected rows.
+    # more information can be found in the DuckDB documentation https://duckdb.org/docs/stable/sql/statements/create_table.html.
+    try:
+        # copy content
+        annotation_copy = annotation.model_copy(update={"annotation_label_id": annotation_label_id})
+
+        annotation_type = annotation_copy.annotation_type
+
+        annotation_tags = [
+            AnnotationTagLinkTable(
+                annotation_sample_id=annotation_copy.sample_id,
+                tag_id=tag.tag_id,
+            )
+            for tag in annotation.tags
+        ]
+
+        # we need to create a new annotation details before committing
+        # because copy will be gone with the commit
+        instance_segmentation = (
+            InstanceSegmentationAnnotationTable(
+                sample_id=annotation_copy.sample_id,
+                segmentation_mask=annotation_copy.instance_segmentation_details.segmentation_mask,
+                x=annotation_copy.instance_segmentation_details.x,
+                y=annotation_copy.instance_segmentation_details.y,
+                width=annotation_copy.instance_segmentation_details.width,
+                height=annotation_copy.instance_segmentation_details.height,
+            )
+            if annotation_type == "instance_segmentation"
+            and annotation_copy.instance_segmentation_details
+            else None
+        )
+
+        object_detection = (
+            ObjectDetectionAnnotationTable(
+                sample_id=annotation_copy.sample_id,
+                x=annotation_copy.object_detection_details.x,
+                y=annotation_copy.object_detection_details.y,
+                width=annotation_copy.object_detection_details.width,
+                height=annotation_copy.object_detection_details.height,
+            )
+            if annotation_type == "object_detection" and annotation_copy.object_detection_details
+            else None
+        )
+
+        semantic_segmentation = (
+            SemanticSegmentationAnnotationTable(
+                sample_id=annotation_copy.sample_id,
+                segmentation_mask=annotation_copy.semantic_segmentation_details.segmentation_mask,
+            )
+            if annotation_type == "semantic_segmentation"
+            and annotation_copy.semantic_segmentation_details
+            else None
+        )
+
+        # delete
+        annotation_resolver.delete_annotation(session, annotation.sample_id, delete_sample=False)
+
+        new_annotation = AnnotationBaseTable(
+            sample_id=annotation_copy.sample_id,
+            annotation_label_id=annotation_copy.annotation_label_id,
+            annotation_type=annotation_copy.annotation_type,
+            confidence=annotation_copy.confidence,
+            created_at=annotation_copy.created_at,
+            parent_sample_id=annotation_copy.parent_sample_id,
+        )
+
+        session.add(new_annotation)
+
+        if instance_segmentation:
+            session.add(instance_segmentation)
+
+        if object_detection:
+            session.add(object_detection)
+
+        if semantic_segmentation:
+            session.add(semantic_segmentation)
+
+        if annotation_tags:
+            session.add_all(annotation_tags)
+
+        session.commit()
+        session.flush()
+
+        return annotation_copy
+    except Exception:
+        # Explicit rollback to be safe, then re-raise the original error.
+        session.rollback()
+        raise
+
+    return annotation

lightly_studio/resolvers/annotation_resolver/update_bounding_box.py

@@ -0,0 +1,68 @@
+"""Module for handling the update of annotation bounding box coordinates in the database."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable
+from lightly_studio.resolvers import annotation_resolver
+
+
+@dataclass
+class BoundingBoxCoordinates:
+    """Represents bounding box coordinates."""
+
+    x: int
+    y: int
+    width: int
+    height: int
+
+
+def update_bounding_box(
+    session: Session,
+    annotation_id: UUID,
+    coordinates: BoundingBoxCoordinates,
+) -> AnnotationBaseTable:
+    """Update the bounding box coordinates of an annotation.
+
+    Args:
+        session: Database session for executing the operation.
+        annotation_id: UUID of the annotation to update.
+        coordinates: New bounding box coordinates.
+
+    Returns:
+        The updated annotation with the new bounding box coordinates.
+
+    Raises:
+        ValueError: If the annotation is not found.
+    """
+    annotation = annotation_resolver.get_by_id(session, annotation_id)
+    if not annotation:
+        raise ValueError(f"Annotation with ID {annotation_id} not found.")
+
+    try:
+        if annotation.object_detection_details:
+            annotation.object_detection_details.x = coordinates.x
+            annotation.object_detection_details.y = coordinates.y
+            annotation.object_detection_details.width = coordinates.width
+            annotation.object_detection_details.height = coordinates.height
+            session.add(annotation.object_detection_details)
+
+        elif annotation.instance_segmentation_details:
+            annotation.instance_segmentation_details.x = coordinates.x
+            annotation.instance_segmentation_details.y = coordinates.y
+            annotation.instance_segmentation_details.width = coordinates.width
+            annotation.instance_segmentation_details.height = coordinates.height
+            session.add(annotation.instance_segmentation_details)
+        else:
+            raise ValueError("Annotation type does not support bounding boxes.")
+
+        session.commit()
+        session.refresh(annotation)
+        return annotation
+    except Exception:
+        session.rollback()
+        raise

lightly_studio/resolvers/annotations/__init__.py

@@ -0,0 +1 @@
+

lightly_studio/resolvers/annotations/annotations_filter.py

@@ -0,0 +1,88 @@
+"""Filtering functionality for annotations."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from pydantic import BaseModel, Field
+from sqlmodel import col
+
+from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable, AnnotationType
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.tag import TagTable
+from lightly_studio.type_definitions import QueryType
+
+
+class AnnotationsFilter(BaseModel):
+    """Handles filtering for annotation queries."""
+
+    annotation_types: list[AnnotationType] | None = Field(
+        default=None,
+        description="Types of annotation to filter (e.g., 'object_detection')",
+    )
+    dataset_ids: list[UUID] | None = Field(default=None, description="List of dataset UUIDs")
+    annotation_label_ids: list[UUID] | None = Field(
+        default=None, description="List of annotation label UUIDs"
+    )
+    annotation_tag_ids: list[UUID] | None = Field(default=None, description="List of tag UUIDs")
+    sample_tag_ids: list[UUID] | None = Field(
+        default=None,
+        description="List of sample tag UUIDs to filter annotations by",
+    )
+    sample_ids: list[UUID] | None = Field(
+        default=None, description="List of sample UUIDs to filter annotations by"
+    )
+
+    def apply(
+        self,
+        query: QueryType,
+    ) -> QueryType:
+        """Apply filters to an annotation query.
+
+        Args:
+            query: The base query to apply filters to
+            annotation_table: The SQLModel table class for the annotation type
+
+        Returns:
+            The query with filters applied
+        """
+        # Filter by dataset
+        if self.dataset_ids:
+            query = query.join(AnnotationBaseTable.sample).where(
+                col(SampleTable.dataset_id).in_(self.dataset_ids)
+            )
+
+        # Filter by annotation label
+        if self.annotation_label_ids:
+            query = query.where(
+                col(AnnotationBaseTable.annotation_label_id).in_(self.annotation_label_ids)
+            )
+
+        # Filter by annotation tags
+        if self.annotation_tag_ids:
+            query = (
+                query.join(AnnotationBaseTable.tags)
+                .where(
+                    AnnotationBaseTable.tags.any(col(TagTable.tag_id).in_(self.annotation_tag_ids))
+                )
+                .distinct()
+            )
+
+        # Filter by sample tags
+        if self.sample_tag_ids:
+            query = (
+                query.join(AnnotationBaseTable.parent_sample)
+                .join(SampleTable.tags)
+                .where(SampleTable.tags.any(col(TagTable.tag_id).in_(self.sample_tag_ids)))
+                .distinct()
+            )
+
+        # Filter by sample ids
+        if self.sample_ids:
+            query = query.where(col(AnnotationBaseTable.parent_sample_id).in_(self.sample_ids))
+
+        # Filter by annotation type
+        if self.annotation_types:
+            query = query.where(col(AnnotationBaseTable.annotation_type).in_(self.annotation_types))
+
+        return query