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

lightly_studio/resolvers/embedding_model_resolver.py

@@ -21,6 +21,28 @@ def create(session: Session, embedding_model: EmbeddingModelCreate) -> Embedding
     return db_embedding_model
 
 
+def get_or_create(session: Session, embedding_model: EmbeddingModelCreate) -> EmbeddingModelTable:
+    """Retrieve an existing EmbeddingModel by hash or create a new one if it does not exist."""
+    db_model = get_by_model_hash(
+        session=session, embedding_model_hash=embedding_model.embedding_model_hash
+    )
+    if db_model is None:
+        return create(session=session, embedding_model=embedding_model)
+
+    # Validate that the existing model matches the provided data.
+    if (
+        db_model.name != embedding_model.name
+        or db_model.parameter_count_in_mb != embedding_model.parameter_count_in_mb
+        or db_model.embedding_dimension != embedding_model.embedding_dimension
+        # TODO(Michal, 09/2025): Allow same model for different datasets.
+        or db_model.dataset_id != embedding_model.dataset_id
+    ):
+        raise ValueError(
+            "An embedding model with the same hash but different parameters already exists."
+        )
+    return db_model
+
+
 def get_all_by_dataset_id(session: Session, dataset_id: UUID) -> list[EmbeddingModelTable]:
     """Retrieve all embedding models."""
     embedding_models = session.exec(

lightly_studio/resolvers/sample_resolver.py

@@ -7,9 +7,11 @@ from datetime import datetime, timezone
 from uuid import UUID
 
 from pydantic import BaseModel
+from sqlalchemy.orm import joinedload, selectinload
 from sqlmodel import Session, col, func, select
 from sqlmodel.sql.expression import Select
 
+from lightly_studio.api.routes.api.validators import Paginated
 from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable
 from lightly_studio.models.annotation_label import AnnotationLabelTable
 from lightly_studio.models.embedding_model import EmbeddingModelTable
@@ -36,6 +38,22 @@ def create_many(session: Session, samples: list[SampleCreate]) -> list[SampleTab
     return db_samples
 
 
+def filter_new_paths(session: Session, file_paths_abs: list[str]) -> tuple[list[str], list[str]]:
+    """Return a) file_path_abs that do not already exist in the database and b) those that do."""
+    existing_file_paths_abs = set(
+        session.exec(
+            select(col(SampleTable.file_path_abs)).where(
+                col(SampleTable.file_path_abs).in_(file_paths_abs)
+            )
+        ).all()
+    )
+    file_paths_abs_set = set(file_paths_abs)
+    return (
+        list(file_paths_abs_set - existing_file_paths_abs),  # paths that are not in the DB
+        list(file_paths_abs_set & existing_file_paths_abs),  # paths that are already in the DB
+    )
+
+
 def get_by_id(session: Session, dataset_id: UUID, sample_id: UUID) -> SampleTable | None:
     """Retrieve a single sample by ID."""
     return session.exec(
@@ -45,6 +63,13 @@ def get_by_id(session: Session, dataset_id: UUID, sample_id: UUID) -> SampleTabl
     ).one_or_none()
 
 
+def count_by_dataset_id(session: Session, dataset_id: UUID) -> int:
+    """Count the number of samples in a dataset."""
+    return session.exec(
+        select(func.count()).select_from(SampleTable).where(SampleTable.dataset_id == dataset_id)
+    ).one()
+
+
 def get_many_by_id(session: Session, sample_ids: list[UUID]) -> list[SampleTable]:
     """Retrieve multiple samples by their IDs.
 
@@ -63,19 +88,33 @@ class GetAllSamplesByDatasetIdResult(BaseModel):
 
     samples: Sequence[SampleTable]
     total_count: int
+    next_cursor: int | None = None
 
 
 def get_all_by_dataset_id(  # noqa: PLR0913
     session: Session,
     dataset_id: UUID,
-    offset: int = 0,
-    limit: int | None = None,
+    pagination: Paginated | None = None,
     filters: SampleFilter | None = None,
     text_embedding: list[float] | None = None,
     sample_ids: list[UUID] | None = None,
 ) -> GetAllSamplesByDatasetIdResult:
     """Retrieve samples for a specific dataset with optional filtering."""
-    samples_query = select(SampleTable).where(SampleTable.dataset_id == dataset_id)
+    samples_query = (
+        select(SampleTable)
+        .options(
+            selectinload(SampleTable.annotations).options(
+                joinedload(AnnotationBaseTable.annotation_label),
+                joinedload(AnnotationBaseTable.object_detection_details),
+                joinedload(AnnotationBaseTable.instance_segmentation_details),
+                joinedload(AnnotationBaseTable.semantic_segmentation_details),
+            ),
+            selectinload(SampleTable.tags),
+            # Ignore type checker error below as it's a false positive caused by TYPE_CHECKING.
+            joinedload(SampleTable.metadata_dict),  # type: ignore[arg-type]
+        )
+        .where(SampleTable.dataset_id == dataset_id)
+    )
     total_count_query = (
         select(func.count()).select_from(SampleTable).where(SampleTable.dataset_id == dataset_id)
     )
@@ -120,15 +159,20 @@ def get_all_by_dataset_id(  # noqa: PLR0913
             col(SampleTable.created_at).asc(), col(SampleTable.sample_id).asc()
         )
 
-    # paginate query when offset or limit are set/positive
-    if offset > 0:
-        samples_query = samples_query.offset(offset)
-    if limit is not None:
-        samples_query = samples_query.limit(limit)
+    # Apply pagination if provided
+    if pagination is not None:
+        samples_query = samples_query.offset(pagination.offset).limit(pagination.limit)
+
+    total_count = session.exec(total_count_query).one()
+
+    next_cursor = None
+    if pagination and pagination.offset + pagination.limit < total_count:
+        next_cursor = pagination.offset + pagination.limit
 
     return GetAllSamplesByDatasetIdResult(
         samples=session.exec(samples_query).all(),
-        total_count=session.exec(total_count_query).one(),
+        total_count=total_count,
+        next_cursor=next_cursor,
     )
 
 

lightly_studio/resolvers/tag_resolver.py

@@ -274,3 +274,26 @@ def remove_annotation_ids_from_tag_id(
     session.commit()
     session.refresh(tag)
     return tag
+
+
+def get_or_create_sample_tag_by_name(
+    session: Session,
+    dataset_id: UUID,
+    tag_name: str,
+) -> TagTable:
+    """Get an existing sample tag by name or create a new one if it doesn't exist.
+
+    Args:
+        session: Database session for executing queries.
+        dataset_id: The dataset ID to search/create the tag for.
+        tag_name: Name of the tag to get or create.
+
+    Returns:
+        The existing or newly created sample tag.
+    """
+    existing_tag = get_by_name(session=session, tag_name=tag_name, dataset_id=dataset_id)
+    if existing_tag:
+        return existing_tag
+
+    new_tag = TagCreate(name=tag_name, dataset_id=dataset_id, kind="sample")
+    return create(session=session, tag=new_tag)

lightly_studio/selection/select.py

@@ -1,96 +1,105 @@
-"""Provides the user python interface to selection."""
+"""Provides the user python interface to selection bound to sample ids."""
 
 from __future__ import annotations
 
+from collections.abc import Iterable
+from typing import Final
 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,
+    MetadataWeightingStrategy,
     SelectionConfig,
     SelectionStrategy,
 )
 
 
 class Selection:
-    """User selection interface for the dataset."""
+    """Selection interface for candidate sample ids."""
 
-    # 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.
+    def __init__(
+        self,
+        dataset_id: UUID,
+        session: Session,
+        input_sample_ids: Iterable[UUID],
+    ) -> None:
+        """Create the selection interface.
 
         Args:
-            dataset_id: The ID of the dataset to select from.
-            session: The database session to use for selection.
+            dataset_id: Dataset in which the selection is performed.
+            session: Database session to resolve selection dependencies.
+            input_sample_ids: Candidate sample ids considered for selection.
+                The iterable is consumed immediately to capture a stable snapshot.
+        """
+        self._dataset_id: Final[UUID] = dataset_id
+        self._session: Final[Session] = session
+        self._input_sample_ids: list[UUID] = list(input_sample_ids)
 
+    def metadata_weighting(
+        self,
+        n_samples_to_select: int,
+        selection_result_tag_name: str,
+        metadata_key: str,
+    ) -> None:
+        """Select a subset based on numeric metadata weights.
+
+        Args:
+            n_samples_to_select: Number of samples to select.
+            selection_result_tag_name: Tag name for the selection result.
+            metadata_key: Metadata key used as weights (float or int values).
         """
-        self.dataset_id = dataset_id
-        self.session = session
+        strategy = MetadataWeightingStrategy(metadata_key=metadata_key)
+        self.multi_strategies(
+            n_samples_to_select=n_samples_to_select,
+            selection_result_tag_name=selection_result_tag_name,
+            selection_strategies=[strategy],
+        )
 
     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.
+        """Select a diverse subset using embeddings.
 
         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.
+            n_samples_to_select: Number of samples to select.
+            selection_result_tag_name: Tag name for the selection result.
+            embedding_model_name: Optional embedding model name. If None, uses the only
+                available model or raises if multiple exist.
         """
         strategy = EmbeddingDiversityStrategy(embedding_model_name=embedding_model_name)
-        selection_config = SelectionConfig(
-            dataset_id=self.dataset_id,
+        self.multi_strategies(
             n_samples_to_select=n_samples_to_select,
             selection_result_tag_name=selection_result_tag_name,
-            sample_filter=sample_filter,
-            strategies=[strategy],
+            selection_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.
+        """Select a subset based on multiple 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.
-
+            n_samples_to_select: Number of samples to select.
+            selection_result_tag_name: Tag name for the selection result.
+            selection_strategies: Strategies to compose for selection.
         """
         config = SelectionConfig(
-            dataset_id=self.dataset_id,
+            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)
+        select_via_database(
+            session=self._session,
+            config=config,
+            input_sample_ids=self._input_sample_ids,
+        )

lightly_studio/selection/select_via_db.py

@@ -3,29 +3,33 @@
 from __future__ import annotations
 
 import datetime
+from uuid import UUID
 
 from sqlmodel import Session
 
 from lightly_studio.models.tag import TagCreate
 from lightly_studio.resolvers import (
     embedding_model_resolver,
+    metadata_resolver,
     sample_embedding_resolver,
-    sample_resolver,
     tag_resolver,
 )
 from lightly_studio.selection.mundig import Mundig
 from lightly_studio.selection.selection_config import (
     EmbeddingDiversityStrategy,
+    MetadataWeightingStrategy,
     SelectionConfig,
 )
 
 
-def select_via_database(session: Session, config: SelectionConfig) -> None:
-    """Runs selection and all database interactions of it.
+def select_via_database(
+    session: Session, config: SelectionConfig, input_sample_ids: list[UUID]
+) -> None:
+    """Run selection using the provided candidate sample ids.
 
-    First resolves the selection config to actual database values.
+    First resolves the selection config to concrete database values.
     Then calls Mundig to run the selection with pure values.
-    Last creates a tag for the selected set.
+    Finally creates a tag for the selected set.
     """
     # Check if the tag name is already used
     existing_tag = tag_resolver.get_by_name(
@@ -40,18 +44,7 @@ def select_via_database(session: Session, config: SelectionConfig) -> None:
         )
         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))
