lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/resolvers/annotation_label_resolver/update.py

@@ -0,0 +1,38 @@
+"""Update annotation label functionality."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.models.annotation_label import (
+    AnnotationLabelCreate,
+    AnnotationLabelTable,
+)
+
+from .get_by_id import get_by_id
+
+
+def update(
+    session: Session, label_id: UUID, label_data: AnnotationLabelCreate
+) -> AnnotationLabelTable | None:
+    """Update an existing annotation label.
+
+    Args:
+        session (Session): The database session.
+        label_id (UUID): The identifier of the annotation label to update.
+        label_data (AnnotationLabelCreate): The new data.
+
+    Returns:
+        AnnotationLabelTable | None: The updated annotation label if it exists,
+        otherwise None.
+    """
+    label = get_by_id(session=session, label_id=label_id)
+    if not label:
+        return None
+
+    label.annotation_label_name = label_data.annotation_label_name
+    session.commit()
+    session.refresh(label)
+    return label

lightly_studio/resolvers/annotation_resolver/__init__.py

@@ -0,0 +1,40 @@
+"""Resolvers for database operations."""
+
+from lightly_studio.resolvers.annotation_resolver.count_annotations_by_dataset import (
+    count_annotations_by_dataset,
+)
+from lightly_studio.resolvers.annotation_resolver.create_many import create_many
+from lightly_studio.resolvers.annotation_resolver.delete_annotation import (
+    delete_annotation,
+)
+from lightly_studio.resolvers.annotation_resolver.delete_annotations import (
+    delete_annotations,
+)
+from lightly_studio.resolvers.annotation_resolver.get_all import get_all
+from lightly_studio.resolvers.annotation_resolver.get_all_with_payload import (
+    get_all_with_payload,
+)
+from lightly_studio.resolvers.annotation_resolver.get_by_id import get_by_id, get_by_ids
+from lightly_studio.resolvers.annotation_resolver.get_by_id_with_payload import (
+    get_by_id_with_payload,
+)
+from lightly_studio.resolvers.annotation_resolver.update_annotation_label import (
+    update_annotation_label,
+)
+from lightly_studio.resolvers.annotation_resolver.update_bounding_box import (
+    update_bounding_box,
+)
+
+__all__ = [
+    "count_annotations_by_dataset",
+    "create_many",
+    "delete_annotation",
+    "delete_annotations",
+    "get_all",
+    "get_all_with_payload",
+    "get_by_id",
+    "get_by_id_with_payload",
+    "get_by_ids",
+    "update_annotation_label",
+    "update_bounding_box",
+]

lightly_studio/resolvers/annotation_resolver/count_annotations_by_dataset.py

@@ -0,0 +1,129 @@
+"""Handler for database operations related to annotations."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, col, func, 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 count_annotations_by_dataset(  # noqa: PLR0913 // FIXME: refactor to use proper pydantic
+    session: Session,
+    dataset_id: UUID,
+    filtered_labels: list[str] | None = None,
+    min_width: int | None = None,
+    max_width: int | None = None,
+    min_height: int | None = None,
+    max_height: int | None = None,
+    tag_ids: list[UUID] | None = None,
+) -> list[tuple[str, int, int]]:
+    """Count annotations for a specific dataset.
+
+    Annotations for a specific dataset are grouped by annotation
+    label name and counted for total and filtered.
+    """
+    # Query for total counts (unfiltered)
+    total_counts_query = (
+        select(
+            AnnotationLabelTable.annotation_label_name,
+            func.count(col(AnnotationBaseTable.sample_id)).label("total_count"),
+        )
+        .join(
+            AnnotationBaseTable,
+            col(AnnotationBaseTable.annotation_label_id)
+            == col(AnnotationLabelTable.annotation_label_id),
+        )
+        .join(
+            ImageTable,
+            col(ImageTable.sample_id) == col(AnnotationBaseTable.parent_sample_id),
+        )
+        .join(
+            SampleTable,
+            col(SampleTable.sample_id) == col(ImageTable.sample_id),
+        )
+        .where(SampleTable.dataset_id == dataset_id)
+        .group_by(AnnotationLabelTable.annotation_label_name)
+        .order_by(col(AnnotationLabelTable.annotation_label_name).asc())
+    )
+
+    total_counts = {row[0]: row[1] for row in session.exec(total_counts_query).all()}
+
+    # Build filtered query for current counts
+    filtered_query = (
+        select(
+            AnnotationLabelTable.annotation_label_name,
+            func.count(col(AnnotationBaseTable.sample_id)).label("current_count"),
+        )
+        .join(
+            AnnotationBaseTable,
+            col(AnnotationBaseTable.annotation_label_id)
+            == col(AnnotationLabelTable.annotation_label_id),
+        )
+        .join(
+            ImageTable,
+            col(ImageTable.sample_id) == col(AnnotationBaseTable.parent_sample_id),
+        )
+        .join(
+            SampleTable,
+            col(SampleTable.sample_id) == col(ImageTable.sample_id),
+        )
+        .where(SampleTable.dataset_id == dataset_id)
+    )
+
+    # Add dimension filters
+    if min_width is not None:
+        filtered_query = filtered_query.where(ImageTable.width >= min_width)
+    if max_width is not None:
+        filtered_query = filtered_query.where(ImageTable.width <= max_width)
+    if min_height is not None:
+        filtered_query = filtered_query.where(ImageTable.height >= min_height)
+    if max_height is not None:
+        filtered_query = filtered_query.where(ImageTable.height <= max_height)
+
+    # Add label filter if specified
+    if filtered_labels:
+        filtered_query = filtered_query.where(
+            col(ImageTable.sample_id).in_(
+                select(ImageTable.sample_id)
+                .join(
+                    AnnotationBaseTable,
+                    col(ImageTable.sample_id) == col(AnnotationBaseTable.parent_sample_id),
+                )
+                .join(
+                    AnnotationLabelTable,
+                    col(AnnotationBaseTable.annotation_label_id)
+                    == col(AnnotationLabelTable.annotation_label_id),
+                )
+                .where(col(AnnotationLabelTable.annotation_label_name).in_(filtered_labels))
+            )
+        )
+
+    # filter by tag_ids
+    if tag_ids:
+        filtered_query = (
+            filtered_query.join(AnnotationBaseTable.tags)
+            .where(AnnotationBaseTable.tags.any(col(TagTable.tag_id).in_(tag_ids)))
+            .distinct()
+        )
+
+    # Group by label name and sort
+    filtered_query = filtered_query.group_by(AnnotationLabelTable.annotation_label_name).order_by(
+        col(AnnotationLabelTable.annotation_label_name).asc()
+    )
+
+    _rows = session.exec(filtered_query).all()
+
+    current_counts = {row[0]: row[1] for row in _rows}
+
+    return [
+        (label, current_counts.get(label, 0), total_count)
+        for label, total_count in total_counts.items()
+    ]

