lightly-studio 0.3.1__py3-none-any.whl

lightly_studio/models/tag.py

@@ -0,0 +1,96 @@
+"""This module contains the Tag model and related enumerations."""
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, List, Literal, Optional
+from uuid import UUID, uuid4
+
+from sqlalchemy import UniqueConstraint
+from sqlalchemy.orm import Mapped
+from sqlmodel import Field, Relationship, SQLModel, String
+
+from lightly_studio.models.annotation.links import AnnotationTagLinkTable
+from lightly_studio.models.sample import SampleTagLinkTable
+
+if TYPE_CHECKING:
+    from lightly_studio.models.annotation.annotation_base import (
+        AnnotationBaseTable,
+    )
+    from lightly_studio.models.sample import SampleTable
+
+else:
+    SampleTable = object
+    TagTable = object
+    AnnotationBaseTable = object
+
+
+# TagKind is the kind of tag we support.
+TagKind = Literal[
+    "sample",
+    "annotation",
+]
+
+
+class TagBase(SQLModel):
+    """Base class for the Tag model."""
+
+    name: str
+    description: Optional[str] = ""
+    kind: TagKind = "sample"
+
+
+class TagCreate(TagBase):
+    """Tag model when creating."""
+
+    dataset_id: UUID
+
+
+class TagCreateBody(TagBase):
+    """Tag model when creating."""
+
+
+class TagUpdate(TagBase):
+    """Tag model when updating."""
+
+
+class TagUpdateBody(TagBase):
+    """Tag model when updating."""
+
+    dataset_id: Optional[UUID] = None
+
+
+class TagView(TagBase):
+    """Tag model when retrieving."""
+
+    tag_id: UUID
+    kind: TagKind
+    created_at: datetime
+    updated_at: datetime
+
+
+class TagTable(TagBase, table=True):
+    """This class defines the Tag model."""
+
+    __tablename__ = "tags"
+    # ensure there can only be one tag named "lightly_studio" per dataset
+    __table_args__ = (
+        UniqueConstraint("dataset_id", "kind", "name", name="unique_name_constraint"),
+    )
+    tag_id: UUID = Field(default_factory=uuid4, primary_key=True)
+    dataset_id: UUID
+    kind: TagKind = Field(sa_type=String)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), index=True)
+    updated_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+    )
+
+    """The sample ids associated with the tag."""
+    samples: Mapped[List["SampleTable"]] = Relationship(
+        back_populates="tags",
+        link_model=SampleTagLinkTable,
+    )
+
+    """The annotation ids associated with the tag (legacy bounding box)."""
+    annotations: Mapped[List["AnnotationBaseTable"]] = Relationship(
+        back_populates="tags",
+        link_model=AnnotationTagLinkTable,
+    )

lightly_studio/resolvers/__init__.py

@@ -0,0 +1,7 @@
+"""Resolvers for database operations."""
+
+from lightly_studio.resolvers import metadata_resolver
+
+__all__ = [
+    "metadata_resolver",
+]

lightly_studio/resolvers/annotation_label_resolver/__init__.py

@@ -0,0 +1,21 @@
+"""Handler for database operations related to annotation labels."""
+
+from .create import create
+from .delete import delete
+from .get_all import get_all
+from .get_by_id import get_by_id
+from .get_by_ids import get_by_ids
+from .get_by_label_name import get_by_label_name
+from .names_by_ids import names_by_ids
+from .update import update
+
+__all__ = [
+    "create",
+    "delete",
+    "get_all",
+    "get_by_id",
+    "get_by_ids",
+    "get_by_label_name",
+    "names_by_ids",
+    "update",
+]

lightly_studio/resolvers/annotation_label_resolver/create.py

@@ -0,0 +1,27 @@
+"""Create annotation label functionality."""
+
+from __future__ import annotations
+
+from sqlmodel import Session
+
+from lightly_studio.models.annotation_label import (
+    AnnotationLabelCreate,
+    AnnotationLabelTable,
+)
+
+
+def create(session: Session, label: AnnotationLabelCreate) -> AnnotationLabelTable:
+    """Create a new annotation label in the database.
+
+    Args:
+        session (Session): The database session.
+        label (AnnotationLabelCreate): The annotation label data to be created.
+
+    Returns:
+        AnnotationLabelTable: The newly created annotation label record.
+    """
+    db_label = AnnotationLabelTable.model_validate(label)
+    session.add(db_label)
+    session.commit()
+    session.refresh(db_label)
+    return db_label

lightly_studio/resolvers/annotation_label_resolver/delete.py

