lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/api/routes/api/dataset.py

@@ -0,0 +1,191 @@
+"""This module contains the API routes for managing datasets."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import List
+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
+from typing_extensions import Annotated
+
+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,
+    DatasetOverviewView,
+    DatasetTable,
+    DatasetView,
+    DatasetViewWithCount,
+)
+from lightly_studio.resolvers import dataset_resolver
+from lightly_studio.resolvers.dataset_resolver.export import ExportFilter
+
+dataset_router = APIRouter()
+
+
+def get_and_validate_dataset_id(
+    session: SessionDep,
+    dataset_id: UUID,
+) -> DatasetTable:
+    """Get and validate the existence of a dataset on a route."""
+    dataset = dataset_resolver.get_by_id(session=session, dataset_id=dataset_id)
+    if not dataset:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail=f""" Dataset with {dataset_id} not found.""",
+        )
+    return dataset
+
+
+@dataset_router.get("/datasets", response_model=List[DatasetView])
+def read_datasets(
+    session: SessionDep,
+    paginated: Annotated[Paginated, Query()],
+) -> list[DatasetTable]:
+    """Retrieve a list of datasets from the database."""
+    return dataset_resolver.get_all(session=session, offset=paginated.offset, limit=paginated.limit)
+
+
+@dataset_router.get("/datasets/{dataset_id}/root_dataset", response_model=DatasetView)
+def read_root_dataset(
+    session: SessionDep,
+    dataset_id: Annotated[UUID, Path(title="Dataset Id")],
+) -> DatasetTable:
+    """Retrieve the root dataset for a given dataset."""
+    return dataset_resolver.get_root_dataset(session=session, dataset_id=dataset_id)
+
+
+@dataset_router.get("/datasets/{dataset_id}/hierarchy", response_model=List[DatasetView])
+def read_dataset_hierarchy(
+    session: SessionDep,
+    dataset_id: Annotated[UUID, Path(title="Root Dataset Id")],
+) -> list[DatasetTable]:
+    """Retrieve the dataset hierarchy from the database, starting with the root node."""
+    return dataset_resolver.get_hierarchy(session=session, root_dataset_id=dataset_id)
+
+
+@dataset_router.get("/datasets/overview", response_model=List[DatasetOverviewView])
+def read_datasets_overview(session: SessionDep) -> list[DatasetOverviewView]:
+    """Retrieve root datasets with metadata for dashboard display."""
+    return dataset_resolver.get_root_datasets_overview(session=session)
+
+
+@dataset_router.get("/datasets/{dataset_id}", response_model=DatasetViewWithCount)
+def read_dataset(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(get_and_validate_dataset_id),
+    ],
+) -> DatasetViewWithCount:
+    """Retrieve a single dataset from the database."""
+    return dataset_resolver.get_dataset_details(session=session, dataset=dataset)
+
+
+@dataset_router.put("/datasets/{dataset_id}")
+def update_dataset(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(get_and_validate_dataset_id),
+    ],
+    dataset_input: DatasetCreate,
+) -> DatasetTable:
+    """Update an existing dataset in the database."""
+    return dataset_resolver.update(
+        session=session,
+        dataset_id=dataset.dataset_id,
+        dataset_data=dataset_input,
+    )
+
+
+@dataset_router.delete("/datasets/{dataset_id}")
+def delete_dataset(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(get_and_validate_dataset_id),
+    ],
+) -> dict[str, str]:
+    """Delete a dataset from the database."""
+    dataset_resolver.delete(session=session, dataset_id=dataset.dataset_id)
+    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."""
+
+    include: ExportFilter | None = Field(
+        None, description="include filter for sample_ids or tag_ids"
+    )
+    exclude: ExportFilter | None = Field(
+        None, description="exclude filter for sample_ids or tag_ids"
+    )
+
+
+# This endpoint should be a GET, however due to the potential huge size
+# 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",
+)
+def export_dataset_to_absolute_paths(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(get_and_validate_dataset_id),
+    ],
+    body: ExportBody,
+) -> PlainTextResponse:
+    """Export dataset from the database."""
+    # export dataset to absolute paths
+    exported = dataset_resolver.export(
+        session=session,
+        dataset_id=dataset.dataset_id,
+        include=body.include,
+        exclude=body.exclude,
+    )
+
+    # Create a response with the exported data
+    response = PlainTextResponse("\n".join(exported))
+
+    # Add the Content-Disposition header to force download
+    filename = f"{dataset.name}_exported_{datetime.now(timezone.utc)}.txt"
+    response.headers["Access-Control-Expose-Headers"] = "Content-Disposition"
+    response.headers["Content-Disposition"] = f"attachment; filename={filename}"
+
+    return response
+
+
+# TODO(Michal, 09/2025): Move to export.py
+@dataset_router.post(
+    "/datasets/{dataset_id}/export/stats",
+)
+def export_dataset_stats(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(get_and_validate_dataset_id),
+    ],
+    body: ExportBody,
+) -> int:
+    """Get statistics about the export query."""
+    return dataset_resolver.get_filtered_samples_count(
+        session=session,
+        dataset_id=dataset.dataset_id,
+        include=body.include,
+        exclude=body.exclude,
+    )