lightly_studio/resolvers/annotation_resolver/create_many.py

@@ -0,0 +1,124 @@
+"""Handler for database operations related to annotations."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+    AnnotationCreate,
+)
+from lightly_studio.models.annotation.instance_segmentation import (
+    InstanceSegmentationAnnotationTable,
+)
+from lightly_studio.models.annotation.object_detection import (
+    ObjectDetectionAnnotationTable,
+)
+from lightly_studio.models.annotation.semantic_segmentation import (
+    SemanticSegmentationAnnotationTable,
+)
+from lightly_studio.models.dataset import SampleType
+from lightly_studio.models.sample import SampleCreate
+from lightly_studio.resolvers import dataset_resolver, sample_resolver
+
+
+def create_many(
+    session: Session,
+    parent_dataset_id: UUID,
+    annotations: list[AnnotationCreate],
+) -> list[UUID]:
+    """Create multiple annotations in bulk with their respective type-specific details.
+
+    Creates base annotations and their associated type-specific details (object detection,
+    instance segmentation, or semantic segmentation) in the annotation dataset child of
+    the provided parent dataset.
+
+    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: SQLAlchemy session for database operations.
+        parent_dataset_id: UUID of the parent dataset.
+        annotations: List of annotation objects to create.
+
+    Returns:
+        List of created annotation IDs.
+    """
+    # Step 1: Create all base annotations
+    base_annotations = []
+    object_detection_annotations = []
+    instance_segmentation_annotations = []
+    semantic_segmentation_annotations = []
+    annotation_dataset_id = dataset_resolver.get_or_create_child_dataset(
+        session=session, dataset_id=parent_dataset_id, sample_type=SampleType.ANNOTATION
+    )
+
+    sample_ids = sample_resolver.create_many(
+        session=session,
+        samples=[SampleCreate(dataset_id=annotation_dataset_id) for _ in annotations],
+    )
+    for annotation_create, sample_id in zip(annotations, sample_ids):
+        # Create base annotation
+        db_base_annotation = AnnotationBaseTable(
+            sample_id=sample_id,
+            annotation_label_id=annotation_create.annotation_label_id,
+            annotation_type=annotation_create.annotation_type,
+            confidence=annotation_create.confidence,
+            dataset_id=annotation_dataset_id,
+            parent_sample_id=annotation_create.parent_sample_id,
+        )
+
+        # Set other relationship details to None
+        db_base_annotation.instance_segmentation_details = None
+        db_base_annotation.semantic_segmentation_details = None
+        db_base_annotation.object_detection_details = None
+
+        base_annotations.append(db_base_annotation)
+
+    # Bulk save base annotations and flush to get IDs
+    session.bulk_save_objects(base_annotations)
+    session.flush()
+
+    # Step 2: Create specific annotation details
+    for i, annotation_create in enumerate(annotations):
+        # Create object detection details
+        if base_annotations[i].annotation_type == "object_detection":
+            db_object_detection = ObjectDetectionAnnotationTable(
+                sample_id=base_annotations[i].sample_id,
+                x=annotation_create.x,
+                y=annotation_create.y,
+                width=annotation_create.width,
+                height=annotation_create.height,
+            )
+            object_detection_annotations.append(db_object_detection)
+
+        # Create instance segmentation details
+        elif base_annotations[i].annotation_type == "instance_segmentation":
+            db_instance_segmentation = InstanceSegmentationAnnotationTable(
+                sample_id=base_annotations[i].sample_id,
+                segmentation_mask=annotation_create.segmentation_mask,
+                x=annotation_create.x,
+                y=annotation_create.y,
+                width=annotation_create.width,
+                height=annotation_create.height,
+            )
+            instance_segmentation_annotations.append(db_instance_segmentation)
+        elif base_annotations[i].annotation_type == "semantic_segmentation":
+            db_semantic_segmentation = SemanticSegmentationAnnotationTable(
+                sample_id=base_annotations[i].sample_id,
+                segmentation_mask=annotation_create.segmentation_mask,
+            )
+            semantic_segmentation_annotations.append(db_semantic_segmentation)
+
+    # Bulk save object detection annotations
+    session.bulk_save_objects(object_detection_annotations)
+    session.bulk_save_objects(instance_segmentation_annotations)
+    session.bulk_save_objects(semantic_segmentation_annotations)
+
+    # Commit everything
+    session.commit()
+
+    return [annotation.sample_id for annotation in base_annotations]

lightly_studio/resolvers/annotation_resolver/delete_annotation.py

@@ -0,0 +1,87 @@
+"""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.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.models.sample import SampleTable
+from lightly_studio.resolvers import annotation_resolver
+
+
+def delete_annotation(
+    session: Session,
+    annotation_id: UUID,
+    delete_sample: bool = True,
+) -> None:
+    """Delete all annotations and their tag links using filters.
+
+    Args:
+        session: Database session.
+        annotation_id: Annotation ID to filter by.
+        delete_sample: Whether to also delete the annotation's sample. Defaults to True.
+                      Set to False when updating an annotation (to reuse the sample).
+    """
+    # Find annotation_ids to delete
+    annotation = annotation_resolver.get_by_id(
+        session,
+        annotation_id=annotation_id,
+    )
+    if not annotation:
+        raise ValueError(f"Annotation {annotation_id} not found")
+
+    # Store the annotation's sample_id before deletion
+    annotation_sample_id = annotation.sample_id
+
+    session.exec(  # type: ignore
+        delete(ObjectDetectionAnnotationTable).where(
+            col(ObjectDetectionAnnotationTable.sample_id) == annotation.sample_id
+        )
+    )
+    session.exec(  # type: ignore
+        delete(InstanceSegmentationAnnotationTable).where(
+            col(InstanceSegmentationAnnotationTable.sample_id) == annotation.sample_id
+        )
+    )
+    session.exec(  # type: ignore
+        delete(SemanticSegmentationAnnotationTable).where(
+            col(SemanticSegmentationAnnotationTable.sample_id) == annotation.sample_id
+        )
+    )
+    session.commit()
+
+    # Delete tag links
+    session.exec(  # type: ignore
+        delete(AnnotationTagLinkTable).where(
+            col(AnnotationTagLinkTable.annotation_sample_id).in_([annotation.sample_id])
+        )
+    )
+    session.commit()
+
+    # Delete the annotation using explicit DELETE to avoid relationship cascade issues
+    session.exec(  # type: ignore
+        delete(AnnotationBaseTable).where(
+            col(AnnotationBaseTable.sample_id) == annotation.sample_id
+        )
+    )
+    session.commit()
+
+    # Then delete the annotation's sample (created specifically for this annotation)
+    # unless we're keeping it for reuse (e.g., when updating annotation label)
+    if delete_sample:
+        annotation_sample = session.get(SampleTable, annotation_sample_id)
+        if annotation_sample:
+            session.delete(annotation_sample)
+            session.commit()

