lightly-studio 0.3.1__py3-none-any.whl

lightly_studio/resolvers/annotation_resolver/delete_annotations.py

@@ -0,0 +1,56 @@
+"""Handler for database operations related to annotations."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, col, delete
+
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+)
+from lightly_studio.models.annotation.links import AnnotationTagLinkTable
+from lightly_studio.resolvers import annotation_resolver
+from lightly_studio.resolvers.annotations.annotations_filter import (
+    AnnotationsFilter,
+)
+
+
+def delete_annotations(
+    session: Session,
+    annotation_task_ids: list[UUID] | None,
+    annotation_label_ids: list[UUID] | None,
+) -> None:
+    """Delete all annotations and their tag links using filters.
+
+    Args:
+        session: Database session.
+        annotation_task_ids: List of annotation task IDs to filter by.
+        annotation_label_ids: List of annotation label IDs to filter by.
+    """
+    # Find annotation_ids to delete
+    annotations = annotation_resolver.get_all(
+        session,
+        filters=AnnotationsFilter(
+            annotation_label_ids=annotation_label_ids,
+            annotation_task_ids=annotation_task_ids,
+        ),
+    ).annotations
+    annotation_ids = [annotation.annotation_id for annotation in annotations]
+    # TODO(Horatiu, 06/2025): Check if there is a way to delete the links
+    # automatically using SQLModel/SQLAlchemy.
+    if annotation_ids:
+        # Delete tag links first
+        session.exec(  # type: ignore
+            delete(AnnotationTagLinkTable).where(
+                col(AnnotationTagLinkTable.annotation_id).in_(annotation_ids)
+            )
+        )
+        session.commit()
+        # Now delete the annotations themselves
+        session.exec(  # type: ignore
+            delete(AnnotationBaseTable).where(
+                col(AnnotationBaseTable.annotation_id).in_(annotation_ids)
+            )
+        )
+        session.commit()

lightly_studio/resolvers/annotation_resolver/get_all.py

@@ -0,0 +1,74 @@
+"""Handler for database operations related to annotations."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from pydantic import BaseModel
+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.resolvers.annotations.annotations_filter import (
+    AnnotationsFilter,
+)
+
+
+class GetAllAnnotationsResult(BaseModel):
+    """Result of getting all annotations."""
+
+    annotations: Sequence[AnnotationBaseTable]
+
+    total_count: int
+
+    next_cursor: int | None = None
+
+
+def get_all(
+    session: Session,
+    pagination: Paginated | None = None,
+    filters: AnnotationsFilter | None = None,
+) -> GetAllAnnotationsResult:
+    """Get all annotations from the database.
+
+    Args:
+        session: Database session
+        pagination: Optional pagination parameters
+        filters: Optional filters to apply to the query
+
+    Returns:
+        List of annotations matching the filters
+    """
+    annotations_statement = select(AnnotationBaseTable)
+
+    annotations_statement = annotations_statement.order_by(
+        col(AnnotationBaseTable.created_at).asc(),
+        col(AnnotationBaseTable.annotation_id).asc(),
+    )
+
+    total_count_statement = select(func.count()).select_from(AnnotationBaseTable)
+
+    # Apply filters if provided
+    if filters is not None:
+        annotations_statement = filters.apply(annotations_statement)
+        total_count_statement = filters.apply(total_count_statement)
+
+    # Apply pagination if provided
+    if pagination is not None:
+        annotations_statement = annotations_statement.offset(pagination.offset).limit(
+            pagination.limit
+        )
+
+    total_count = session.exec(total_count_statement).one()
+
+    next_cursor = None
+    if pagination and pagination.offset + pagination.limit < total_count:
+        next_cursor = pagination.offset + pagination.limit
+
+    return GetAllAnnotationsResult(
+        annotations=session.exec(annotations_statement).all(),
+        total_count=total_count,
+        next_cursor=next_cursor,
+    )

lightly_studio/resolvers/annotation_resolver/get_by_id.py

@@ -0,0 +1,18 @@
+"""Handler for database operations related to annotations."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, 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.annotation_id == annotation_id)
+    ).one_or_none()

lightly_studio/resolvers/annotation_resolver/update_annotation_label.py

@@ -0,0 +1,144 @@
+"""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_id=annotation_copy.annotation_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(
+                annotation_id=annotation_copy.annotation_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(
+                annotation_id=annotation_copy.annotation_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(
+                annotation_id=annotation_copy.annotation_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.annotation_id)
+
+        new_annotation = AnnotationBaseTable(
+            annotation_id=annotation_copy.annotation_id,
+            annotation_label_id=annotation_copy.annotation_label_id,
+            annotation_type=annotation_copy.annotation_type,
+            annotation_task_id=annotation_copy.annotation_task_id,
+            confidence=annotation_copy.confidence,
+            created_at=annotation_copy.created_at,
+            dataset_id=annotation_copy.dataset_id,
+            sample_id=annotation_copy.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/annotation_task_resolver.py

@@ -0,0 +1,31 @@
+"""Resolver for annotation tasks."""
+
+from typing import List, Optional, Sequence
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.annotation_task import AnnotationTaskTable
+
+
+def create(session: Session, annotation_task: AnnotationTaskTable) -> AnnotationTaskTable:
+    """Create a new annotation task."""
+    session.add(annotation_task)
+    session.commit()
+    session.refresh(annotation_task)
+    return annotation_task
+
+
+def get_by_id(session: Session, annotation_task_id: UUID) -> Optional[AnnotationTaskTable]:
+    """Get an annotation task by ID."""
+    statement = select(AnnotationTaskTable).where(
+        AnnotationTaskTable.annotation_task_id == annotation_task_id
+    )
+    return session.exec(statement).first()
+
+
+def get_all(session: Session) -> List[AnnotationTaskTable]:
+    """Get all annotation tasks."""
+    statement = select(AnnotationTaskTable).order_by(col(AnnotationTaskTable.created_at).asc())
+    results: Sequence[AnnotationTaskTable] = session.exec(statement).all()
+    return list(results)

lightly_studio/resolvers/annotations/__init__.py

@@ -0,0 +1 @@
+

lightly_studio/resolvers/annotations/annotations_filter.py

@@ -0,0 +1,89 @@
+"""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
+from lightly_studio.models.annotation_task import 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",
+    )
+    annotation_task_ids: list[UUID] | None = Field(
+        default=None, description="List of annotation task UUIDs"
+    )
+
+    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.where(col(AnnotationBaseTable.dataset_id).in_(self.dataset_ids))
+
+        # Filter by annotation task
+        if self.annotation_task_ids:
+            query = query.where(
+                col(AnnotationBaseTable.annotation_task_id).in_(self.annotation_task_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.sample)
+                .join(SampleTable.tags)
+                .where(SampleTable.tags.any(col(TagTable.tag_id).in_(self.sample_tag_ids)))
+                .distinct()
+            )
+
+        # Filter by annotation type
+        if self.annotation_types:
+            query = query.where(col(AnnotationBaseTable.annotation_type).in_(self.annotation_types))
+
+        return query