lightly-studio 0.3.1__py3-none-any.whl

lightly_studio/resolvers/tag_resolver.py

@@ -0,0 +1,276 @@
+"""Handler for database operations related to tags."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from uuid import UUID
+
+import sqlmodel
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable
+from lightly_studio.models.annotation.links import AnnotationTagLinkTable
+from lightly_studio.models.sample import SampleTable, SampleTagLinkTable
+from lightly_studio.models.tag import TagCreate, TagTable, TagUpdate
+
+
+def create(session: Session, tag: TagCreate) -> TagTable:
+    """Create a new tag in the database."""
+    db_tag = TagTable.model_validate(tag)
+    session.add(db_tag)
+    session.commit()
+    session.refresh(db_tag)
+    return db_tag
+
+
+# TODO(Michal, 06/2025): Use Paginated struct instead of offset/limit.
+def get_all_by_dataset_id(
+    session: Session, dataset_id: UUID, offset: int = 0, limit: int | None = None
+) -> list[TagTable]:
+    """Retrieve all tags with pagination."""
+    query = (
+        select(TagTable)
+        .where(TagTable.dataset_id == dataset_id)
+        .order_by(col(TagTable.created_at).asc(), col(TagTable.tag_id).asc())
+        .offset(offset)
+    )
+    if limit is not None:
+        query = query.limit(limit)
+    tags = session.exec(query).all()
+    return list(tags) if tags else []
+
+
+def get_by_id(session: Session, tag_id: UUID) -> TagTable | None:
+    """Retrieve a single tag by ID."""
+    return session.exec(select(TagTable).where(TagTable.tag_id == tag_id)).one_or_none()
+
+
+def get_by_name(session: Session, tag_name: str, dataset_id: UUID | None) -> TagTable | None:
+    """Retrieve a single tag by ID."""
+    if dataset_id:
+        return session.exec(
+            select(TagTable)
+            .where(TagTable.dataset_id == dataset_id)
+            .where(TagTable.name == tag_name)
+        ).one_or_none()
+    return session.exec(select(TagTable).where(TagTable.name == tag_name)).one_or_none()
+
+
+def update(session: Session, tag_id: UUID, tag_data: TagUpdate) -> TagTable | None:
+    """Update an existing tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag:
+        return None
+
+    # due to duckdb/OLAP optimisations, update operations effecting unique
+    # constraints (e.g colums) will lead to a unique constraint violation.
+    # This is due to a update is implemented as delete+insert. The error
+    # happens only within the same session.
+    # To fix it, we can delete, commit + insert a new tag.
+    # https://duckdb.org/docs/sql/indexes#over-eager-unique-constraint-checking
+    session.delete(tag)
+    session.commit()
+
+    # create clone of tag with updated values
+    tag_updated = TagTable.model_validate(tag)
+    tag_updated.name = tag_data.name
+    tag_updated.description = tag_data.description
+    tag_updated.updated_at = datetime.now(timezone.utc)
+
+    session.add(tag_updated)
+    session.commit()
+    session.refresh(tag_updated)
+    return tag_updated
+
+
+def delete(session: Session, tag_id: UUID) -> bool:
+    """Delete a tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag:
+        return False
+
+    session.delete(tag)
+    session.commit()
+    return True
+
+
+def add_tag_to_sample(
+    session: Session,
+    tag_id: UUID,
+    sample: SampleTable,
+) -> SampleTable | None:
+    """Add a tag to a sample."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "sample":
+        raise ValueError(f"Tag {tag_id} is not of kind 'sample'")
+
+    sample.tags.append(tag)
+    session.add(sample)
+    session.commit()
+    session.refresh(sample)
+    return sample
+
+
+def remove_tag_from_sample(
+    session: Session,
+    tag_id: UUID,
+    sample: SampleTable,
+) -> SampleTable | None:
+    """Remove a tag from a sample."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "sample":
+        raise ValueError(f"Tag {tag_id} is not of kind 'sample'")
+
+    sample.tags.remove(tag)
+    session.add(sample)
+    session.commit()
+    session.refresh(sample)
+    return sample
+
+
+def add_tag_to_annotation(
+    session: Session,
+    tag_id: UUID,
+    annotation: AnnotationBaseTable,
+) -> AnnotationBaseTable | None:
+    """Add a tag to a annotation."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "annotation":
+        raise ValueError(f"Tag {tag_id} is not of kind 'annotation'")
+
+    annotation.tags.append(tag)
+    session.add(annotation)
+    session.commit()
+    session.refresh(annotation)
+    return annotation
+
+
+def assign_tag_to_annotation(
+    session: Session,
+    tag: TagTable,
+    annotation: AnnotationBaseTable,
+) -> AnnotationBaseTable:
+    """Add a tag to a annotation."""
+    annotation.tags.append(tag)
+    session.add(annotation)
+    session.commit()
+    session.refresh(annotation)
+    return annotation
+
+
+def remove_tag_from_annotation(
+    session: Session,
+    tag_id: UUID,
+    annotation: AnnotationBaseTable,
+) -> AnnotationBaseTable | None:
+    """Remove a tag from a annotation."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "annotation":
+        raise ValueError(f"Tag {tag_id} is not of kind 'annotation'")
+
+    annotation.tags.remove(tag)
+    session.add(annotation)
+    session.commit()
+    session.refresh(annotation)
+    return annotation
+
+
+def add_sample_ids_to_tag_id(
+    session: Session,
+    tag_id: UUID,
+    sample_ids: list[UUID],
+) -> TagTable | None:
+    """Add a list of sample_ids to a tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "sample":
+        raise ValueError(f"Tag {tag_id} is not of kind 'sample'")
+
+    for sample_id in sample_ids:
+        session.merge(SampleTagLinkTable(sample_id=sample_id, tag_id=tag_id))
+
+    session.commit()
+    session.refresh(tag)
+    return tag
+
+
+def remove_sample_ids_from_tag_id(
+    session: Session,
+    tag_id: UUID,
+    sample_ids: list[UUID],
+) -> TagTable | None:
+    """Remove a list of sample_ids to a tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "sample":
+        raise ValueError(f"Tag {tag_id} is not of kind 'sample'")
+
+    session.exec(  # type:ignore[call-overload]
+        sqlmodel.delete(SampleTagLinkTable).where(
+            col(SampleTagLinkTable.tag_id) == tag_id,
+            col(SampleTagLinkTable.sample_id).in_(sample_ids),
+        )
+    )
+
+    session.commit()
+    session.refresh(tag)
+    return tag
+
+
+def add_annotation_ids_to_tag_id(
+    session: Session,
+    tag_id: UUID,
+    annotation_ids: list[UUID],
+) -> TagTable | None:
+    """Add a list of annotation_ids to a tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "annotation":
+        raise ValueError(f"Tag {tag_id} is not of kind 'annotation'")
+
+    for annotation_id in annotation_ids:
+        session.merge(
+            AnnotationTagLinkTable(
+                tag_id=tag_id,
+                annotation_id=annotation_id,
+            )
+        )
+
+    session.commit()
+    session.refresh(tag)
+    return tag
+
+
+def remove_annotation_ids_from_tag_id(
+    session: Session,
+    tag_id: UUID,
+    annotation_ids: list[UUID],
+) -> TagTable | None:
+    """Remove a list of things to a tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "annotation":
+        raise ValueError(f"Tag {tag_id} is not of kind 'annotation'")
+
+    session.exec(  # type:ignore[call-overload]
+        sqlmodel.delete(AnnotationTagLinkTable).where(
+            col(AnnotationTagLinkTable.tag_id) == tag_id,
+            col(AnnotationTagLinkTable.annotation_id).in_(annotation_ids),
+        )
+    )
+
+    session.commit()
+    session.refresh(tag)
+    return tag

lightly_studio/selection/README.md

@@ -0,0 +1,6 @@
+# Selection
+
+The architecture follows the design document. 
+See https://docs.google.com/document/d/1ZRICdFmfJmxUBy3FFoeUWsAgsCNWDHg8CK5MJiGmX74/edit?tab=t.kbfvnrepsuf#bookmark=id.4wbh2eqhgokk
+The design document also contains an architecture diagram showing the
+relationship between the various components of the system.

lightly_studio/selection/mundig.py

@@ -0,0 +1,105 @@
+"""Python interface to the Mundig selection algorithm."""
+
+from __future__ import annotations
+
+from typing import Iterable
+
+# TODO(Malte, 08/2025): About the type ignore:
+# Use pyo3 typing stubs once they are implemented.
+# See https://github.com/PyO3/pyo3/issues/510
+# Or remove the type ignore once typing stubs were added manually.
+import lightly_mundig  # type: ignore[import-untyped]
+import numpy as np
+from environs import Env
+
+
+class Mundig:
+    """Python wrapper for the Mundig selection algorithm.
+
+    This class provides a Python interface to the lightly_mundig Rust library
+    for sample selection.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the Mundig selection interface."""
+        # Read LIGHTLY_STUDIO_LICENSE_KEY with .env file support
+        env = Env()
+        env.read_env()
+        license_key = env.str("LIGHTLY_STUDIO_LICENSE_KEY", default=None)
+        if license_key is None:
+            raise ValueError(
+                "LIGHTLY_STUDIO_LICENSE_KEY environment variable is not set. "
+                "Please set it to your LightlyStudio license key."
+            )
+
+        self.mundig = lightly_mundig.Selection(token=license_key)
+
+        self.n_input_samples: int | None = None
+
+    def run(self, n_samples: int) -> list[int]:
+        """Run the selection algorithm and return selected sample indices.
+
+        Args:
+            n_samples: The number of samples to select.
+
+        Returns:
+            A list of indices of the selected samples.
+        """
+        selected: list[int] = self.mundig.run_selection(
+            n_total_samples=self.n_input_samples, n_samples_to_select=n_samples
+        )
+        return selected
+
+    def add_diversity(self, embeddings: Iterable[Iterable[float]], strength: float = 1.0) -> None:
+        """Add diversity-based selection using sample embeddings.
+
+        Args:
+            embeddings:
+                The embeddings of each sample.
+                First dimension is over the samples, the second dimension is
+                the embedding size. The embedding size must be the same for
+                all samples.
+            strength:
+                The strength of the diversity strategy.
+
+        """
+        # Convert to ndarray with float32 dtype if not already
+        if isinstance(embeddings, np.ndarray) and embeddings.dtype == np.float32:
+            embeddings_ndarray = embeddings
+        else:
+            embeddings_ndarray = np.array(embeddings, dtype=np.float32)
+        self._check_consistent_input_size(embeddings_ndarray.shape[0])
+        self.mundig.add_diversifying_strategy(embeddings=embeddings_ndarray, strength=strength)
+
+    def add_weighting(self, weights: Iterable[float], strength: float = 1.0) -> None:
+        """Add a weighting strategy.
+
+        Args:
+            weights:
+                The weight or importance or utility of each sample.
+            strength:
+                The strength of the weighting strategy.
+        """
+        weights_ndarray = np.array(weights, dtype=np.float32)
+        self._check_consistent_input_size(weights_ndarray.shape[0])
+        self.mundig.add_weighting_strategy(weights=weights_ndarray, strength=strength)
+
+    def _check_consistent_input_size(self, n_input_samples_strategy: int) -> None:
+        """Assert that input samples count is consistent across strategies.
+
+        Args:
+            n_input_samples_strategy:
+                The number of input samples in the currently added strategy.
+
+        Raises:
+            ValueError:
+                If the number of input samples in the new strategy differs
+                from the one used in previous strategies.
+        """
+        if self.n_input_samples is None:
+            self.n_input_samples = n_input_samples_strategy
+        elif self.n_input_samples != n_input_samples_strategy:
+            raise ValueError(
+                f"Expected {self.n_input_samples} input samples, "
+                f"but the latest strategy passed {n_input_samples_strategy}."
+            )

lightly_studio/selection/select.py

@@ -0,0 +1,96 @@
+"""Provides the user python interface to selection."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.resolvers.samples_filter import SampleFilter
+from lightly_studio.selection.select_via_db import select_via_database
+from lightly_studio.selection.selection_config import (
+    EmbeddingDiversityStrategy,
+    SelectionConfig,
+    SelectionStrategy,
+)
+
+
+class Selection:
+    """User selection interface for the dataset."""
+
+    # TODO(Malte, 08/2025): Create this class within the DatasetView.
+    # Then the arguments can be passed directly from the DatasetView.
+    # Example:
+    # class DatasetView:
+    #     def __init__(self, dataset_id: UUID, session: Session):
+    #         self.select = Select(dataset_id, session)
+    # User interface:
+    # dataset_view = ...
+    # dataset_view.select.diverse(...)
+    #
+    # See https://docs.google.com/document/d/1ZRICdFmfJmxUBy3FFoeUWsAgsCNWDHg8CK5MJiGmX74/edit?tab=t.kbfvnrepsuf#bookmark=id.8klhhwr5q4dp
+
+    def __init__(self, dataset_id: UUID, session: Session):
+        """Creates the interface to run selection.
+
+        Args:
+            dataset_id: The ID of the dataset to select from.
+            session: The database session to use for selection.
+
+        """
+        self.dataset_id = dataset_id
+        self.session = session
+
+    def diverse(
+        self,
+        n_samples_to_select: int,
+        selection_result_tag_name: str,
+        embedding_model_name: str | None = None,
+        sample_filter: SampleFilter | None = None,
+    ) -> None:
+        """Selects a diverse subset of the dataset.
+
+        Args:
+            n_samples_to_select: The number of samples to select.
+            selection_result_tag_name: The tag name to use for the selection result.
+            embedding_model_name:
+                The name of the embedding model to use.
+                If None, assert that there is only one embedding model and uses it.
+            sample_filter: An optional filter to apply to the samples.
+        """
+        strategy = EmbeddingDiversityStrategy(embedding_model_name=embedding_model_name)
+        selection_config = SelectionConfig(
+            dataset_id=self.dataset_id,
+            n_samples_to_select=n_samples_to_select,
+            selection_result_tag_name=selection_result_tag_name,
+            sample_filter=sample_filter,
+            strategies=[strategy],
+        )
+        select_via_database(session=self.session, config=selection_config)
+
+    def multi_strategies(
+        self,
+        n_samples_to_select: int,
+        selection_result_tag_name: str,
+        selection_strategies: list[SelectionStrategy],
+        sample_filter: SampleFilter | None = None,
+    ) -> None:
+        """Select a subset of the dataset based on multiple selection strategies.
+
+        Args:
+            n_samples_to_select: The number of samples to select.
+            selection_result_tag_name: The tag name to use for the selection result.
+            selection_strategies:
+                Selection strategies to use for the selection. They can be created after
+                importing them from `lightly_studio.selection.selection_config`.
+            sample_filter: An optional filter to apply to the samples.
+
+        """
+        config = SelectionConfig(
+            dataset_id=self.dataset_id,
+            n_samples_to_select=n_samples_to_select,
+            selection_result_tag_name=selection_result_tag_name,
+            sample_filter=sample_filter,
+            strategies=selection_strategies,
+        )
+        select_via_database(session=self.session, config=config)