lightly_studio/api/routes/api/dataset_tag.py

@@ -0,0 +1,266 @@
+"""This module contains the API routes for managing tags."""
+
+from __future__ import annotations
+
+from typing import List
+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
+from typing_extensions import Annotated
+
+from lightly_studio.api.routes.api.dataset import get_and_validate_dataset_id
+from lightly_studio.api.routes.api.status import (
+    HTTP_STATUS_CONFLICT,
+    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.tag import (
+    TagCreate,
+    TagCreateBody,
+    TagTable,
+    TagUpdate,
+    TagUpdateBody,
+    TagView,
+)
+from lightly_studio.resolvers import tag_resolver
+
+tag_router = APIRouter()
+
+
+@tag_router.post(
+    "/datasets/{dataset_id}/tags",
+    response_model=TagView,
+    status_code=HTTP_STATUS_CREATED,
+)
+def create_tag(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(get_and_validate_dataset_id),
+    ],
+    body: TagCreateBody,
+) -> TagTable:
+    """Create a new tag in the database."""
+    dataset_id = dataset.dataset_id
+    try:
+        return tag_resolver.create(
+            session=session,
+            tag=TagCreate(**body.model_dump(exclude_unset=True), dataset_id=dataset_id),
+        )
+    except IntegrityError as e:
+        raise HTTPException(
+            status_code=HTTP_STATUS_CONFLICT,
+            detail=f"""
+                Tag with name {body.name} already exists
+                in the dataset {dataset_id}.
+            """,
+        ) from e
+
+
+@tag_router.get("/datasets/{dataset_id}/tags", response_model=List[TagView])
+def read_tags(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(get_and_validate_dataset_id),
+    ],
+    paginated: Annotated[Paginated, Query()],
+) -> list[TagTable]:
+    """Retrieve a list of tags from the database."""
+    return tag_resolver.get_all_by_dataset_id(
+        session=session,
+        dataset_id=dataset.dataset_id,
+        offset=paginated.offset,
+        limit=paginated.limit,
+    )
+
+
+@tag_router.get("/datasets/{dataset_id}/tags/{tag_id}")
+def read_tag(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(get_and_validate_dataset_id),
+    ],
+    tag_id: Annotated[UUID, Path(title="Tag Id")],
+) -> TagTable:
+    """Retrieve a single tag from the database."""
+    tag = tag_resolver.get_by_id(session=session, tag_id=tag_id)
+    if not tag:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail=f"""
+            Tag with id {tag_id} for dataset {dataset.dataset_id} not found.
+            """,
+        )
+    return tag
+
+
+@tag_router.put("/datasets/{dataset_id}/tags/{tag_id}")
+def update_tag(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(get_and_validate_dataset_id),
+    ],
+    tag_id: Annotated[UUID, Path(title="Tag Id")],
+    body: TagUpdateBody,
+) -> TagTable:
+    """Update an existing tag in the database."""
+    try:
+        tag = tag_resolver.update(
+            session=session,
+            tag_id=tag_id,
+            tag_data=TagUpdate(
+                **body.model_dump(exclude_unset=True),
+            ),
+        )
+        if not tag:
+            raise HTTPException(
+                status_code=HTTP_STATUS_NOT_FOUND,
+                detail=f"Tag with id {tag_id} not found.",
+            )
+    except IntegrityError as e:
+        raise HTTPException(
+            status_code=HTTP_STATUS_CONFLICT,
+            detail=f"""
+                Cannot update tag. Tag with name {body.name}
+                already exists in the dataset {dataset.dataset_id}.
+            """,
+        ) from e
+    return tag
+
+
+@tag_router.delete("/datasets/{dataset_id}/tags/{tag_id}")
+def delete_tag(
+    session: SessionDep,
+    tag_id: Annotated[UUID, Path(title="Tag Id")],
+) -> dict[str, str]:
+    """Delete a tag from the database."""
+    if not tag_resolver.delete(session=session, tag_id=tag_id):
+        raise HTTPException(status_code=HTTP_STATUS_NOT_FOUND, detail="tag not found")
+    return {"status": "deleted"}
+
+
+class SampleIdsBody(BaseModel):
+    """body parameters for adding or removing thing_ids."""
+
+    sample_ids: list[UUID] | None = Field(None, description="sample ids to add/remove")
+
+
+@tag_router.post(
+    "/datasets/{dataset_id}/tags/{tag_id}/add/samples",
+    status_code=HTTP_STATUS_CREATED,
+)
+def add_sample_ids_to_tag_id(
+    session: SessionDep,
+    # dataset_id is needed for the generator
+    dataset_id: Annotated[  # noqa: ARG001
+        UUID,
+        Path(title="Dataset Id", description="The ID of the dataset"),
+    ],
+    tag_id: UUID,
+    body: SampleIdsBody,
+) -> bool:
+    """Add sample_ids to a tag_id."""
+    tag = tag_resolver.get_by_id(session=session, tag_id=tag_id)
+    if not tag:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail=f"Tag {tag_id} not found, can't add sample_ids.",
+        )
+
+    sample_ids = body.sample_ids if body.sample_ids else []
+    tag_resolver.add_sample_ids_to_tag_id(session=session, tag_id=tag_id, sample_ids=sample_ids)
+    return True
+
+
+@tag_router.delete(
+    "/datasets/{dataset_id}/tags/{tag_id}/remove/samples",
+)
+def remove_thing_ids_to_tag_id(
+    session: SessionDep,
+    tag_id: UUID,
+    body: SampleIdsBody,
+) -> bool:
+    """Add thing_ids to a tag_id."""
+    tag = tag_resolver.get_by_id(session=session, tag_id=tag_id)
+    if not tag:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail=f"Tag {tag_id} not found, can't remove samples.",
+        )
+
+    sample_ids = body.sample_ids if body.sample_ids else []
+    tag_resolver.remove_sample_ids_from_tag_id(
+        session=session, tag_id=tag_id, sample_ids=sample_ids
+    )
+    return True
+
+
+class AnnotationIdsBody(BaseModel):
+    """body parameters for adding or removing annotation_ids."""
+
+    annotation_ids: list[UUID] | None = Field(None, description="annotation ids to add/remove")
+
+
+@tag_router.post(
+    "/datasets/{dataset_id}/tags/{tag_id}/add/annotations",
+    status_code=HTTP_STATUS_CREATED,
+)
+def add_annotation_ids_to_tag_id(
+    session: SessionDep,
+    # dataset_id is needed for the generator
+    dataset_id: Annotated[  # noqa: ARG001
+        UUID,
+        Path(title="Dataset Id", description="The ID of the dataset"),
+    ],
+    tag_id: UUID,
+    body: AnnotationIdsBody,
+) -> bool:
+    """Add thing_ids to a tag_id."""
+    tag = tag_resolver.get_by_id(session=session, tag_id=tag_id)
+    if not tag:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail=f"Tag {tag_id} not found, can't add annotations.",
+        )
+
+    annotation_ids = body.annotation_ids if body.annotation_ids else []
+    tag_resolver.add_annotation_ids_to_tag_id(
+        session=session, tag_id=tag_id, annotation_ids=annotation_ids
+    )
+    return True
+
+
+@tag_router.delete(
+    "/datasets/{dataset_id}/tags/{tag_id}/remove/annotations",
+)
+def remove_annotation_ids_to_tag_id(
+    session: SessionDep,
+    tag_id: UUID,
+    body: AnnotationIdsBody,
+) -> bool:
+    """Add thing_ids to a tag_id."""
+    tag = tag_resolver.get_by_id(session=session, tag_id=tag_id)
+    if not tag:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail=f"Tag {tag_id} not found, can't remove annotations.",
+        )
+
+    annotation_ids = body.annotation_ids if body.annotation_ids else []
+    tag_resolver.remove_annotation_ids_from_tag_id(
+        session=session, tag_id=tag_id, annotation_ids=annotation_ids
+    )
+    return True

