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

lightly_studio/__init__.py

@@ -1,11 +1,11 @@
 # Set up logging before importing any other modules.
 # Add noqa to silence unused import and unsorted imports linter warnings.
 from . import setup_logging  # noqa: F401 I001
-
-# 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.core.dataset import Dataset
 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.annotation_base import AnnotationType
 
-__all__ = ["Dataset", "DatasetLoader", "start_gui"]
+__all__ = ["AnnotationType", "Dataset", "DatasetLoader", "start_gui"]

lightly_studio/api/app.py

@@ -11,19 +11,20 @@ from fastapi.routing import APIRoute
 from sqlmodel import Session
 from typing_extensions import Annotated
 
-from lightly_studio.api.db import db_manager
+from lightly_studio import db_manager
 from lightly_studio.api.routes import healthz, images, webapp
 from lightly_studio.api.routes.api import (
     annotation,
     annotation_label,
-    annotation_task,
     classifier,
     dataset,
     dataset_tag,
+    embeddings2d,
+    export,
     features,
     metadata,
-    metrics,
     sample,
+    selection,
     settings,
     text_embedding,
 )
@@ -84,16 +85,17 @@ 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(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

@@ -7,16 +7,16 @@ from uuid import UUID
 from fastapi import APIRouter, Body, Depends, HTTPException, Path
 from fastapi.params import Query
 from pydantic import BaseModel
-from sqlmodel import Session
 from typing_extensions import Annotated
 
-from lightly_studio.api.db import get_session
+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,
     HTTP_STATUS_NOT_FOUND,
 )
 from lightly_studio.api.routes.api.validators import Paginated, PaginatedWithCursor
+from lightly_studio.db_manager import SessionDep
 from lightly_studio.models.annotation.annotation_base import (
     AnnotationBaseTable,
     AnnotationDetailsView,
@@ -27,6 +27,7 @@ from lightly_studio.resolvers import annotation_resolver, tag_resolver
 from lightly_studio.resolvers.annotation_resolver.get_all import (
     GetAllAnnotationsResult,
 )
+from lightly_studio.resolvers.annotation_resolver.update_bounding_box import BoundingBoxCoordinates
 from lightly_studio.resolvers.annotations.annotations_filter import (
     AnnotationsFilter,
 )
@@ -36,7 +37,7 @@ from lightly_studio.services.annotations_service.update_annotation import (
 )
 
 annotations_router = APIRouter(prefix="/datasets/{dataset_id}", tags=["annotations"])
-SessionDep = Annotated[Session, Depends(get_session)]
+annotations_router.include_router(annotations_module.create_annotation_router)
 
 
 @annotations_router.get("/annotations/count")
@@ -134,11 +135,8 @@ class AnnotationUpdateInput(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
 
 
 @annotations_router.put("/annotations/{annotation_id}")
@@ -161,10 +159,7 @@ def update_annotation(
             annotation_id=annotation_id,
             dataset_id=dataset_id,
             label_name=annotation_update_input.label_name,
-            x=annotation_update_input.x,
-            y=annotation_update_input.y,
-            width=annotation_update_input.width,
-            height=annotation_update_input.height,
+            bounding_box=annotation_update_input.bounding_box,
         ),
     )
 
@@ -188,10 +183,7 @@ def update_annotations(
                 annotation_id=annotation_update_input.annotation_id,
                 dataset_id=dataset_id,
                 label_name=annotation_update_input.label_name,
-                x=annotation_update_input.x,
-                y=annotation_update_input.y,
-                width=annotation_update_input.width,
-                height=annotation_update_input.height,
+                bounding_box=annotation_update_input.bounding_box,
             )
             for annotation_update_input in annotation_update_inputs
         ],
@@ -231,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/annotation_label.py

@@ -4,15 +4,13 @@ from __future__ import annotations
 
 from uuid import UUID
 
-from fastapi import APIRouter, Depends, HTTPException
-from sqlmodel import Session
-from typing_extensions import Annotated
+from fastapi import APIRouter, HTTPException
 
