lightly-studio 0.3.1__py3-none-any.whl

lightly_studio/__init__.py

@@ -0,0 +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
+
+
+__all__ = ["Dataset", "DatasetLoader", "start_gui"]

lightly_studio/api/app.py

@@ -0,0 +1,110 @@
+"""This module contains the FastAPI app configuration."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+
+from fastapi import APIRouter, Depends, FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+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.api.routes import healthz, images, webapp
+from lightly_studio.api.routes.api import (
+    annotation,
+    annotation_label,
+    annotation_task,
+    classifier,
+    dataset,
+    dataset_tag,
+    features,
+    metadata,
+    metrics,
+    sample,
+    settings,
+    text_embedding,
+)
+from lightly_studio.api.routes.api.exceptions import (
+    register_exception_handlers,
+)
+from lightly_studio.dataset.env import LIGHTLY_STUDIO_DEBUG
+
+SessionDep = Annotated[Session, Depends(db_manager.session)]
+
+
+@asynccontextmanager
+async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
+    """Lifespan context for initializing and cleaning up resources.
+
+    Args:
+        _: The FastAPI application instance.
+
+    Yields:
+        None when the application is ready.
+    """
+    yield
+
+
+if LIGHTLY_STUDIO_DEBUG:
+    import logging
+
+    logging.basicConfig()
+    logger = logging.getLogger("sqlalchemy.engine")
+    logger.setLevel(logging.DEBUG)
+
+"""Create the FastAPI app."""
+app = FastAPI(lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+def use_route_names_as_operation_ids(app: FastAPI) -> None:
+    """Use API function name for operation IDs.
+
+    Should be called only after all routes have been added.
+    """
+    for route in app.routes:
+        if isinstance(route, APIRoute):
+            route.operation_id = route.name  # in this case, 'read_items'
+
+
+register_exception_handlers(app)
+
+# api routes
+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(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(features.features_router)
+api_router.include_router(metadata.metadata_router)
+api_router.include_router(metrics.metrics_router)
+
+
+app.include_router(api_router)
+# images serving
+app.include_router(images.app_router, prefix="/images")
+
+
+# health status check
+app.include_router(healthz.health_router)
+
+# webapp routes
+app.include_router(webapp.app_router)
+
+use_route_names_as_operation_ids(app)

lightly_studio/api/cache.py

@@ -0,0 +1,77 @@
+"""This module contains the FastAPI cache configuration for static files."""
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+from os import PathLike, stat_result
+
+from fastapi import Response
+from fastapi.staticfiles import StaticFiles
+from starlette.types import Scope
+
+from .routes.api.status import HTTP_STATUS_UNSUPPORTED_MEDIA_TYPE
+
+
+class StaticFilesCache(StaticFiles):
+    """StaticFiles class with cache headers."""
+
+    days_to_expire = 1
+
+    def __init__(  # noqa: PLR0913 (too-many-arguments)
+        self,
+        *,
+        directory: str | PathLike[str] | None = None,
+        packages: list[str | tuple[str, str]] | None = None,
+        html: bool = False,
+        check_dir: bool = True,
+        follow_symlink: bool = False,
+        cachecontrol: str | None = None,
+    ) -> None:
+        """Initialize the StaticFilesCache class."""
+        self.cachecontrol = cachecontrol or f"private, max-age={self.days_to_expire * 24 * 60 * 60}"
+        super().__init__(
+            directory=directory,
+            packages=packages,
+            html=html,
+            check_dir=check_dir,
+            follow_symlink=follow_symlink,
+        )
+
+    def file_response(
+        self,
+        full_path: str | PathLike[str],
+        stat_result: stat_result,
+        scope: Scope,
+        status_code: int = 200,
+    ) -> Response:
+        """Override the file_response method to add cache headers."""
+        allowed_extensions = (
+            # Images
+            ".png",
+            ".jpg",
+            ".jpeg",
+            ".gif",
+            ".webp",
+            ".bmp",
+            ".tiff",
+            # Movies
+            ".mov",
+            ".mp4",
+            ".avi",
+        )
+
+        if not str(full_path).lower().endswith(allowed_extensions):
+            return Response(
+                status_code=HTTP_STATUS_UNSUPPORTED_MEDIA_TYPE
+            )  # Unsupported Media Type
+        resp: Response = super().file_response(full_path, stat_result, scope, status_code)
+        resp.headers.setdefault("Cache-Control", self.cachecontrol)
+
+        # Calculate expiration date
+        expire_date = datetime.now(timezone.utc) + timedelta(days=self.days_to_expire)
+        resp.headers.setdefault("Expires", expire_date.strftime("%a, %d %b %Y %H:%M:%S GMT"))
+
+        # Add Vary header to make sure caches respect the query parameters
+        resp.headers.setdefault("Vary", "Accept-Encoding, Origin, v")
+
+        return resp

lightly_studio/api/db.py

@@ -0,0 +1,133 @@
+"""Module provides functions to initialize and manage the DuckDB."""
+
+from __future__ import annotations
+
+import logging
+import os
+from contextlib import contextmanager
+from typing import Generator
+
+from sqlalchemy.engine import Engine
+from sqlmodel import Session, SQLModel, create_engine
+
+import lightly_studio.api.db_tables  # noqa: F401, required for SQLModel to work properly
+
+
+class MockDatabaseManager:
+    """Mock version of DatabaseManager."""
+
+    _persistent_session: Session | None = None
+
+    def __init__(self) -> None:
+        """Create a new instance of the MockDatabaseManager."""
+        self.engine = create_engine("duckdb:///:memory:")
+        self._session_instance = Session(self.engine, close_resets_only=False)
+        # Initialize tables
+        SQLModel.metadata.create_all(self.engine)
+
+    @contextmanager
+    def session(self) -> Generator[Session, None, None]:
+        """Return the database session."""
+        try:
+            yield self._session_instance
+            if not self._session_instance.in_transaction():
+                self._session_instance.commit()
+        except Exception:
+            self._session_instance.rollback()
+            raise
+        finally:
+            self._session_instance.close()
+
+    def persistent_session(self) -> Session:
+        """Create a persistent session."""
+        if self._persistent_session is None:
+            self._persistent_session = Session(self.engine, close_resets_only=False)
+        return self._persistent_session
+
+
+class DatabaseManager:
+    """Manages database connections and ensures proper resource handling."""
+
+    _instance: DatabaseManager | None = None
+    engine: Engine | None = None
+    _persistent_session: Session | None = None
+
+    @staticmethod
+    def database_exists(db_file: str) -> bool:
+        """Check if database file exists.
+
+        Args:
+            db_file: Path to the database file
+
+        Returns:
+            True if database file exists, False otherwise
+        """
+        return os.path.exists(db_file)
+
+    @staticmethod
+    def cleanup_database(db_file: str) -> None:
+        """Delete database file if it exists.
+
+        Args:
+            db_file: Path to the database file to delete
+        """
+        if DatabaseManager.database_exists(db_file):
+            os.remove(db_file)
+            logging.info(f"Deleted existing database: {db_file}")
+
+    def __new__(
+        cls, db_file: str = "lightly_studio.db", cleanup_existing: bool = False
+    ) -> DatabaseManager:
+        """Create a new instance of the DatabaseManager.
+
+        Args:
+            db_file: Path to the database file
+            cleanup_existing: If True, deletes existing database
+            before creating a new one
+
+        Returns:
+            DatabaseManager instance
+        """
+        if cleanup_existing:
+            cls.cleanup_database(db_file)
+
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            # File-based DuckDB
+            cls._instance.engine = create_engine(f"duckdb:///{db_file}")
+            # Initialize tables
+            SQLModel.metadata.create_all(cls._instance.engine)
+        return cls._instance
+
+    @contextmanager
+    def session(self) -> Generator[Session, None, None]:
+        """Create a new database session."""
+        session = Session(self.engine, close_resets_only=False)
+        try:
+            yield session
+            session.commit()
+        except Exception:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    def persistent_session(self) -> Session:
+        """Create a persistent session."""
+        if self._persistent_session is None:
+            # TODO(Michal, 08/2025): Consider setting close_resets_only=True
+            # before the release. The current solution is more strict and prevents
+            # bugs but it might be too restrictive for users.
+            self._persistent_session = Session(self.engine, close_resets_only=False)
+        return self._persistent_session
+
+
+# Global instance
+db_manager = DatabaseManager(cleanup_existing=True)
+
+
+# For FastAPI dependency injection
+def get_session() -> Generator[Session, None, None]:
+    """Yield a new session for database operations."""
+    with db_manager.session() as session:
+        yield session

lightly_studio/api/db_tables.py

@@ -0,0 +1,32 @@
+"""Module provides functions to initialize and manage the DuckDB."""
+
+from lightly_studio.models.annotation.annotation_base import (
+    AnnotationBaseTable,  # noqa: F401, required for SQLModel to work properly
+)
+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
+)
+from lightly_studio.models.embedding_model import (
+    EmbeddingModelTable,  # noqa: F401, required for SQLModel to work properly
+)
+from lightly_studio.models.metadata import (
+    SampleMetadataTable,  # noqa: F401, required for SQLModel to work properly
+)
+from lightly_studio.models.sample import (
+    SampleTable,  # noqa: F401, required for SQLModel to work properly
+)
+from lightly_studio.models.sample_embedding import (
+    SampleEmbeddingTable,  # noqa: F401, required for SQLModel to work properly
+)
+from lightly_studio.models.settings import (
+    SettingTable,  # noqa: F401, required for SQLModel to work properly
+)
+from lightly_studio.models.tag import (
+    TagTable,  # noqa: F401, required for SQLModel to work properly
+)

lightly_studio/api/features.py

@@ -0,0 +1,7 @@
+"""Global list fo the active features."""
+
+# TODO(Kondrat 04/25): Pass the feature flag to app
+# https://linear.app/lightly/issue/LIG-6708/introduce-apifeatures-endpoint
+from typing import List
+
+lightly_studio_active_features: List[str] = []

lightly_studio/api/routes/api/annotation.py

@@ -0,0 +1,233 @@
+"""This module contains the API routes for managing annotations."""
+
+from __future__ import annotations
+
+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.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.models.annotation.annotation_base import (
+    AnnotationBaseTable,
+    AnnotationDetailsView,
+    AnnotationViewsWithCount,
+)
+from lightly_studio.models.dataset import DatasetTable
+from lightly_studio.resolvers import annotation_resolver, tag_resolver
+from lightly_studio.resolvers.annotation_resolver.get_all import (
+    GetAllAnnotationsResult,
+)
+from lightly_studio.resolvers.annotations.annotations_filter import (
+    AnnotationsFilter,
+)
+from lightly_studio.services import annotations_service
+from lightly_studio.services.annotations_service.update_annotation import (
+    AnnotationUpdate,
+)
+
+annotations_router = APIRouter(prefix="/datasets/{dataset_id}", tags=["annotations"])
+SessionDep = Annotated[Session, Depends(get_session)]
+
+
+@annotations_router.get("/annotations/count")
+def count_annotations_by_dataset(  # noqa: PLR0913 // FIXME: refactor to use proper pydantic
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(get_and_validate_dataset_id),
+    ],
+    session: SessionDep,
+    filtered_labels: Annotated[list[str] | None, Query()] = None,
+    min_width: Annotated[int | None, Query(ge=0)] = None,
+    max_width: Annotated[int | None, Query(ge=0)] = None,
+    min_height: Annotated[int | None, Query(ge=0)] = None,
+    max_height: Annotated[int | None, Query(ge=0)] = None,
+    tag_ids: list[UUID] | None = None,
+) -> list[dict[str, str | int]]:
+    """Get annotation counts for a specific dataset.
+
+    Returns a list of dictionaries with label name and count.
+    """
+    counts = annotation_resolver.count_annotations_by_dataset(
+        session=session,
+        dataset_id=dataset.dataset_id,
+        filtered_labels=filtered_labels,
+        min_width=min_width,
+        max_width=max_width,
+        min_height=min_height,
+        max_height=max_height,
+        tag_ids=tag_ids,
+    )
+    return [
+        {
+            "label_name": label_name,
+            "current_count": current_count,
+            "total_count": total_count,
+        }
+        for label_name, current_count, total_count in counts
+    ]
+
+
+@annotations_router.get(
+    "/annotations",
+    response_model=AnnotationViewsWithCount,
+)
+def read_annotations(
+    dataset_id: Annotated[UUID, Path(title="Dataset Id", description="The ID of the dataset")],
+    session: SessionDep,
+    pagination: Annotated[PaginatedWithCursor, Depends()],
+    annotation_label_ids: Annotated[list[UUID] | None, Query()] = None,
+    tag_ids: Annotated[list[UUID] | None, Query()] = None,
+) -> GetAllAnnotationsResult:
+    """Retrieve a list of annotations from the database."""
+    return annotation_resolver.get_all(
+        session=session,
+        pagination=Paginated(
+            offset=pagination.offset,
+            limit=pagination.limit,
+        ),
+        filters=AnnotationsFilter(
+            dataset_ids=[dataset_id],
+            annotation_label_ids=annotation_label_ids,
+            annotation_tag_ids=tag_ids,
+        ),
+    )
+
+
+@annotations_router.post(
+    "/annotations/{annotation_id}/tag/{tag_id}",
+    status_code=HTTP_STATUS_CREATED,
+)
+def add_tag_to_annotation(
+    session: SessionDep,
+    annotation_id: UUID,
+    tag_id: UUID,
+) -> bool:
+    """Add annotation to a tag."""
+    annotation = annotation_resolver.get_by_id(session=session, annotation_id=annotation_id)
+    if not annotation:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail=f"Annotation {annotation_id} not found",
+        )
+
+    if not tag_resolver.add_tag_to_annotation(
+        session=session, tag_id=tag_id, annotation=annotation
+    ):
+        raise HTTPException(status_code=HTTP_STATUS_NOT_FOUND, detail=f"Tag {tag_id} not found")
+
+    return True
+
+
+class AnnotationUpdateInput(BaseModel):
+    """API input model for updating an annotation."""
+
+    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
+
+
+@annotations_router.put("/annotations/{annotation_id}")
+def update_annotation(
+    session: SessionDep,
+    dataset_id: Annotated[
+        UUID,
+        Path(title="Dataset Id"),
+    ],
+    annotation_id: Annotated[
+        UUID,
+        Path(title="Annotation ID", description="ID of the annotation to update"),
+    ],
+    annotation_update_input: Annotated[AnnotationUpdateInput, Body()],
+) -> AnnotationBaseTable:
+    """Update an existing annotation in the database."""
+    return annotations_service.update_annotation(
+        session=session,
+        annotation_update=AnnotationUpdate(
+            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,
+        ),
+    )
+
+
+@annotations_router.put(
+    "/annotations",
+)
+def update_annotations(
+    session: SessionDep,
+    dataset_id: Annotated[
+        UUID,
+        Path(title="Dataset Id"),
+    ],
+    annotation_update_inputs: Annotated[list[AnnotationUpdateInput], Body()],
+) -> list[AnnotationBaseTable]:
+    """Update multiple annotations in the database."""
+    return annotations_service.update_annotations(
+        session=session,
+        annotation_updates=[
+            AnnotationUpdate(
+                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,
+            )
+            for annotation_update_input in annotation_update_inputs
+        ],
+    )
+
+
+@annotations_router.get("/annotations/{annotation_id}", response_model=AnnotationDetailsView)
+def get_annotation(
+    session: SessionDep,
+    dataset_id: Annotated[  # noqa: ARG001
+        UUID,
+        Path(title="Dataset Id", description="The ID of the dataset"),
+    ],  # We need dataset_id because otherwise the path would not match
+    annotation_id: Annotated[UUID, Path(title="Annotation ID")],
+) -> AnnotationBaseTable:
+    """Retrieve an existing annotation from the database."""
+    return annotations_service.get_annotation_by_id(session=session, annotation_id=annotation_id)
+
+
+@annotations_router.delete("/annotations/{annotation_id}/tag/{tag_id}")
+def remove_tag_from_annotation(
+    session: SessionDep,
+    tag_id: UUID,
+    annotation_id: UUID,
+) -> bool:
+    """Remove annotation from a tag."""
+    annotation = annotation_resolver.get_by_id(session=session, annotation_id=annotation_id)
+    if not annotation:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail=f"Annotation {annotation_id} not found",
+        )
+
+    if not tag_resolver.remove_tag_from_annotation(
+        session=session, tag_id=tag_id, annotation=annotation
+    ):
+        raise HTTPException(status_code=HTTP_STATUS_NOT_FOUND, detail=f"Tag {tag_id} not found")
+
+    return True

lightly_studio/api/routes/api/annotation_label.py

@@ -0,0 +1,90 @@
+"""This module contains the API routes for managing annotation labels."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from fastapi import APIRouter, Depends, HTTPException
+from sqlmodel import Session
+from typing_extensions import Annotated
+
+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.models.annotation_label import (
+    AnnotationLabelCreate,
+    AnnotationLabelTable,
+)
+from lightly_studio.resolvers import annotation_label_resolver
+
+annotations_label_router = APIRouter()
+SessionDep = Annotated[Session, Depends(get_session)]
+
+
+@annotations_label_router.post(
+    "/annotation_labels",
+    status_code=HTTP_STATUS_CREATED,
+)
+def create_annotation_label(
+    input_label: AnnotationLabelCreate,
+    session: SessionDep,
+) -> AnnotationLabelTable:
+    """Create a new annotation label in the database."""
+    return annotation_label_resolver.create(session=session, label=input_label)
+
+
+@annotations_label_router.get("/annotation_labels")
+def read_annotation_labels(
+    session: SessionDep,
+) -> list[AnnotationLabelTable]:
+    """Retrieve a list of annotation labels from the database."""
+    return annotation_label_resolver.get_all(session=session)
+
+
+@annotations_label_router.get("/annotation_labels/{label_id}")
+def read_annotation_label(
+    label_id: UUID,
+    session: SessionDep,
+) -> AnnotationLabelTable:
+    """Retrieve a single annotation label from the database."""
+    label = annotation_label_resolver.get_by_id(session=session, label_id=label_id)
+    if not label:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail="Annotation label not found",
+        )
+    return label
+
+
+@annotations_label_router.put("/annotation_labels/{label_id}")
+def update_annotation_label(
+    label_id: UUID,
+    label_input: AnnotationLabelCreate,
+    session: SessionDep,
+) -> AnnotationLabelTable:
+    """Update an existing annotation label in the database."""
+    label = annotation_label_resolver.update(
+        session=session, label_id=label_id, label_data=label_input
+    )
+    if not label:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail="Annotation label not found",
+        )
+    return label
+
+
+@annotations_label_router.delete("/annotation_labels/{label_id}")
+def delete_annotation_label(
+    label_id: UUID,
+    session: SessionDep,
+) -> dict[str, str]:
+    """Delete an annotation label from the database."""
+    if not annotation_label_resolver.delete(session=session, label_id=label_id):
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail="Annotation label not found",
+        )
+    return {"status": "deleted"}