lightly_studio/resolvers/annotation_resolver/delete_annotations.py

@@ -0,0 +1,60 @@
+"""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_label_ids: list[UUID] | None,
+) -> None:
+    """Delete all annotations and their tag links using filters.
+
+    Args:
+        session: Database session.
+        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,
+        ),
+    ).annotations
+    for annotation in annotations:
+        if annotation.object_detection_details:
+            session.delete(annotation.object_detection_details)
+        if annotation.instance_segmentation_details:
+            session.delete(annotation.instance_segmentation_details)
+        if annotation.semantic_segmentation_details:
+            session.delete(annotation.semantic_segmentation_details)
+    annotation_ids = [annotation.sample_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_sample_id).in_(annotation_ids)
+            )
+        )
+        session.commit()
+        # Now delete the annotations themselves
+        session.exec(  # type: ignore
+            delete(AnnotationBaseTable).where(
+                col(AnnotationBaseTable.sample_id).in_(annotation_ids)
+            )
+        )
+        session.commit()

lightly_studio/resolvers/annotation_resolver/get_all.py

@@ -0,0 +1,85 @@
+"""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.models.image import ImageTable
+from lightly_studio.models.video import VideoFrameTable, VideoTable
+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.outerjoin(
+            ImageTable, col(ImageTable.sample_id) == col(AnnotationBaseTable.parent_sample_id)
+        )
+        .outerjoin(
+            VideoFrameTable,
+            col(VideoFrameTable.sample_id) == col(AnnotationBaseTable.parent_sample_id),
+        )
+        .outerjoin(VideoTable, col(VideoTable.sample_id) == col(VideoFrameTable.parent_sample_id))
+        .order_by(
+            func.coalesce(ImageTable.file_path_abs, VideoTable.file_path_abs, "").asc(),
+            col(AnnotationBaseTable.created_at).asc(),
+            col(AnnotationBaseTable.sample_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,
+    )