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

lightly_studio/__init__.py

@@ -6,6 +6,6 @@ from lightly_studio.core.start_gui import start_gui
 
 # TODO (Jonas 08/25): This will be removed as soon as the new interface is used in the examples
 from lightly_studio.dataset.loader import DatasetLoader
-from lightly_studio.models.annotation_task import AnnotationType
+from lightly_studio.models.annotation.annotation_base import AnnotationType
 
 __all__ = ["AnnotationType", "Dataset", "DatasetLoader", "start_gui"]

lightly_studio/api/app.py

@@ -16,14 +16,16 @@ from lightly_studio.api.routes import healthz, images, webapp
 from lightly_studio.api.routes.api import (
     annotation,
     annotation_label,
-    annotation_task,
+    caption,
     classifier,
     dataset,
     dataset_tag,
+    embeddings2d,
+    export,
     features,
     metadata,
-    metrics,
     sample,
+    selection,
     settings,
     text_embedding,
 )
@@ -84,16 +86,18 @@ api_router = APIRouter(prefix="/api", tags=["api"])
 
 api_router.include_router(dataset.dataset_router)
 api_router.include_router(dataset_tag.tag_router)
+api_router.include_router(export.export_router)
 api_router.include_router(sample.samples_router)
 api_router.include_router(annotation_label.annotations_label_router)
 api_router.include_router(annotation.annotations_router)
+api_router.include_router(caption.captions_router)
 api_router.include_router(text_embedding.text_embedding_router)
-api_router.include_router(annotation_task.router)
 api_router.include_router(settings.settings_router)
 api_router.include_router(classifier.classifier_router)
+api_router.include_router(embeddings2d.embeddings2d_router)
 api_router.include_router(features.features_router)
 api_router.include_router(metadata.metadata_router)
-api_router.include_router(metrics.metrics_router)
+api_router.include_router(selection.selection_router)
 
 
 app.include_router(api_router)

lightly_studio/api/db_tables.py

