lightly-studio 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

lightly_studio/metadata/compute_typicality.py

@@ -0,0 +1,67 @@
+"""Computes typicality from embeddings."""
+
+from uuid import UUID
+
+from lightly_mundig import Typicality  # type: ignore[import-untyped]
+from sqlmodel import Session
+
+from lightly_studio.dataset.env import LIGHTLY_STUDIO_LICENSE_KEY
+from lightly_studio.resolvers import (
+    metadata_resolver,
+    sample_embedding_resolver,
+)
+
+DEFAULT_NUM_NEAREST_NEIGHBORS = 20
+
+
+def compute_typicality_metadata(
+    session: Session,
+    dataset_id: UUID,
+    embedding_model_id: UUID,
+    metadata_name: str = "typicality",
+) -> None:
+    """Computes typicality for each sample in the dataset from embeddings.
+
+    Typicality is a measure of how representative a sample is of the dataset.
+    It is calculated for each sample from its K-nearest neighbors in the
+    embedding space.
+
+    The computed typicality values are stored as metadata for each sample.
+
+    Args:
+        session:
+            The database session.
+        dataset_id:
+            The ID of the dataset for which to compute the typicality.
+        embedding_model_id:
+            The ID of the embedding model to use for the computation.
+        metadata_name:
+            The name of the metadata field to store the typicality values in.
+            Defaults to "typicality".
+    """
+    license_key = LIGHTLY_STUDIO_LICENSE_KEY
+    if license_key is None:
+        raise ValueError(
+            "LIGHTLY_STUDIO_LICENSE_KEY environment variable is not set. "
+            "Please set it to your LightlyStudio license key."
+        )
+
+    samples = sample_embedding_resolver.get_all_by_dataset_id(
+        session=session, dataset_id=dataset_id, embedding_model_id=embedding_model_id
+    )
+
+    embeddings = [sample.embedding for sample in samples]
+    typicality = Typicality(embeddings=embeddings, token=license_key)
+    typicality_values = typicality.calculate_typicality(
+        num_nearest_neighbors=DEFAULT_NUM_NEAREST_NEIGHBORS
+    )
+    assert len(samples) == len(typicality_values), (
+        "The number of samples and computed typicality values must match"
+    )
+
+    metadata = [
+        (sample.sample_id, {metadata_name: typicality})
+        for sample, typicality in zip(samples, typicality_values)
+    ]
+
+    metadata_resolver.bulk_update_metadata(session, metadata)

lightly_studio/models/annotation/annotation_base.py

@@ -1,6 +1,7 @@
 """This module defines the base annotation model."""
 
 from datetime import datetime, timezone
+from enum import Enum
 from typing import TYPE_CHECKING, List, Optional
 from uuid import UUID, uuid4
 
@@ -22,10 +23,6 @@ from lightly_studio.models.annotation.semantic_segmentation import (
     SemanticSegmentationAnnotationTable,
     SemanticSegmentationAnnotationView,
 )