@@ -0,0 +1,28 @@
+"""Delete annotation label functionality."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from .get_by_id import get_by_id
+
+
+def delete(session: Session, label_id: UUID) -> bool:
+    """Delete an annotation label.
+
+    Args:
+        session (Session): The database session.
+        label_id (UUID): The unique identifier of the annotation label.
+
+    Returns:
+        bool: True if the label was deleted, False if the label was not found.
+    """
+    label = get_by_id(session=session, label_id=label_id)
+    if not label:
+        return False
+
+    session.delete(label)
+    session.commit()
+    return True

lightly_studio/resolvers/annotation_label_resolver/get_all.py

@@ -0,0 +1,22 @@
+"""Get all annotation labels functionality."""
+
+from __future__ import annotations
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.annotation_label import AnnotationLabelTable
+
+
+def get_all(session: Session) -> list[AnnotationLabelTable]:
+    """Retrieve all annotation labels.
+
+    Args:
+        session (Session): The database session.
+
+    Returns:
+        list[AnnotationLabelTable]: A list of annotation labels.
+    """
+    labels = session.exec(
+        select(AnnotationLabelTable).order_by(col(AnnotationLabelTable.created_at).asc())
+    ).all()
+    return list(labels) if labels else []

lightly_studio/resolvers/annotation_label_resolver/get_by_id.py

@@ -0,0 +1,24 @@
+"""Get annotation label by ID functionality."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, select
+
+from lightly_studio.models.annotation_label import AnnotationLabelTable
+
+
+def get_by_id(session: Session, label_id: UUID) -> AnnotationLabelTable | None:
+    """Retrieve a single annotation label by ID.
+
+    Args:
+        session (Session): The database session used to execute the query.
+        label_id (UUID): The unique identifier of the annotation label to get.
+
+    Returns:
+        AnnotationLabelTable | None: The annotation label if found, or None.
+    """
+    return session.exec(
+        select(AnnotationLabelTable).where(AnnotationLabelTable.annotation_label_id == label_id)
+    ).one_or_none()

lightly_studio/resolvers/annotation_label_resolver/get_by_ids.py

@@ -0,0 +1,25 @@
+"""Get annotation labels by IDs functionality."""
+
+from __future__ import annotations
+
+from typing import Sequence
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.annotation_label import AnnotationLabelTable
+
+
+def get_by_ids(session: Session, ids: Sequence[UUID]) -> list[AnnotationLabelTable]:
+    """Retrieve annotation labels by their IDs.
+
+    Output order matches the input order.
+    """
+    results = session.exec(
+        select(AnnotationLabelTable).where(
+            col(AnnotationLabelTable.annotation_label_id).in_(list(ids))
+        )
+    ).all()
+    # Return labels in the same order as the input ids.
+    label_map = {label.annotation_label_id: label for label in results}
+    return [label_map[id_] for id_ in ids if id_ in label_map]

lightly_studio/resolvers/annotation_label_resolver/get_by_label_name.py

@@ -0,0 +1,24 @@
+"""Handler for database operations related to annotation labels."""
+
+from __future__ import annotations
+
+from sqlmodel import Session, select
+
+from lightly_studio.models.annotation_label import (
+    AnnotationLabelTable,
+)
+
+
+def get_by_label_name(session: Session, label_name: str) -> AnnotationLabelTable | None:
+    """Retrieve a single annotation label by its name.
+
+    Args:
+        session: The database session to use for the query.
+        label_name: The name of the annotation label to retrieve.
+
+    Returns:
+        The AnnotationLabelTable instance if found, None otherwise.
+    """
+    return session.exec(
+        select(AnnotationLabelTable).where(AnnotationLabelTable.annotation_label_name == label_name)
+    ).one_or_none()

lightly_studio/resolvers/annotation_label_resolver/names_by_ids.py

