lightly-studio 0.3.1__py3-none-any.whl

lightly_studio/api/routes/api/annotation_task.py

@@ -0,0 +1,38 @@
+"""API endpoints for annotation tasks."""
+
+from typing import List
+from uuid import UUID
+
+from fastapi import APIRouter, Depends, HTTPException, status
+from sqlmodel import Session
+
+from lightly_studio.api.db import get_session
+from lightly_studio.models.annotation_task import AnnotationTaskTable
+from lightly_studio.resolvers import annotation_task_resolver
+
+router = APIRouter(prefix="/annotationtasks", tags=["annotationtasks"])
+
+
+@router.get("/", response_model=List[AnnotationTaskTable])
+def get_annotation_tasks(
+    session: Session = Depends(get_session),  # noqa: B008
+) -> List[AnnotationTaskTable]:
+    """Get all annotation tasks."""
+    return annotation_task_resolver.get_all(session=session)
+
+
+@router.get("/{annotation_task_id}", response_model=AnnotationTaskTable)
+def get_annotation_task(
+    annotation_task_id: UUID,
+    session: Session = Depends(get_session),  # noqa: B008
+) -> AnnotationTaskTable:
+    """Get an annotation task by ID."""
+    task = annotation_task_resolver.get_by_id(
+        session=session, annotation_task_id=annotation_task_id
+    )
+    if task is None:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"Annotation task with ID {annotation_task_id} not found",
+        )
+    return task

lightly_studio/api/routes/api/classifier.py