+    n_samples_to_select = min(config.n_samples_to_select, len(input_sample_ids))
     if n_samples_to_select == 0:
         print("No samples available for selection.")
         return
@@ -66,16 +59,27 @@ def select_via_database(session: Session, config: SelectionConfig) -> None:
             ).embedding_model_id
             embedding_tables = sample_embedding_resolver.get_by_sample_ids(
                 session=session,
-                sample_ids=sample_ids,
+                sample_ids=input_sample_ids,
                 embedding_model_id=embedding_model_id,
             )
             embeddings = [e.embedding for e in embedding_tables]
             mundig.add_diversity(embeddings=embeddings, strength=strat.strength)
+        elif isinstance(strat, MetadataWeightingStrategy):
+            key = strat.metadata_key
+            weights = []
+            for sample_id in input_sample_ids:
+                weight = metadata_resolver.get_value_for_sample(session, sample_id, key)
+                if not isinstance(weight, (float, int)):
+                    raise ValueError(
+                        f"Metadata {key} is not a number, only numbers can be used as weights"
+                    )
+                weights.append(float(weight))
+            mundig.add_weighting(weights, 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]
+    selected_sample_ids = [input_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"

lightly_studio/selection/selection_config.py

@@ -6,14 +6,11 @@ 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]
@@ -29,3 +26,9 @@ class EmbeddingDiversityStrategy(SelectionStrategy):
     """Selection strategy based on embedding diversity."""
 
     embedding_model_name: str | None