lightly_studio/api/routes/api/embeddings2d.py

@@ -0,0 +1,90 @@
+"""Routes delivering 2D embeddings for visualization."""
+
+from __future__ import annotations
+
+import io
+from uuid import UUID
+
+import pyarrow as pa
+from fastapi import APIRouter, Response
+from pyarrow import ipc
+from pydantic import BaseModel, Field
+from sqlmodel import select
+
+from lightly_studio.db_manager import SessionDep
+from lightly_studio.models.embedding_model import EmbeddingModelTable
+from lightly_studio.resolvers import image_resolver, twodim_embedding_resolver
+from lightly_studio.resolvers.image_filter import ImageFilter
+
+embeddings2d_router = APIRouter()
+
+
+class GetEmbeddings2DRequest(BaseModel):
+    """Request body for retrieving 2D embeddings."""
+
+    filters: ImageFilter = Field(description="Filter parameters identifying matching samples")
+
+
+@embeddings2d_router.post("/embeddings2d/default")
+def get_2d_embeddings(
+    session: SessionDep,
+    body: GetEmbeddings2DRequest,
+) -> Response:
+    """Return 2D embeddings serialized as an Arrow stream."""
+    dataset_id = (
+        body.filters.sample_filter.dataset_id if body.filters.sample_filter is not None else None
+    )
+    if dataset_id is None:
+        raise ValueError("Dataset ID must be provided in filters.")
+
+    # TODO(Malte, 09/2025): Support choosing the embedding model via API parameter.
+    embedding_model = session.exec(
+        select(EmbeddingModelTable).where(EmbeddingModelTable.dataset_id == dataset_id).limit(1)
+    ).first()
+    if embedding_model is None:
+        raise ValueError("No embedding model configured.")
+
+    x_array, y_array, sample_ids = twodim_embedding_resolver.get_twodim_embeddings(
+        session=session,
+        dataset_id=dataset_id,
+        embedding_model_id=embedding_model.embedding_model_id,
+    )
+
+    matching_sample_ids: set[UUID] | None = None
+    filters = body.filters if body else None
+    if filters:
+        matching_samples_result = image_resolver.get_all_by_dataset_id(
+            session=session,
+            dataset_id=dataset_id,
+            filters=filters,
+        )
+        matching_sample_ids = {sample.sample_id for sample in matching_samples_result.samples}
+
+    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]
+
+    table = pa.table(
+        {
+            "x": pa.array(x_array, type=pa.float32()),
+            "y": pa.array(y_array, 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",
+        },
+    )

lightly_studio/api/routes/api/exceptions.py

@@ -0,0 +1,114 @@
+"""This module contains the FastAPI app configuration."""
+
+from __future__ import annotations
+
+import logging
+
+from fastapi import FastAPI, Request
+from fastapi.exceptions import RequestValidationError, ResponseValidationError
+from fastapi.responses import JSONResponse
+from sqlalchemy.exc import DataError, IntegrityError, OperationalError
+
+from lightly_studio.api.routes.api.status import (
+    HTTP_STATUS_BAD_REQUEST,
+    HTTP_STATUS_CONFLICT,
+    HTTP_STATUS_INTERNAL_SERVER_ERROR,
+    HTTP_STATUS_UNPROCESSABLE_ENTITY,
+)
+
+# Set up logger for error handling
+logger = logging.getLogger("lightly_studio.api.exceptions")
+
+
+def _log_error_details(
+    exc: Exception,
+    status_code: int,
+) -> None:
+    """Log detailed error information with request context."""
+    # Log the error with different levels based on status code
+    logger.error(f"Server Error {status_code}: {exc}")
+
+
+def register_exception_handlers(app: FastAPI) -> None:
+    """Register exception handlers for the FastAPI app."""
+
+    @app.exception_handler(IntegrityError)
+    async def _integrity_error_handler(_request: Request, _exc: IntegrityError) -> JSONResponse:
+        """Handle database integrity errors."""
+        _log_error_details(
+            exc=_exc,
+            status_code=HTTP_STATUS_CONFLICT,
+        )
+
+        return JSONResponse(
+            status_code=HTTP_STATUS_CONFLICT,
+            content={"error": _exc.statement or "Database constraint violated."},
+        )
+
+    @app.exception_handler(OperationalError)
+    async def _operational_error_handler(_request: Request, _exc: OperationalError) -> JSONResponse:
+        """Handle database operational errors."""
+        _log_error_details(
+            exc=_exc,
+            status_code=HTTP_STATUS_INTERNAL_SERVER_ERROR,
+        )
+        return JSONResponse(
+            status_code=HTTP_STATUS_INTERNAL_SERVER_ERROR,
+            content={"error": _exc.statement or "Database operation failed."},
+        )
+
+    @app.exception_handler(DataError)
+    async def _data_error_handler(_request: Request, _exc: DataError) -> JSONResponse:
+        """Handle database data errors."""
+        _log_error_details(
+            exc=_exc,
+            status_code=HTTP_STATUS_BAD_REQUEST,
+        )
+        return JSONResponse(
+            status_code=HTTP_STATUS_BAD_REQUEST,
+            content={"error": _exc.statement or "Invalid response."},
+        )
+
+    @app.exception_handler(ResponseValidationError)
+    async def _data_validation_error_handler(
+        _request: Request, _exc: ResponseValidationError
+    ) -> JSONResponse:
+        """Handle database data errors."""
+        error_details = _exc.errors()
+        if error_details:
+            detail = error_details[0].get("msg", "Invalid data provided.")
+        else:
+            detail = "Invalid data provided."
+
+        _log_error_details(
+            exc=_exc,
+            status_code=HTTP_STATUS_BAD_REQUEST,
+        )
+
+        return JSONResponse(status_code=HTTP_STATUS_BAD_REQUEST, content={"error": detail})
+
+    @app.exception_handler(ValueError)
+    async def _value_error_handler(_request: Request, _exc: ValueError) -> JSONResponse:
+        """Handle value errors."""
+        _log_error_details(
+            exc=_exc,
+            status_code=HTTP_STATUS_INTERNAL_SERVER_ERROR,
+        )
+        return JSONResponse(status_code=HTTP_STATUS_BAD_REQUEST, content={"error": str(_exc)})
+
+    @app.exception_handler(RequestValidationError)
+    async def _request_validation_error_handler(
+        request: Request, exc: RequestValidationError
+    ) -> JSONResponse:
+        body = (await request.body()).decode("utf-8", errors="replace")
+        logger.warning(
+            "Request validation error on %s?%s | errors=%s | body=%s",
+            request.url.path,
+            request.url.query,
+            exc.errors(),
+            body[:500],  # don't log huge bodies
+        )
+        return JSONResponse(
+            status_code=HTTP_STATUS_UNPROCESSABLE_ENTITY,
+            content={"detail": exc.errors()},
+        )