@@ -0,0 +1,387 @@
+"""This module contains the API routes for managing classifiers."""
+
+from __future__ import annotations
+
+import io
+from pathlib import Path
+from uuid import UUID
+
+from fastapi import APIRouter, Depends, 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.few_shot_classifier.classifier import (
+    ExportType,
+)
+from lightly_studio.few_shot_classifier.classifier_manager import (
+    ClassifierManagerProvider,
+)
+from lightly_studio.models.classifier import EmbeddingClassifier
+
+classifier_router = APIRouter()
+SessionDep = Annotated[Session, Depends(get_session)]
+
+
+class GetNegativeSamplesRequest(BaseModel):
+    """Request for getting negative samples for classifier training."""
+
+    positive_sample_ids: list[UUID]
+    dataset_id: UUID
+
+
+class GetNegativeSamplesResponse(BaseModel):
+    """Response for getting negative samples for classifier training."""
+
+    negative_sample_ids: list[UUID]
+
+
+@classifier_router.post("/classifiers/get_negative_samples")
+def get_negative_samples(
+    request: GetNegativeSamplesRequest, session: SessionDep
+) -> GetNegativeSamplesResponse:
+    """Get negative samples for classifier training.
+
+    Args:
+        request: The request containing negative sample parameters.
+        session: Database session.
+
+    Returns:
+        The response containing negative sample IDs.
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    negative_samples = classifier_manager.provide_negative_samples(
+        session=session,
+        dataset_id=request.dataset_id,
+        selected_samples=request.positive_sample_ids,
+    )
+    # Extract just the sample IDs from the returned Sample objects
+    negative_sample_ids = [sample.sample_id for sample in negative_samples]
+    return GetNegativeSamplesResponse(negative_sample_ids=negative_sample_ids)
+
+
+class SamplesToRefineResponse(BaseModel):
+    """Response for samples for classifier refinement.
+
+    Maps class names to lists of sample IDs. First class gets high confidence
+    samples, second class gets low confidence samples.
+    """
+
+    samples: dict[str, list[UUID]]
+
+
+@classifier_router.get("/classifiers/{classifier_id}/samples_to_refine")
+def samples_to_refine(
+    classifier_id: UUID,
+    dataset_id: UUID,
+    session: SessionDep,
+) -> SamplesToRefineResponse:
+    """Get samples for classifier refinement.
+
+    Args:
+        classifier_id: The ID of the classifier.
+        dataset_id: The ID of the dataset.
+        session: Database session.
+
+    Returns:
+        The response containing sample IDs for refinement.
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    samples = classifier_manager.get_samples_for_fine_tuning(
+        session=session, classifier_id=classifier_id, dataset_id=dataset_id
+    )
+    return SamplesToRefineResponse(samples=samples)
+
+
+@classifier_router.get("/classifiers/{classifier_id}/sample_history")
+def sample_history(
+    classifier_id: UUID,
+) -> SamplesToRefineResponse:
+    """Get all samples used in the classifier training.
+
+    Args:
+        classifier_id: The ID of the classifier.
+
+    Returns:
+        The response containing sample IDs used in the training.
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    samples = classifier_manager.get_annotations(classifier_id=classifier_id)
+    return SamplesToRefineResponse(samples=samples)
+
+
+@classifier_router.post(
+    "/classifiers/{classifier_id}/commit_temp_classifier",
+)
+def commit_temp_classifier(
+    classifier_id: UUID,
+) -> None:
+    """Commit the classifier.
+
+    Args:
+        classifier_id: The ID of the classifier.
+
+    Returns:
+        None
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    classifier_manager.commit_temp_classifier(classifier_id=classifier_id)
+
+
+@classifier_router.delete(
+    "/classifiers/{classifier_id}/drop_temp_classifier",
+)
+def drop_temp_classifier(
+    classifier_id: UUID,
+) -> None:
+    """Drop the classifier.
+
+    Args:
+        classifier_id: The ID of the classifier.
+
+    Returns:
+        None
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    classifier_manager.drop_temp_classifier(classifier_id=classifier_id)
+
+
+class SaveClassifierRequest(BaseModel):
+    """Request for saving classifier to a file."""
+
+    file_path: str
+
+
+@classifier_router.post(
+    "/classifiers/{classifier_id}/save_classifier_to_file/{export_type}",
+)
+def save_classifier_to_file(
+    classifier_id: UUID,
+    export_type: ExportType,
+) -> StreamingResponse:
+    """Save the classifier to a file.
+
+    Args:
+        classifier_id: The ID of the classifier.
+        export_type: The type of export (e.g., "sklearn", "lightly").
+
+    Returns:
+        StreamingResponse containing the pickled classifier file.
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    # Use BytesIO to capture the file content and send it as a response.
+    buffer = io.BytesIO()
+    classifier_manager.save_classifier_to_buffer(
+        classifier_id=classifier_id, buffer=buffer, export_type=export_type
+    )
+    buffer.seek(0)
+
+    # Get classifier name for the filename
+    classifier = classifier_manager.get_classifier_by_id(classifier_id=classifier_id)
+    filename = f"{classifier.classifier_name}.pkl"
+    headers = {
+        "Content-Disposition": f'attachment; filename="{filename}"',
+        "Content-Type": "application/octet-stream",
+        "Access-Control-Expose-Headers": "Content-Disposition",
+    }
+
+    return StreamingResponse(buffer, headers=headers, media_type="application/octet-stream")
+
+
+class LoadClassifierRequest(BaseModel):
+    """Request for loading classifier from a file."""
+
+    file_path: str
+
+
+class LoadClassifierResponse(BaseModel):
+    """Response for loading classifier from a file."""
+
+    classifier_id: UUID
+
+
+@classifier_router.post(
+    "/classifiers/load_classifier_from_file",
+)
+def load_classifier_from_file(
+    request: LoadClassifierRequest,
+    session: SessionDep,
+) -> LoadClassifierResponse:
+    """Load the classifier from a file.
+
+    Args:
+        request: The request containing the file path.
+        session: Database session.
+
+    Returns:
+        Response with the ID of the loaded classifier.
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    classifier = classifier_manager.load_classifier_from_file(
+        session=session, file_path=Path(request.file_path)
+    )
+    return LoadClassifierResponse(classifier_id=classifier.classifier_id)
+
+
+@classifier_router.post(
+    "/classifiers/load_classifier_from_buffer",
+)
+def load_classifier_from_buffer(
+    file: UploadFile,
+    session: SessionDep,
+) -> UUID:
+    """Load a classifier from an uploaded file buffer.
+
+    Args:
+        file: The uploaded classifier file.
+        session: Database session.
+
+    Returns:
+        The ID of the loaded classifier.
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+
+    # Read file into buffer
+    buffer = io.BytesIO(file.file.read())
+
+    # Load classifier from buffer
+    classifier = classifier_manager.load_classifier_from_buffer(session=session, buffer=buffer)
+    return classifier.classifier_id
+
+
+@classifier_router.post(
+    "/classifiers/{classifier_id}/train_classifier",
+)
+def train_classifier(
+    classifier_id: UUID,
+    session: SessionDep,
+) -> None:
+    """Train the classifier.
+
+    Args:
+        classifier_id: The ID of the classifier.
+        session: Database session.
+
+    Returns:
+        None
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    classifier_manager.train_classifier(session=session, classifier_id=classifier_id)
+
+
+class UpdateAnnotationsRequest(BaseModel):
+    """Request for updating classifier annotations."""
+
+    annotations: dict[str, list[UUID]]
+
+
+@classifier_router.post(
+    "/classifiers/{classifier_id}/update_annotations",
+)
+def update_classifiers_annotations(
+    classifier_id: UUID,
+    request: UpdateAnnotationsRequest,
+) -> None:
+    """Update the annotations for a classifier.
+
+    Args:
+        classifier_id: The ID of the classifier.
+        request: The request containing the new annotations.
+
+    Returns:
+        None
+
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    classifier_manager.update_classifiers_annotations(
+        classifier_id=classifier_id,
+        new_annotations=request.annotations,
+    )
+
+
+class CreateClassifierRequest(BaseModel):
+    """Request model for creating a classifier."""
+
+    name: str
+    class_list: list[str]
+    dataset_id: UUID
+
+
+class CreateClassifierResponse(BaseModel):
+    """Response model for creating a classifier."""
+
+    name: str
+    classifier_id: str
+
+
+@classifier_router.post("/classifiers/create")
+def create_classifier(
+    request: CreateClassifierRequest, session: SessionDep
+) -> CreateClassifierResponse:
+    """Create a new classifier.
+
+    Args:
+        request: The request containing classifier creation parameters.
+        session: Database session.
+
+    Returns:
+        Response with the name and ID of the classifier.
+
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    classifier = classifier_manager.create_classifier(
+        session=session,
+        name=request.name,
+        class_list=request.class_list,
+        dataset_id=request.dataset_id,
+    )
+    return CreateClassifierResponse(
+        name=classifier.few_shot_classifier.name,
+        classifier_id=str(classifier.classifier_id),
+    )
+
+
+class GetAllClassifiersResponse(BaseModel):
+    """Response model for getting all active classifiers."""
+
+    classifiers: list[EmbeddingClassifier]
+
+
+@classifier_router.get("/classifiers/get_all_classifiers")
+def get_all_classifiers() -> GetAllClassifiersResponse:
+    """Get all active classifiers.
+
+    Returns:
+        Response with list of tuples containing classifier names and IDs.
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    classifiers = classifier_manager.get_all_classifiers()
+    return GetAllClassifiersResponse(classifiers=classifiers)
+
+
+@classifier_router.post(
+    "/classifiers/{classifier_id}/run_on_dataset/{dataset_id}",
+)
+def run_classifier_route(
+    classifier_id: UUID,
+    dataset_id: UUID,
+    session: SessionDep,
+) -> None:
+    """Run the classifier on a dataset.
+
+    Args:
+        dataset_id: The ID of the dataset to run the classifier on.
+        classifier_id: The ID of the classifier.
+        session: Database session.
+
+    Returns:
+        None
+    """
+    classifier_manager = ClassifierManagerProvider.get_classifier_manager()
+    classifier_manager.run_classifier(
+        session=session,
+        classifier_id=classifier_id,
+        dataset_id=dataset_id,
+    )

lightly_studio/api/routes/api/dataset.py

@@ -0,0 +1,182 @@
+"""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, Session
+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.models.dataset import (
+    DatasetCreate,
+    DatasetTable,
+    DatasetView,
+)
+from lightly_studio.resolvers import dataset_resolver
+from lightly_studio.resolvers.dataset_resolver import (
+    ExportFilter,
+)
+
+dataset_router = APIRouter()
+SessionDep = Annotated[Session, Depends(get_session)]
+
+
+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.post(
+    "/datasets",
+    response_model=DatasetView,
+    status_code=201,
+)
+def create_dataset(
+    dataset_input: DatasetCreate,
+    session: SessionDep,
+) -> DatasetTable:
+    """Create a new dataset in the database."""
+    return dataset_resolver.create(session=session, dataset=dataset_input)
+
+
+@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}")
+def read_dataset(
+    dataset: Annotated[
+        DatasetTable,
+        Path(title="Dataset Id"),
+        Depends(get_and_validate_dataset_id),
+    ],
+) -> DatasetTable:
+    """Retrieve a single dataset from the database."""
+    return 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"}
+
+
+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/
+@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
+
+
+"""
+Endpoint to export samples from a dataset.
+"""
+
+
+@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,
+    )