lightly_studio/selection/select_via_db.py

@@ -0,0 +1,93 @@
+"""Database selection functions for the selection process."""
+
+from __future__ import annotations
+
+import datetime
+
+from sqlmodel import Session
+
+from lightly_studio.models.tag import TagCreate
+from lightly_studio.resolvers import (
+    embedding_model_resolver,
+    sample_embedding_resolver,
+    sample_resolver,
+    tag_resolver,
+)
+from lightly_studio.selection.mundig import Mundig
+from lightly_studio.selection.selection_config import (
+    EmbeddingDiversityStrategy,
+    SelectionConfig,
+)
+
+
+def select_via_database(session: Session, config: SelectionConfig) -> None:
+    """Runs selection and all database interactions of it.
+
+    First resolves the selection config to actual database values.
+    Then calls Mundig to run the selection with pure values.
+    Last creates a tag for the selected set.
+    """
+    # Check if the tag name is already used
+    existing_tag = tag_resolver.get_by_name(
+        session=session,
+        tag_name=config.selection_result_tag_name,
+        dataset_id=config.dataset_id,
+    )
+    if existing_tag:
+        msg = (
+            f"Tag with name {config.selection_result_tag_name} already exists in the "
+            f"dataset {config.dataset_id}. Please use a different tag name."
+        )
+        raise ValueError(msg)
+
+    # TODO(Malte, 08/2025): Use a DatasetQuery instead of SampleFilter once
+    # the latter is implemented.
+    # See https://linear.app/lightly/issue/LIG-7292/story-python-ui-mvp1-without-datasetquery-and-sample
+    samples = sample_resolver.get_all_by_dataset_id(
+        session,
+        limit=None,
+        dataset_id=config.dataset_id,
+        filters=config.sample_filter,
+    ).samples
+    sample_ids = [s.sample_id for s in samples]
+
+    n_samples_to_select = min(config.n_samples_to_select, len(sample_ids))
+    if n_samples_to_select == 0:
+        print("No samples available for selection.")
+        return
+
+    mundig = Mundig()
+    for strat in config.strategies:
+        if isinstance(strat, EmbeddingDiversityStrategy):
+            embedding_model_id = embedding_model_resolver.get_by_name(
+                session=session,
+                dataset_id=config.dataset_id,
+                embedding_model_name=strat.embedding_model_name,
+            ).embedding_model_id
+            embedding_tables = sample_embedding_resolver.get_by_sample_ids(
+                session=session,
+                sample_ids=sample_ids,
+                embedding_model_id=embedding_model_id,
+            )
+            embeddings = [e.embedding for e in embedding_tables]
+            mundig.add_diversity(embeddings=embeddings, strength=strat.strength)
+        else:
+            raise ValueError(f"Selection strategy of type {type(strat)} is unknown.")
+
+    selected_indices = mundig.run(n_samples=n_samples_to_select)
+    selected_sample_ids = [sample_ids[i] for i in selected_indices]
+
+    datetime_str = datetime.datetime.now(tz=datetime.timezone.utc).isoformat()
+    tag_description = f"Selected at {datetime_str} UTC"
+    tag = tag_resolver.create(
+        session=session,
+        tag=TagCreate(
+            dataset_id=config.dataset_id,
+            name=config.selection_result_tag_name,
+            kind="sample",
+            description=tag_description,
+        ),
+    )
+    tag_resolver.add_sample_ids_to_tag_id(
+        session=session, tag_id=tag.tag_id, sample_ids=selected_sample_ids
+    )