-from lightly_studio.api.db import get_session
 from lightly_studio.api.routes.api.status import (
     HTTP_STATUS_CREATED,
     HTTP_STATUS_NOT_FOUND,
 )
+from lightly_studio.db_manager import SessionDep
 from lightly_studio.models.annotation_label import (
     AnnotationLabelCreate,
     AnnotationLabelTable,
@@ -20,7 +18,6 @@ from lightly_studio.models.annotation_label import (
 from lightly_studio.resolvers import annotation_label_resolver
 
 annotations_label_router = APIRouter()
-SessionDep = Annotated[Session, Depends(get_session)]
 
 
 @annotations_label_router.post(

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/classifier.py

@@ -6,13 +6,11 @@ import io
 from pathlib import Path
 from uuid import UUID
 
-from fastapi import APIRouter, Depends, UploadFile
+from fastapi import APIRouter, UploadFile
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
-from sqlmodel import Session
-from typing_extensions import Annotated
 
-from lightly_studio.api.db import get_session
+from lightly_studio.db_manager import SessionDep
 from lightly_studio.few_shot_classifier.classifier import (
     ExportType,
 )
@@ -22,7 +20,6 @@ from lightly_studio.few_shot_classifier.classifier_manager import (
 from lightly_studio.models.classifier import EmbeddingClassifier
 
 classifier_router = APIRouter()
-SessionDep = Annotated[Session, Depends(get_session)]
 
 
 class GetNegativeSamplesRequest(BaseModel):

lightly_studio/api/routes/api/dataset.py

@@ -9,12 +9,12 @@ from uuid import UUID
 from fastapi import APIRouter, Depends, HTTPException, Path, Query
 from fastapi.responses import PlainTextResponse
 from pydantic import BaseModel
-from sqlmodel import Field, Session
+from sqlmodel import Field
 from typing_extensions import Annotated
 
-from lightly_studio.api.db import get_session
 from lightly_studio.api.routes.api.status import HTTP_STATUS_NOT_FOUND
 from lightly_studio.api.routes.api.validators import Paginated
+from lightly_studio.db_manager import SessionDep
 from lightly_studio.models.dataset import (
     DatasetCreate,
     DatasetTable,
@@ -26,7 +26,6 @@ from lightly_studio.resolvers.dataset_resolver import (
 )
 
 dataset_router = APIRouter()
-SessionDep = Annotated[Session, Depends(get_session)]
 
 
 def get_and_validate_dataset_id(
@@ -109,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."""
 
@@ -124,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",
 )
@@ -156,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/dataset_tag.py

@@ -8,10 +8,9 @@ from uuid import UUID
 from fastapi import APIRouter, Depends, HTTPException, Path, Query
 from pydantic import BaseModel
 from sqlalchemy.exc import IntegrityError
-from sqlmodel import Field, Session
+from sqlmodel import Field
 from typing_extensions import Annotated
 
-from lightly_studio.api.db import get_session
 from lightly_studio.api.routes.api.dataset import get_and_validate_dataset_id
 from lightly_studio.api.routes.api.status import (
     HTTP_STATUS_CONFLICT,
@@ -19,6 +18,7 @@ from lightly_studio.api.routes.api.status import (
     HTTP_STATUS_NOT_FOUND,
 )
 from lightly_studio.api.routes.api.validators import Paginated
+from lightly_studio.db_manager import SessionDep
 from lightly_studio.models.dataset import DatasetTable
 from lightly_studio.models.tag import (
     TagCreate,
@@ -31,7 +31,6 @@ from lightly_studio.models.tag import (
 from lightly_studio.resolvers import tag_resolver
 
 tag_router = APIRouter()
-SessionDep = Annotated[Session, Depends(get_session)]
 
 
 @tag_router.post(

lightly_studio/api/routes/api/embeddings2d.py

@@ -0,0 +1,104 @@
+"""Routes delivering 2D embeddings for visualization."""
+
+from __future__ import annotations
+
+import io
+
+import numpy as np
+import pyarrow as pa
+from fastapi import APIRouter, HTTPException, Response
+from numpy.typing import NDArray
+from pyarrow import ipc
+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
+
+embeddings2d_router = APIRouter()
+
+
+@embeddings2d_router.get("/embeddings2d/tsne")
+def get_embeddings2d__tsne(session: SessionDep) -> 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.")
+
+    # TODO(Malte, 09/2025): Support choosing a subset of samples via API parameter.
+    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]
+
+    # 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()),
+        }
+    )
+
+    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,18 +5,16 @@ from __future__ import annotations
 from typing import List
 from uuid import UUID
 
-from fastapi import APIRouter, Depends, Path
-from sqlmodel import Session
+from fastapi import APIRouter, Path
 from typing_extensions import Annotated
 
-from lightly_studio.api.db import get_session
+from lightly_studio.db_manager import SessionDep
 from lightly_studio.models.metadata import MetadataInfoView
 from lightly_studio.resolvers.metadata_resolver.sample.get_metadata_info import (
     get_all_metadata_keys_and_schema,
 )
 
 metadata_router = APIRouter(prefix="/datasets/{dataset_id}", tags=["metadata"])
-SessionDep = Annotated[Session, Depends(get_session)]
 
 
 @metadata_router.get("/metadata/info", response_model=List[MetadataInfoView])

lightly_studio/api/routes/api/sample.py

@@ -6,16 +6,15 @@ from uuid import UUID
 
 from fastapi import APIRouter, Depends, HTTPException, Path, Query
 from pydantic import BaseModel, Field
-from sqlmodel import Session
 from typing_extensions import Annotated
 
-from lightly_studio.api.db import get_session
 from lightly_studio.api.routes.api.dataset import get_and_validate_dataset_id
 from lightly_studio.api.routes.api.status import (
     HTTP_STATUS_CREATED,
     HTTP_STATUS_NOT_FOUND,
 )
 from lightly_studio.api.routes.api.validators import Paginated
+from lightly_studio.db_manager import SessionDep
 from lightly_studio.models.dataset import DatasetTable
 from lightly_studio.models.sample import (
     SampleCreate,
@@ -27,12 +26,12 @@ from lightly_studio.resolvers import (
     sample_resolver,
     tag_resolver,
 )
+from lightly_studio.resolvers.sample_resolver import GetAllSamplesByDatasetIdResult
 from lightly_studio.resolvers.samples_filter import (
     SampleFilter,
 )
 
 samples_router = APIRouter(prefix="/datasets/{dataset_id}", tags=["samples"])
-SessionDep = Annotated[Session, Depends(get_session)]
 
 
 @samples_router.post("/samples", response_model=SampleView)
@@ -60,7 +59,7 @@ def read_samples(
     session: SessionDep,
     dataset_id: Annotated[UUID, Path(title="Dataset Id")],
     body: ReadSamplesRequest,
-) -> SampleViewsWithCount:
+) -> GetAllSamplesByDatasetIdResult:
     """Retrieve a list of samples from the database with optional filtering.
 
     Args:
@@ -71,22 +70,15 @@ def read_samples(
     Returns:
         A list of filtered samples.
     """
-    result = sample_resolver.get_all_by_dataset_id(
+    return sample_resolver.get_all_by_dataset_id(
         session=session,
         dataset_id=dataset_id,
-        # TODO(Michal, 06/2025): Pass the Paginated object directly.
-        offset=body.pagination.offset if body.pagination else 0,
-        limit=body.pagination.limit if body.pagination else 10,
+        pagination=body.pagination,
         filters=body.filters,
         text_embedding=body.text_embedding,
         sample_ids=body.sample_ids,
     )
 
-    return SampleViewsWithCount(
-        data=result.samples,
-        total_count=result.total_count,
-    )
-
 
 @samples_router.get("/samples/dimensions")
 def get_sample_dimensions(