@@ -6,9 +6,6 @@ from lightly_studio.models.annotation.annotation_base import (
 from lightly_studio.models.annotation_label import (
     AnnotationLabelTable,  # noqa: F401, required for SQLModel to work properly
 )
-from lightly_studio.models.annotation_task import (
-    AnnotationTaskTable,  # noqa: F401, required for SQLModel to work properly
-)
 from lightly_studio.models.dataset import (
     DatasetTable,  # noqa: F401, required for SQLModel to work properly
 )

lightly_studio/api/routes/api/annotation.py

@@ -9,6 +9,7 @@ from fastapi.params import Query
 from pydantic import BaseModel
 from typing_extensions import Annotated
 
+from lightly_studio.api.routes.api import annotations as annotations_module
 from lightly_studio.api.routes.api.dataset import get_and_validate_dataset_id
 from lightly_studio.api.routes.api.status import (
     HTTP_STATUS_CREATED,
@@ -36,6 +37,7 @@ from lightly_studio.services.annotations_service.update_annotation import (
 )
 
 annotations_router = APIRouter(prefix="/datasets/{dataset_id}", tags=["annotations"])
+annotations_router.include_router(annotations_module.create_annotation_router)
 
 
 @annotations_router.get("/annotations/count")
@@ -221,3 +223,27 @@ def remove_tag_from_annotation(
         raise HTTPException(status_code=HTTP_STATUS_NOT_FOUND, detail=f"Tag {tag_id} not found")
 
     return True
+
+
+@annotations_router.delete("/annotations/{annotation_id}")
+def delete_annotation(
+    session: SessionDep,
+    # We need dataset_id because generator doesn't include it
+    # actuall path for this route is /datasets/{dataset_id}/annotations/{annotation_id}
+    dataset_id: Annotated[  # noqa: ARG001
+        UUID,
+        Path(title="Dataset Id", description="The ID of the dataset"),
+    ],
+    annotation_id: Annotated[
+        UUID, Path(title="Annotation ID", description="ID of the annotation to delete")
+    ],
+) -> dict[str, str]:
+    """Delete an annotation from the database."""
+    try:
+        annotations_service.delete_annotation(session=session, annotation_id=annotation_id)
+        return {"status": "deleted"}
+    except ValueError as e:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail="Annotation not found",
+        ) from e

lightly_studio/api/routes/api/annotations/__init__.py

@@ -0,0 +1,7 @@
+from .create_annotation import (
+    create_annotation_router,
+)
+
+__all__ = [
+    "create_annotation_router",
+]

lightly_studio/api/routes/api/annotations/create_annotation.py

@@ -0,0 +1,52 @@
+"""Create annotation route."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from fastapi import APIRouter, Path
+from fastapi.params import Body
+from pydantic import BaseModel
+from typing_extensions import Annotated
+
+from lightly_studio.db_manager import SessionDep
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+    AnnotationType,
+    AnnotationView,
+)
+from lightly_studio.services import annotations_service
+from lightly_studio.services.annotations_service.create_annotation import AnnotationCreateParams
+
+create_annotation_router = APIRouter()
+
+
+class AnnotationCreateInput(BaseModel):
+    """API interface to create annotation."""
+
+    annotation_label_id: UUID
+    annotation_type: AnnotationType
+    sample_id: UUID
+    x: int | None = None
+    y: int | None = None
+    width: int | None = None
+    height: int | None = None
+    segmentation_mask: list[int] | None = None
+
+
+@create_annotation_router.post(
+    "/annotations",
+    response_model=AnnotationView,
+)
+def create_annotation(
+    dataset_id: Annotated[UUID, Path(title="Dataset Id", description="The ID of the dataset")],
+    session: SessionDep,
+    create_annotation_input: Annotated[AnnotationCreateInput, Body()],
+) -> AnnotationBaseTable:
+    """Create a new annotation."""
+    return annotations_service.create_annotation(
+        session=session,
+        annotation=AnnotationCreateParams(
+            dataset_id=dataset_id, **create_annotation_input.model_dump()
+        ),
+    )

lightly_studio/api/routes/api/caption.py

@@ -0,0 +1,30 @@
+"""API routes for dataset captions."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from fastapi import APIRouter, Depends, Path
+from typing_extensions import Annotated
+
+from lightly_studio.api.routes.api.validators import Paginated, PaginatedWithCursor
+from lightly_studio.db_manager import SessionDep
+from lightly_studio.models.caption import CaptionsListView
+from lightly_studio.resolvers import caption_resolver
+from lightly_studio.resolvers.caption_resolver import GetAllCaptionsResult
+
+captions_router = APIRouter(prefix="/datasets/{dataset_id}", tags=["captions"])
+
+
+@captions_router.get("/captions", response_model=CaptionsListView)
+def read_captions(
+    dataset_id: Annotated[UUID, Path(title="Dataset Id")],
+    session: SessionDep,
+    pagination: Annotated[PaginatedWithCursor, Depends()],
+) -> GetAllCaptionsResult:
+    """Retrieve captions for a dataset."""
+    return caption_resolver.get_all(
+        session=session,
+        dataset_id=dataset_id,
+        pagination=Paginated(offset=pagination.offset, limit=pagination.limit),
+    )

lightly_studio/api/routes/api/dataset.py

@@ -108,6 +108,7 @@ def delete_dataset(
     return {"status": "deleted"}
 
 
+# TODO(Michal, 09/2025): Move to export.py
 class ExportBody(BaseModel):
     """body parameters for including or excluding tag_ids or sample_ids."""
 
@@ -123,6 +124,7 @@ class ExportBody(BaseModel):
 # of sample_ids, it is a POST request to avoid URL length limitations.
 # A body with a GET request is supported by fastAPI however it has undefined
 # behavior: https://fastapi.tiangolo.com/tutorial/body/
+# TODO(Michal, 09/2025): Move to export.py
 @dataset_router.post(
     "/datasets/{dataset_id}/export",
 )
@@ -155,11 +157,7 @@ def export_dataset_to_absolute_paths(
     return response
 
 
-"""
-Endpoint to export samples from a dataset.
-"""
-
-
+# TODO(Michal, 09/2025): Move to export.py
 @dataset_router.post(
     "/datasets/{dataset_id}/export/stats",
 )

lightly_studio/api/routes/api/embeddings2d.py

@@ -0,0 +1,136 @@
+"""Routes delivering 2D embeddings for visualization."""
+
+from __future__ import annotations
+
+import io
+from uuid import UUID
+
+import numpy as np
+import pyarrow as pa
+from fastapi import APIRouter, HTTPException, Response
+from numpy.typing import NDArray
+from pyarrow import ipc
+from pydantic import BaseModel, Field
+from sklearn.manifold import TSNE
+from sqlmodel import select
+
+from lightly_studio.db_manager import SessionDep
+from lightly_studio.models.dataset import DatasetTable
+from lightly_studio.models.embedding_model import EmbeddingModelTable
+from lightly_studio.resolvers import sample_embedding_resolver, sample_resolver
+from lightly_studio.resolvers.samples_filter import SampleFilter
+
+embeddings2d_router = APIRouter()
+
+
+class GetEmbeddings2DRequest(BaseModel):
+    """Request body for retrieving 2D embeddings."""
+
+    filters: SampleFilter | None = Field(
+        None,
+        description="Filter parameters identifying matching samples",
+    )
+
+
+@embeddings2d_router.post("/embeddings2d/tsne")
+def get_embeddings2d__tsne(
+    session: SessionDep,
+    body: GetEmbeddings2DRequest | None = None,
+) -> Response:
+    """Return 2D embeddings serialized as an Arrow stream."""
+    # TODO(Malte, 09/2025): Support choosing the dataset via API parameter.
+    dataset = session.exec(select(DatasetTable).limit(1)).first()
+    if dataset is None:
+        raise HTTPException(status_code=404, detail="No dataset configured.")
+
+    # TODO(Malte, 09/2025): Support choosing the embedding model via API parameter.
+    embedding_model = session.exec(
+        select(EmbeddingModelTable)
+        .where(EmbeddingModelTable.dataset_id == dataset.dataset_id)
+        .limit(1)
+    ).first()
+    if embedding_model is None:
+        raise HTTPException(status_code=404, detail="No embedding model configured.")
+
+    embeddings = sample_embedding_resolver.get_all_by_dataset_id(
+        session=session,
+        dataset_id=dataset.dataset_id,
+        embedding_model_id=embedding_model.embedding_model_id,
+    )
+
+    embedding_values = np.asarray([e.embedding for e in embeddings], dtype=np.float32)
+    embedding_values_tsne = _calculate_tsne_embeddings(embedding_values)
+    x = embedding_values_tsne[:, 0]
+    y = embedding_values_tsne[:, 1]
+
+    matching_sample_ids: set[UUID] | None = None
+    filters = body.filters if body else None
+    if filters:
+        matching_samples_result = sample_resolver.get_all_by_dataset_id(
+            session=session,
+            dataset_id=dataset.dataset_id,
+            filters=filters,
+        )
+        matching_sample_ids = {sample.sample_id for sample in matching_samples_result.samples}
+
+    sample_ids = [embedding.sample_id for embedding in embeddings]
+    if matching_sample_ids is None:
+        fulfils_filter = [1] * len(sample_ids)
+    else:
+        fulfils_filter = [1 if sample_id in matching_sample_ids else 0 for sample_id in sample_ids]
+
+    # TODO(Malte, 09/2025): Save the 2D-embeddings in the database to avoid recomputing
+    # them on every request.
+
+    # TODO(Malte, 09/2025): Include a sample identifier in the returned payload.
+    table = pa.table(
+        {
+            "x": pa.array(x, type=pa.float32()),
+            "y": pa.array(y, type=pa.float32()),
+            "fulfils_filter": pa.array(fulfils_filter, type=pa.uint8()),
+            "sample_id": pa.array([str(sample_id) for sample_id in sample_ids], type=pa.string()),
+        }
+    )
+
+    buffer = io.BytesIO()
+    with ipc.new_stream(buffer, table.schema) as writer:
+        writer.write_table(table)
+    buffer.seek(0)
+
+    return Response(
+        content=buffer.getvalue(),
+        media_type="application/vnd.apache.arrow.stream",
+        headers={
+            "Content-Disposition": "inline; filename=embeddings2d.arrow",
+            "Content-Type": "application/vnd.apache.arrow.stream",
+            "X-Content-Type-Options": "nosniff",
+        },
+    )
+
+
+def _calculate_tsne_embeddings(embedding_values: NDArray[np.float32]) -> NDArray[np.float32]:
+    # TODO(Malte, 10/2025): Switch to a better and faster projection method than
+    # scikit-learn's TSNE.
+    # See https://linear.app/lightly/issue/LIG-7678/embedding-plot-investigate-fasterandbetter-2d-computation-options
+    n_samples = embedding_values.shape[0]
+    # For 0, 1 or 2 samples we hard-code deterministic coordinates.
+    if n_samples == 0:
+        return np.zeros((0, 2), dtype=np.float32)
+    if n_samples == 1:
+        return np.asarray([[0.0, 0.0]], dtype=np.float32)
+    if n_samples == 2:  # noqa: PLR2004
+        return np.asarray([[0.0, 0.0], [1.0, 1.0]], dtype=np.float32)
+
+    # Copied from lightly-core:
+    # https://github.com/lightly-ai/lightly-core/blob/b738952516e916eba42fdd28498491ff18df5c1e/appv2/packages/queueworker/src/jobs/embeddings2d/function-source/main.py#L179-L186
+    embeddings_2d: NDArray[np.float32] = TSNE(
+        init="pca",  # changed in https://github.com/scikit-learn/scikit-learn/issues/18018
+        learning_rate="auto",  # changed in https://github.com/scikit-learn/scikit-learn/issues/18018
+        n_components=2,
+        # Perplexity must be _less_ than the number of entries. 30 is the default value.
+        # https://scikit-learn.org/stable/modules/generated/sklearn.manifold.TSNE.html
+        perplexity=min(30.0, float(n_samples - 1)),
+        # Make the computation deterministic.
+        random_state=0,
+    ).fit_transform(embedding_values)
+    return embeddings_2d

lightly_studio/api/routes/api/export.py

@@ -0,0 +1,73 @@
+"""API routes for exporting dataset annotation tasks."""
+
+from __future__ import annotations
+
+from collections.abc import Generator
+from pathlib import Path as PathlibPath
+from tempfile import TemporaryDirectory
+
+from fastapi import APIRouter, Depends, Path
+from fastapi.responses import StreamingResponse
+from typing_extensions import Annotated
+
+from lightly_studio.api.routes.api import dataset as dataset_api
+from lightly_studio.core.dataset_query.dataset_query import DatasetQuery
+from lightly_studio.db_manager import SessionDep
+from lightly_studio.export import export_dataset
+from lightly_studio.models.dataset import DatasetTable
+
+export_router = APIRouter(prefix="/datasets/{dataset_id}", tags=["export"])
+
+
+@export_router.get("/export/annotations")
+def export_dataset_annotations(
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(dataset_api.get_and_validate_dataset_id),
+    ],
+    session: SessionDep,
+) -> StreamingResponse:
+    """Export dataset annotations for an object detection task in COCO format."""
+    # Query to export - all samples in the dataset.
+    dataset_query = DatasetQuery(dataset=dataset, session=session)
+
+    # Create the export in a temporary directory. We cannot use a context manager
+    # because the directory should be deleted only after the file has finished streaming.
+    temp_dir = TemporaryDirectory()
+    output_path = PathlibPath(temp_dir.name) / "coco_export.json"
+
+    try:
+        export_dataset.to_coco_object_detections(
+            session=session,
+            samples=dataset_query,
+            output_json=output_path,
+        )
+    except Exception:
+        temp_dir.cleanup()
+        # Reraise.
+        raise
+
+    return StreamingResponse(
+        content=_stream_export_file(
+            temp_dir=temp_dir,
+            file_path=output_path,
+        ),
+        media_type="application/json",
+        headers={
+            "Access-Control-Expose-Headers": "Content-Disposition",
+            "Content-Disposition": f"attachment; filename={output_path.name}",
+        },
+    )
+
+
+def _stream_export_file(
+    temp_dir: TemporaryDirectory[str],
+    file_path: PathlibPath,
+) -> Generator[bytes, None, None]:
+    """Stream the export file and clean up the temporary directory afterwards."""
+    try:
+        with file_path.open("rb") as file:
+            yield from file
+    finally:
+        temp_dir.cleanup()

lightly_studio/api/routes/api/metadata.py

@@ -5,11 +5,16 @@ from __future__ import annotations
 from typing import List
 from uuid import UUID
 
-from fastapi import APIRouter, Path
+from fastapi import APIRouter, Depends, Path
+from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
+from lightly_studio.api.routes.api.dataset import get_and_validate_dataset_id
 from lightly_studio.db_manager import SessionDep
+from lightly_studio.metadata import compute_typicality
+from lightly_studio.models.dataset import DatasetTable
 from lightly_studio.models.metadata import MetadataInfoView
+from lightly_studio.resolvers import embedding_model_resolver
 from lightly_studio.resolvers.metadata_resolver.sample.get_metadata_info import (
     get_all_metadata_keys_and_schema,
 )
@@ -33,3 +38,54 @@ def get_metadata_info(
         for numerical metadata types.
     """
     return get_all_metadata_keys_and_schema(session=session, dataset_id=dataset_id)
+
+
+class ComputeTypicalityRequest(BaseModel):
+    """Request model for computing typicality metadata."""
+
+    embedding_model_name: str | None = Field(
+        default=None,
+        description="Embedding model name (uses default if not specified)",
+    )
+    metadata_name: str = Field(
+        default="typicality",
+        description="Metadata field name (defaults to 'typicality')",
+    )
+
+
+@metadata_router.post(
+    "/metadata/typicality",
+    status_code=204,
+    response_model=None,
+)
+def compute_typicality_metadata(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Depends(get_and_validate_dataset_id),
+    ],
+    request: ComputeTypicalityRequest,
+) -> None:
+    """Compute typicality metadata for a dataset.
+
+    Args:
+        session: The database session.
+        dataset: The dataset to compute typicality for.
+        request: Request parameters including optional embedding model name
+            and metadata field name.
+
+    Returns:
+        None (204 No Content on success).
+    """
+    embedding_model = embedding_model_resolver.get_by_name(
+        session=session,
+        dataset_id=dataset.dataset_id,
+        embedding_model_name=request.embedding_model_name,
+    )
+
+    compute_typicality.compute_typicality_metadata(
+        session=session,
+        dataset_id=dataset.dataset_id,
+        embedding_model_id=embedding_model.embedding_model_id,
+        metadata_name=request.metadata_name,
+    )

lightly_studio/api/routes/api/selection.py

@@ -0,0 +1,87 @@
+"""This module contains the API routes for managing selections."""
+
+from __future__ import annotations
+
+from typing import Union
+
+from fastapi import APIRouter, Depends, HTTPException
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated
+
+from lightly_studio.api.routes.api.dataset import get_and_validate_dataset_id
+from lightly_studio.db_manager import SessionDep
+from lightly_studio.models.dataset import DatasetTable
+from lightly_studio.resolvers import sample_resolver
+from lightly_studio.selection.select_via_db import select_via_database
+from lightly_studio.selection.selection_config import (
+    EmbeddingDiversityStrategy,
+    MetadataWeightingStrategy,
+    SelectionConfig,
+)
+
+selection_router = APIRouter()
+
+Strategy = Annotated[
+    Union[EmbeddingDiversityStrategy, MetadataWeightingStrategy],
+    Field(discriminator="strategy_name"),
+]
+
+
+class SelectionRequest(BaseModel):
+    """Request model for selection."""
+
+    n_samples_to_select: int = Field(gt=0, description="Number of samples to select")
+    selection_result_tag_name: str = Field(min_length=1, description="Name for the result tag")
+    strategies: list[Strategy]
+
+
+@selection_router.post(
+    "/datasets/{dataset_id}/selection",
+    status_code=204,
+    response_model=None,
+)
+def create_combination_selection(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Depends(get_and_validate_dataset_id),
+    ],
+    request: SelectionRequest,
+) -> None:
+    """Create a combination selection on the dataset.
+
+    This endpoint performs combination selection using embeddings and metadata.
+    The selected samples are tagged with the specified tag name.
+
+    Args:
+        session: Database session dependency.
+        dataset: Dataset to perform selection on.
+        request: Selection parameters including sample count and tag name.
+
+    Returns:
+        None (204 No Content on success).
+
+    Raises:
+        HTTPException: 400 if selection fails due to invalid parameters or other errors.
+    """
+    # Get all samples in dataset as input for selection.
+    all_samples_result = sample_resolver.get_all_by_dataset_id(
+        session=session, dataset_id=dataset.dataset_id
+    )
+    input_sample_ids = [sample.sample_id for sample in all_samples_result.samples]
+    # Validate we have enough samples to select from.
+    if len(input_sample_ids) < request.n_samples_to_select:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Dataset has only {len(input_sample_ids)} samples, "
+            f"cannot select {request.n_samples_to_select}",
+        )
+    # Create SelectionConfig with diversity strategy.
+    config = SelectionConfig(
+        dataset_id=dataset.dataset_id,
+        n_samples_to_select=request.n_samples_to_select,
+        selection_result_tag_name=request.selection_result_tag_name,
+        strategies=request.strategies,
+    )
+    # Perform selection via database.
+    select_via_database(session=session, config=config, input_sample_ids=input_sample_ids)