lightly_studio/selection/selection_config.py

@@ -0,0 +1,31 @@
+"""Pydantic models for the Selection configuration."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from pydantic import BaseModel
+
+from lightly_studio.resolvers.samples_filter import SampleFilter
+
+
+class SelectionConfig(BaseModel):
+    """Configuration for the selection process."""
+
+    dataset_id: UUID
+    sample_filter: SampleFilter | None = None
+    n_samples_to_select: int
+    selection_result_tag_name: str
+    strategies: list[SelectionStrategy]
+
+
+class SelectionStrategy(BaseModel):
+    """Base class for selection strategies."""
+
+    strength: float = 1.0
+
+
+class EmbeddingDiversityStrategy(SelectionStrategy):
+    """Selection strategy based on embedding diversity."""
+
+    embedding_model_name: str | None

lightly_studio/services/annotations_service/__init__.py

@@ -0,0 +1,21 @@
+"""Services for annotations operations."""
+
+from lightly_studio.services.annotations_service.get_annotation_by_id import (
+    get_annotation_by_id,
+)
+from lightly_studio.services.annotations_service.update_annotation import (
+    update_annotation,
+)
+from lightly_studio.services.annotations_service.update_annotation_label import (
+    update_annotation_label,
+)
+from lightly_studio.services.annotations_service.update_annotations import (
+    update_annotations,
+)
+
+__all__ = [
+    "get_annotation_by_id",
+    "update_annotation",
+    "update_annotation_label",
+    "update_annotations",
+]

lightly_studio/services/annotations_service/get_annotation_by_id.py

@@ -0,0 +1,31 @@
+"""Get an annotation by its ID."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+)
+from lightly_studio.resolvers import (
+    annotation_resolver,
+)
+
+
+def get_annotation_by_id(session: Session, annotation_id: UUID) -> AnnotationBaseTable:
+    """Retrieve an annotation by its ID.
+
+    Args:
+        session: Database session for executing the operation.
+        annotation_id: ID of the annotation to retrieve.
+
+    Returns:
+        The retrieved annotation.
+    """
+    annotation = annotation_resolver.get_by_id(session=session, annotation_id=annotation_id)
+    if not annotation:
+        raise ValueError(f"Annotation {annotation_id} not found")
+
+    return annotation