-from lightly_studio.models.annotation_task import (
-    AnnotationTaskTable,
-    AnnotationType,
-)
 
 if TYPE_CHECKING:
     from lightly_studio.models.annotation_label import (
@@ -41,6 +38,15 @@ else:
     AnnotationLabelTable = object
 
 
+class AnnotationType(str, Enum):
+    """The type of annotation task."""
+
+    CLASSIFICATION = "classification"
+    SEMANTIC_SEGMENTATION = "semantic_segmentation"
+    INSTANCE_SEGMENTATION = "instance_segmentation"
+    OBJECT_DETECTION = "object_detection"
+
+
 class AnnotationBaseTable(SQLModel, table=True):
     """Base class for all annotation models."""
 
@@ -51,9 +57,7 @@ class AnnotationBaseTable(SQLModel, table=True):
     annotation_id: UUID = Field(default_factory=uuid4, primary_key=True)
     annotation_type: AnnotationType
     annotation_label_id: UUID = Field(foreign_key="annotation_labels.annotation_label_id")
-    annotation_task_id: UUID = Field(
-        foreign_key="annotation_tasks.annotation_task_id",
-    )
+
     confidence: Optional[float] = None
     dataset_id: UUID = Field(foreign_key="datasets.dataset_id")
     sample_id: UUID = Field(foreign_key="samples.sample_id")
@@ -61,9 +65,6 @@ class AnnotationBaseTable(SQLModel, table=True):
     annotation_label: Mapped["AnnotationLabelTable"] = Relationship(
         sa_relationship_kwargs={"lazy": "select"},
     )
-    annotation_task: Mapped["AnnotationTaskTable"] = Relationship(
-        sa_relationship_kwargs={"lazy": "select"},
-    )
     sample: Mapped[Optional["SampleTable"]] = Relationship(
         sa_relationship_kwargs={"lazy": "select"},
     )
@@ -101,7 +102,6 @@ class AnnotationCreate(SQLModel):
     """ Required properties for all annotations. """
     annotation_label_id: UUID
     annotation_type: AnnotationType
-    annotation_task_id: UUID
     confidence: Optional[float] = None
     dataset_id: UUID
     sample_id: UUID
@@ -140,7 +140,6 @@ class AnnotationView(SQLModel):
     annotation_id: UUID
     annotation_type: AnnotationType
     annotation_label: AnnotationLabel
-    annotation_task_id: UUID
     confidence: Optional[float] = None
 
     object_detection_details: Optional[ObjectDetectionAnnotationView] = None

lightly_studio/models/caption.py

@@ -0,0 +1,73 @@
+"""This module defines the caption model."""
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, List, Optional
+from uuid import UUID, uuid4
+
+from pydantic import BaseModel, ConfigDict
+from pydantic import Field as PydanticField
+from sqlalchemy.orm import Mapped
+from sqlmodel import Field, Relationship, SQLModel
+
+if TYPE_CHECKING:
+    from lightly_studio.models.sample import SampleTable
+
+
+class CaptionTable(SQLModel, table=True):
+    """Class for caption model."""
+
+    __tablename__ = "caption"
+
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), index=True)
+
+    caption_id: UUID = Field(default_factory=uuid4, primary_key=True)
+    dataset_id: UUID = Field(foreign_key="datasets.dataset_id")
+    sample_id: UUID = Field(foreign_key="samples.sample_id")
+
+    sample: Mapped[Optional["SampleTable"]] = Relationship(
+        sa_relationship_kwargs={"lazy": "select"},
+    )
+
+    text: str
+
+
+class CaptionCreate(SQLModel):
+    """Input model for creating captions."""
+
+    dataset_id: UUID
+    sample_id: UUID
+    text: str
+
+
+class CaptionSampleView(SQLModel):
+    """Sample class for caption view."""
+
+    file_path_abs: str
+    file_name: str
+    dataset_id: UUID
+    sample_id: UUID
+
+
+class CaptionView(SQLModel):
+    """Response model for caption."""
+
+    sample_id: UUID
+    dataset_id: UUID
+    caption_id: UUID
+    text: str
+
+
+class CaptionDetailsView(CaptionView):
+    """Response model for caption."""
+
+    sample: CaptionSampleView
+
+
+class CaptionsListView(BaseModel):
+    """Response model for counted captions."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    captions: List[CaptionDetailsView] = PydanticField(..., alias="data")
+    total_count: int
+    next_cursor: Optional[int] = PydanticField(..., alias="nextCursor")

lightly_studio/models/dataset.py

@@ -19,8 +19,7 @@ from lightly_studio.resolvers.samples_filter import SampleFilter
 class DatasetBase(SQLModel):
     """Base class for the Dataset model."""
 
-    name: str
-    directory: str
+    name: str = Field(unique=True, index=True)
 
 
 class DatasetCreate(DatasetBase):

lightly_studio/models/metadata.py

@@ -188,7 +188,7 @@ class SampleMetadataTable(MetadataBase, table=True):
     """This class defines the SampleMetadataTable model."""
 
     __tablename__ = "metadata"
-    sample_id: UUID = Field(foreign_key="samples.sample_id")
+    sample_id: UUID = Field(foreign_key="samples.sample_id", unique=True)
 
     sample: SampleTable = Relationship(back_populates="metadata_dict")
 

lightly_studio/models/sample.py

@@ -170,8 +170,8 @@ class SampleView(SQLModel):
     file_path_abs: str
     sample_id: UUID
     dataset_id: UUID
-    annotations: List["AnnotationView"] = Field([])
-    tags: List[SampleViewTag] = Field([])
+    annotations: List["AnnotationView"]
+    tags: List[SampleViewTag]
     metadata_dict: Optional["SampleMetadataView"] = None
     width: int
     height: int

lightly_studio/resolvers/annotation_label_resolver/__init__.py

@@ -2,7 +2,7 @@
 
 from .create import create
 from .delete import delete
-from .get_all import get_all
+from .get_all import get_all, get_all_sorted_alphabetically
 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
@@ -13,6 +13,7 @@ __all__ = [
     "create",
     "delete",
     "get_all",
+    "get_all_sorted_alphabetically",
     "get_by_id",
     "get_by_ids",
     "get_by_label_name",

lightly_studio/resolvers/annotation_label_resolver/get_all.py

@@ -20,3 +20,18 @@ def get_all(session: Session) -> list[AnnotationLabelTable]:
         select(AnnotationLabelTable).order_by(col(AnnotationLabelTable.created_at).asc())
     ).all()
     return list(labels) if labels else []
+
+
+def get_all_sorted_alphabetically(session: Session) -> list[AnnotationLabelTable]:
+    """Retrieve all annotation labels sorted alphabetically.
+
+    Args:
+        session (Session): The database session.
+
+    Returns:
+        list[AnnotationLabelTable]: A list of annotation labels.
+    """
+    labels = session.exec(
+        select(AnnotationLabelTable).order_by(col(AnnotationLabelTable.annotation_label_name).asc())
+    ).all()
+    return list(labels) if labels else []

lightly_studio/resolvers/annotation_resolver/__init__.py

@@ -3,7 +3,6 @@
 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,
@@ -12,7 +11,7 @@ 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.get_by_id import get_by_id, get_by_ids
 from lightly_studio.resolvers.annotation_resolver.update_annotation_label import (
     update_annotation_label,
 )
@@ -22,12 +21,12 @@ from lightly_studio.resolvers.annotation_resolver.update_bounding_box import (
 
 __all__ = [
     "count_annotations_by_dataset",
-    "create",
     "create_many",
     "delete_annotation",
     "delete_annotations",
     "get_all",
     "get_by_id",
+    "get_by_ids",
     "update_annotation_label",
     "update_bounding_box",
 ]

lightly_studio/resolvers/annotation_resolver/create_many.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
+from uuid import UUID
 
 from sqlmodel import Session
 
@@ -24,7 +25,7 @@ from lightly_studio.models.annotation.semantic_segmentation import (
 def create_many(
     session: Session,
     annotations: list[AnnotationCreate],
-) -> Sequence[AnnotationBaseTable]:
+) -> Sequence[UUID]:
     """Create many annotations with object detection details in bulk."""
     # Step 1: Create all base annotations
     base_annotations = []
@@ -37,7 +38,6 @@ def create_many(
         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,
@@ -93,4 +93,4 @@ def create_many(
     # Commit everything
     session.commit()
 
-    return base_annotations
+    return [annotation.annotation_id for annotation in base_annotations]

lightly_studio/resolvers/annotation_resolver/delete_annotation.py

@@ -26,7 +26,7 @@ def delete_annotation(
         annotation_id=annotation_id,
     )
     if not annotation:
-        return
+        raise ValueError(f"Annotation {annotation_id} not found")
     if annotation.object_detection_details:
         session.delete(annotation.object_detection_details)
     if annotation.instance_segmentation_details:

lightly_studio/resolvers/annotation_resolver/delete_annotations.py

@@ -18,14 +18,12 @@ from lightly_studio.resolvers.annotations.annotations_filter import (
 
 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
@@ -33,9 +31,15 @@ def delete_annotations(
         session,
         filters=AnnotationsFilter(
             annotation_label_ids=annotation_label_ids,
-            annotation_task_ids=annotation_task_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.annotation_id for annotation in annotations]
     # TODO(Horatiu, 06/2025): Check if there is a way to delete the links
     # automatically using SQLModel/SQLAlchemy.

lightly_studio/resolvers/annotation_resolver/get_by_id.py

@@ -2,9 +2,10 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from uuid import UUID
 
-from sqlmodel import Session, select
+from sqlmodel import Session, col, select
 
 from lightly_studio.models.annotation.annotation_base import (
     AnnotationBaseTable,
@@ -16,3 +17,20 @@ def get_by_id(session: Session, annotation_id: UUID) -> AnnotationBaseTable | No
     return session.exec(
         select(AnnotationBaseTable).where(AnnotationBaseTable.annotation_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.annotation_id).in_(annotation_ids)
+        )
+    ).all()

lightly_studio/resolvers/annotation_resolver/update_annotation_label.py

@@ -111,7 +111,6 @@ def update_annotation_label(
             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,

lightly_studio/resolvers/annotations/annotations_filter.py

@@ -7,8 +7,7 @@ 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.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
@@ -30,9 +29,6 @@ class AnnotationsFilter(BaseModel):
         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,
@@ -51,12 +47,6 @@ class AnnotationsFilter(BaseModel):
         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(

lightly_studio/resolvers/caption_resolver.py

@@ -0,0 +1,80 @@
+"""Resolvers for caption."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from uuid import UUID
+
+from pydantic import BaseModel
+from sqlmodel import Session, col, func, select
+
+from lightly_studio.api.routes.api.validators import Paginated
+from lightly_studio.models.caption import CaptionCreate, CaptionTable
+
+
+class GetAllCaptionsResult(BaseModel):
+    """Result wrapper for caption listings."""
+
+    captions: Sequence[CaptionTable]
+    total_count: int
+    next_cursor: int | None = None
+
+
+def create_many(session: Session, captions: Sequence[CaptionCreate]) -> list[CaptionTable]:
+    """Create many captions in bulk.
+
+    Args:
+        session: Database session
+        captions: The captions to create
+
+    Returns:
+        The created captions
+    """
+    if not captions:
+        return []
+
+    db_captions = [CaptionTable.model_validate(caption) for caption in captions]
+    session.bulk_save_objects(db_captions)
+    session.commit()
+    return db_captions
+
+
+def get_all(
+    session: Session,
+    dataset_id: UUID,
+    pagination: Paginated | None = None,
+) -> GetAllCaptionsResult:
+    """Get all captions from the database.
+
+    Args:
+        session: Database session
+        dataset_id: dataset_id parameter to filter the query
+        pagination: Optional pagination parameters
+
+    Returns:
+        List of captions matching the filters, total number of captions, next cursor (pagination)
+    """
+    query = select(CaptionTable).order_by(
+        col(CaptionTable.created_at).asc(),
+        col(CaptionTable.caption_id).asc(),
+    )
+    count_query = select(func.count()).select_from(CaptionTable)
+
+    query = query.where(CaptionTable.dataset_id == dataset_id)
+    count_query = count_query.where(CaptionTable.dataset_id == dataset_id)
+
+    if pagination is not None:
+        query = query.offset(pagination.offset).limit(pagination.limit)
+
+    captions = session.exec(query).all()
+    total_count = session.exec(count_query).one()
+
+    next_cursor: int | None = None
+    if pagination and pagination.offset + pagination.limit < total_count:
+        next_cursor = pagination.offset + pagination.limit
+
+    return GetAllCaptionsResult(
+        captions=captions,
+        total_count=total_count,
+        next_cursor=next_cursor,
+    )