+
+
+class MetadataWeightingStrategy(SelectionStrategy):
+    """Selection strategy based on metadata weighting."""
+
+    metadata_key: str

lightly_studio/services/annotations_service/__init__.py

@@ -6,6 +6,9 @@ from lightly_studio.services.annotations_service.get_annotation_by_id import (
 from lightly_studio.services.annotations_service.update_annotation import (
     update_annotation,
 )
+from lightly_studio.services.annotations_service.update_annotation_bounding_box import (
+    update_annotation_bounding_box,
+)
 from lightly_studio.services.annotations_service.update_annotation_label import (
     update_annotation_label,
 )
@@ -16,6 +19,7 @@ from lightly_studio.services.annotations_service.update_annotations import (
 __all__ = [
     "get_annotation_by_id",
     "update_annotation",
+    "update_annotation_bounding_box",
     "update_annotation_label",
     "update_annotations",
 ]

lightly_studio/services/annotations_service/update_annotation.py

@@ -10,6 +10,7 @@ from sqlmodel import Session
 from lightly_studio.models.annotation.annotation_base import (
     AnnotationBaseTable,
 )
+from lightly_studio.resolvers.annotation_resolver.update_bounding_box import BoundingBoxCoordinates
 from lightly_studio.services import annotations_service
 
 
@@ -18,11 +19,8 @@ class AnnotationUpdate(BaseModel):
 
     annotation_id: UUID
     dataset_id: UUID
-    label_name: str | None
-    x: int | None = None
-    y: int | None = None
-    width: int | None = None
-    height: int | None = None
+    label_name: str | None = None
+    bounding_box: BoundingBoxCoordinates | None = None
 
 
 def update_annotation(session: Session, annotation_update: AnnotationUpdate) -> AnnotationBaseTable:
@@ -36,30 +34,21 @@ def update_annotation(session: Session, annotation_update: AnnotationUpdate) ->
         The updated annotation.
 
     """
-    if annotation_update.label_name is None:
-        raise ValueError("Label name must be provided for updating annotation")
-
-    # comment this out for now so e2e tests will pass
-    # todo: uncomment after passing bbox coordinates on update from frontend
-    # annotation=get_annotation_by_id(session,annotation_update.annotation_id)
-    # if annotation.annotation_type in (
-    #     AnnotationType.OBJECT_DETECTION,
-    #     AnnotationType.INSTANCE_SEGMENTATION,
-    # ) and any(
-    #     [
-    #         annotation_update.x is None,
-    #         annotation_update.y is None,
-    #         annotation_update.width is None,
-    #         annotation_update.height is None,
-    #     ]
-    # ):
-    #     raise ValueError(
-    #         "All bounding box coordinates (x, y, width, height) "
-    #         "must be provided for updating this annotation type"
-    #     )
-
-    return annotations_service.update_annotation_label(
-        session,
-        annotation_update.annotation_id,
-        annotation_update.label_name,
-    )
+    result = None
+    if annotation_update.label_name is not None:
+        result = annotations_service.update_annotation_label(
+            session,
+            annotation_update.annotation_id,
+            annotation_update.label_name,
+        )
+
+    if annotation_update.bounding_box is not None:
+        result = annotations_service.update_annotation_bounding_box(
+            session,
+            annotation_update.annotation_id,
+            bounding_box=annotation_update.bounding_box,
+        )
+
+    if result is None:
+        raise ValueError("No updates provided for the annotation.")
+    return result

lightly_studio/services/annotations_service/update_annotation_bounding_box.py

@@ -0,0 +1,36 @@
+"""Update the bounding box of an annotation."""
+
+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,
+)
+from lightly_studio.resolvers.annotation_resolver.update_bounding_box import BoundingBoxCoordinates
+
+
+def update_annotation_bounding_box(
+    session: Session, annotation_id: UUID, bounding_box: BoundingBoxCoordinates
+) -> AnnotationBaseTable:
+    """Update the bounding box of an annotation.
+
+    Args:
+        session: Database session for executing the operation.
+        annotation_id: UUID of the annotation to update.
+        bounding_box: New bounding box coordinates to assign to the annotation.
+
+    Returns:
+        The updated annotation with the new bounding box assigned.
+
+    """
+    return annotation_resolver.update_bounding_box(
+        session,
+        annotation_id,
+        bounding_box,
+    )