@@ -0,0 +1,25 @@
+"""Get annotation label names by IDs functionality."""
+
+from __future__ import annotations
+
+from typing import Sequence
+from uuid import UUID
+
+from sqlmodel import Session
+
+from .get_by_ids import get_by_ids
+
+
+def names_by_ids(session: Session, ids: Sequence[UUID]) -> dict[str, str]:
+    """Return a dictionary mapping annotation label IDs to their names.
+
+    Args:
+        session (Session): The database session used to query the labels.
+        ids (Sequence[UUID]): A sequence of UUIDs of the annotation label IDs.
+
+    Returns:
+        dict[str, str]: A dictionary with string representations of the UUIDs
+        and the values are the corresponding annotation label names.
+    """
+    labels = get_by_ids(session=session, ids=ids)
+    return {str(label.annotation_label_id): label.annotation_label_name for label in labels}

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,33 @@
+"""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 import create
+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_by_id import get_by_id
+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",
+    "create_many",
+    "delete_annotation",
+    "delete_annotations",
+    "get_all",
+    "get_by_id",
+    "update_annotation_label",
+    "update_bounding_box",
+]

lightly_studio/resolvers/annotation_resolver/count_annotations_by_dataset.py

@@ -0,0 +1,120 @@
+"""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.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.annotation_id)).label("total_count"),
+        )
+        .join(
+            AnnotationBaseTable,
+            col(AnnotationBaseTable.annotation_label_id)
+            == col(AnnotationLabelTable.annotation_label_id),
+        )
+        .join(
+            SampleTable,
+            col(SampleTable.sample_id) == col(AnnotationBaseTable.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.annotation_id)).label("current_count"),
+        )
+        .join(
+            AnnotationBaseTable,
+            col(AnnotationBaseTable.annotation_label_id)
+            == col(AnnotationLabelTable.annotation_label_id),
+        )
+        .join(
+            SampleTable,
+            col(SampleTable.sample_id) == col(AnnotationBaseTable.sample_id),
+        )
+        .where(SampleTable.dataset_id == dataset_id)
+    )
+
+    # Add dimension filters
+    if min_width is not None:
+        filtered_query = filtered_query.where(SampleTable.width >= min_width)
+    if max_width is not None:
+        filtered_query = filtered_query.where(SampleTable.width <= max_width)
+    if min_height is not None:
+        filtered_query = filtered_query.where(SampleTable.height >= min_height)
+    if max_height is not None:
+        filtered_query = filtered_query.where(SampleTable.height <= max_height)
+
+    # Add label filter if specified
+    if filtered_labels:
+        filtered_query = filtered_query.where(
+            col(SampleTable.sample_id).in_(
+                select(SampleTable.sample_id)
+                .join(
+                    AnnotationBaseTable,
+                    col(SampleTable.sample_id) == col(AnnotationBaseTable.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.py

@@ -0,0 +1,19 @@
+"""Handler for database operations related to annotations."""
+
+from __future__ import annotations
+
+from sqlmodel import Session
+
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+    AnnotationCreate,
+)
+
+
+def create(session: Session, annotation: AnnotationCreate) -> AnnotationBaseTable:
+    """Create a new annotation in the database."""
+    db_annotation = AnnotationBaseTable.model_validate(annotation)
+    session.add(db_annotation)
+    session.commit()
+    session.refresh(db_annotation)
+    return db_annotation

lightly_studio/resolvers/annotation_resolver/create_many.py

@@ -0,0 +1,96 @@
+"""Handler for database operations related to annotations."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+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,
+)
+
+
+def create_many(
+    session: Session,
+    annotations: list[AnnotationCreate],
+) -> Sequence[AnnotationBaseTable]:
+    """Create many annotations with object detection details in bulk."""
+    # Step 1: Create all base annotations
+    base_annotations = []
+    object_detection_annotations = []
+    instance_segmentation_annotations = []
+    semantic_segmentation_annotations = []
+
+    for annotation_create in annotations:
+        # Create base annotation
+        db_base_annotation = AnnotationBaseTable(
+            annotation_label_id=annotation_create.annotation_label_id,
+            annotation_type=annotation_create.annotation_type,
+            annotation_task_id=annotation_create.annotation_task_id,
+            confidence=annotation_create.confidence,
+            dataset_id=annotation_create.dataset_id,
+            sample_id=annotation_create.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(
+                annotation_id=base_annotations[i].annotation_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(
+                annotation_id=base_annotations[i].annotation_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(
+                annotation_id=base_annotations[i].annotation_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 base_annotations

lightly_studio/resolvers/annotation_resolver/delete_annotation.py

@@ -0,0 +1,45 @@
+"""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.links import AnnotationTagLinkTable
+from lightly_studio.resolvers import annotation_resolver
+
+
+def delete_annotation(
+    session: Session,
+    annotation_id: UUID,
+) -> None:
+    """Delete all annotations and their tag links using filters.
+
+    Args:
+        session: Database session.
+        annotation_id: Annotation ID to filter by.
+    """
+    # Find annotation_ids to delete
+    annotation = annotation_resolver.get_by_id(
+        session,
+        annotation_id=annotation_id,
+    )
+    if not annotation:
+        return
+    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)
+
+    session.exec(  # type: ignore
+        delete(AnnotationTagLinkTable).where(
+            col(AnnotationTagLinkTable.annotation_id).in_([annotation.annotation_id])
+        )
+    )
+
+    session.commit()
+    session.delete(annotation)
+    session.commit()