lightly_studio/resolvers/dataset_resolver.py

@@ -41,6 +41,9 @@ class ExportFilter(BaseModel):
 
 def create(session: Session, dataset: DatasetCreate) -> DatasetTable:
     """Create a new dataset in the database."""
+    existing = 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()
@@ -69,12 +72,7 @@ def get_by_id(session: Session, dataset_id: UUID) -> DatasetTable | None:
 
 def get_by_name(session: Session, name: str) -> DatasetTable | None:
     """Retrieve a single dataset by name."""
-    datasets = session.exec(select(DatasetTable).where(DatasetTable.name == name)).all()
-    if len(datasets) == 0:
-        return None
-    if len(datasets) > 1:
-        raise ValueError(f"Cannot retrieve a dataset, found multiple with name '{name}'.")
-    return datasets[0]
+    return session.exec(select(DatasetTable).where(DatasetTable.name == name)).one_or_none()
 
 
 def update(session: Session, dataset_id: UUID, dataset_data: DatasetCreate) -> DatasetTable:
@@ -84,7 +82,6 @@ def update(session: Session, dataset_id: UUID, dataset_data: DatasetCreate) -> D
         raise ValueError(f"Dataset ID was not found '{dataset_id}'.")
 
     dataset.name = dataset_data.name
-    dataset.directory = dataset_data.directory
     dataset.updated_at = datetime.now(timezone.utc)
 
     session.commit()

lightly_studio/resolvers/metadata_resolver/__init__.py

@@ -1,14 +1,14 @@
 """Metadata resolver module."""
 
 from lightly_studio.resolvers.metadata_resolver.sample import (
-    bulk_set_metadata,
+    bulk_update_metadata,
     get_by_sample_id,
     get_value_for_sample,
     set_value_for_sample,
 )
 
 __all__ = [
-    "bulk_set_metadata",
+    "bulk_update_metadata",
     "get_by_sample_id",
     "get_value_for_sample",
     "set_value_for_sample",

lightly_studio/resolvers/metadata_resolver/sample/__init__.py

@@ -1,7 +1,7 @@
 """Resolvers for metadata operations."""
 
-from .bulk_set_metadata import (
-    bulk_set_metadata,
+from .bulk_update_metadata import (
+    bulk_update_metadata,
 )
 from .get_by_sample_id import (
     get_by_sample_id,
@@ -14,7 +14,7 @@ from .set_value_for_sample import (
 )
 
 __all__ = [
-    "bulk_set_metadata",
+    "bulk_update_metadata",
     "get_by_sample_id",
     "get_value_for_sample",
     "set_value_for_sample",

lightly_studio/resolvers/metadata_resolver/sample/bulk_update_metadata.py

@@ -0,0 +1,46 @@
+"""Resolver for operations for setting metadata."""
+
+from __future__ import annotations
+
+from typing import Any
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.metadata import SampleMetadataTable
+
+
+def bulk_update_metadata(
+    session: Session,
+    sample_metadata: list[tuple[UUID, dict[str, Any]]],
+) -> None:
+    """Bulk insert or update metadata for multiple samples.
+
+    If a sample does not have metadata, a new metadata row is created.
+    If a sample already has metadata, the new key-value pairs are merged with the existing metadata.
+
+    Args:
+        session: The database session.
+        sample_metadata: List of (sample_id, metadata_dict) tuples.
+    """
+    # TODO(Mihnea, 10/2025): Consider using SQLAlchemy's bulk operations
+    #  (Session.bulk_insert/update_mappings) if performance becomes a bottleneck.
+    if not sample_metadata:
+        return
+
+    # Get all existing metadata rows for the given sample IDs.
+    sample_ids = [s[0] for s in sample_metadata]
+    existing_metadata = session.exec(
+        select(SampleMetadataTable).where(col(SampleMetadataTable.sample_id).in_(sample_ids))
+    ).all()
+    sample_id_to_existing_metadata = {meta.sample_id: meta for meta in existing_metadata}
+
+    for sample_id, new_metadata in sample_metadata:
+        metadata = sample_id_to_existing_metadata.get(
+            sample_id, SampleMetadataTable(sample_id=sample_id)
+        )
+        for key, value in new_metadata.items():
+            metadata.set_value(key, value)
+        session.add(metadata)
+
+    session.commit()