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

lightly_studio/api/app.py

@@ -16,6 +16,7 @@ from lightly_studio.api.routes import healthz, images, webapp
 from lightly_studio.api.routes.api import (
     annotation,
     annotation_label,
+    caption,
     classifier,
     dataset,
     dataset_tag,
@@ -89,6 +90,7 @@ 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(settings.settings_router)
 api_router.include_router(classifier.classifier_router)

lightly_studio/api/features.py

@@ -1,7 +1,5 @@
-"""Global list fo the active features."""
+"""Global list of 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
+from __future__ import annotations
 
-lightly_studio_active_features: List[str] = []
+lightly_studio_active_features: list[str] = []

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_tag.py

@@ -164,6 +164,11 @@ class SampleIdsBody(BaseModel):
 )
 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:
@@ -215,6 +220,11 @@ class AnnotationIdsBody(BaseModel):
 )
 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:

lightly_studio/api/routes/api/embeddings2d.py

@@ -3,25 +3,38 @@
 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 sklearn.manifold import TSNE
+from pydantic import BaseModel, Field
 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
+from lightly_studio.resolvers import sample_embedding_resolver, sample_resolver
+from lightly_studio.resolvers.samples_filter import SampleFilter
+from lightly_studio.resolvers.twodim_embedding_resolver import _calculate_2d_embeddings
 
 embeddings2d_router = APIRouter()
 
 
-@embeddings2d_router.get("/embeddings2d/tsne")
-def get_embeddings2d__tsne(session: SessionDep) -> Response:
+class GetEmbeddings2DRequest(BaseModel):
+    """Request body for retrieving 2D embeddings."""
+
+    filters: SampleFilter | None = Field(
+        None,
+        description="Filter parameters identifying matching samples",
+    )
+
+
+@embeddings2d_router.post("/embeddings2d/default")
+def get_2d_embeddings(
+    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()
@@ -37,17 +50,33 @@ def get_embeddings2d__tsne(session: SessionDep) -> Response:
     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]
+    planar_embeddings = _calculate_2d_embeddings([e.embedding for e in embeddings])
+    # TODO(lukas 10/2025): This can become very slow for large datasets, we should consider using
+    # `pa.chunked_array` to directly construct the xy array.
+    x = (x for (x, _) in planar_embeddings)
+    y = (y for (_, y) in planar_embeddings)
+
+    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.
@@ -57,6 +86,8 @@ def get_embeddings2d__tsne(session: SessionDep) -> Response:
         {
             "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()),
         }
     )
 
@@ -74,31 +105,3 @@ def get_embeddings2d__tsne(session: SessionDep) -> Response:
             "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/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/core/add_samples.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import json
+from collections import defaultdict
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Iterable
@@ -26,10 +28,12 @@ from tqdm import tqdm
 
 from lightly_studio.models.annotation.annotation_base import AnnotationCreate
 from lightly_studio.models.annotation_label import AnnotationLabelCreate
+from lightly_studio.models.caption import CaptionCreate
 from lightly_studio.models.sample import SampleCreate, SampleTable
 from lightly_studio.resolvers import (
     annotation_label_resolver,
     annotation_resolver,
+    caption_resolver,
     sample_resolver,
 )
 
@@ -218,6 +222,111 @@ def load_into_dataset_from_labelformat(
     return created_sample_ids
 
 
+def load_into_dataset_from_coco_captions(
+    session: Session,
+    dataset_id: UUID,
+    annotations_json: Path,
+    images_path: Path,
+) -> list[UUID]:
+    """Load samples and captions from a COCO captions file into the dataset.
+
+    Args:
+        session: Database session used for resolver operations.
+        dataset_id: Identifier of the dataset that receives the samples.
+        annotations_json: Path to the COCO captions annotations file.
+        images_path: Directory containing the referenced images.
+
+    Returns:
+        The list of newly created sample identifiers.
+    """
+    with fsspec.open(str(annotations_json), "r") as file:
+        coco_payload = json.load(file)
+
+    images: list[dict[str, object]] = coco_payload.get("images", [])
+    annotations: list[dict[str, object]] = coco_payload.get("annotations", [])
+
+    captions_by_image_id: dict[int, list[str]] = defaultdict(list)
+    for annotation in annotations:
+        image_id = annotation["image_id"]
+        caption = annotation["caption"]
+        if not isinstance(image_id, int):
+            continue
+        if not isinstance(caption, str):
+            continue
+        caption_text = caption.strip()
+        if not caption_text:
+            continue
+        captions_by_image_id[image_id].append(caption_text)
+
+    logging_context = _LoadingLoggingContext(
+        n_samples_to_be_inserted=len(images),
+        n_samples_before_loading=sample_resolver.count_by_dataset_id(
+            session=session, dataset_id=dataset_id
+        ),
+    )
+
+    captions_to_create: list[CaptionCreate] = []
+    samples_to_create: list[SampleCreate] = []
+    created_sample_ids: list[UUID] = []
+    image_path_to_captions: dict[str, list[str]] = {}
+
+    for image_info in tqdm(images, desc="Processing images", unit=" images"):
+        if isinstance(image_info["id"], int):
+            image_id_raw = image_info["id"]
+        else:
+            continue
+        file_name_raw = str(image_info["file_name"])
+
+        width = image_info["width"] if isinstance(image_info["width"], int) else 0
+        height = image_info["height"] if isinstance(image_info["height"], int) else 0
+        sample = SampleCreate(
+            file_name=file_name_raw,
+            file_path_abs=str(images_path / file_name_raw),
+            width=width,
+            height=height,
+            dataset_id=dataset_id,
+        )
+        samples_to_create.append(sample)
+        image_path_to_captions[sample.file_path_abs] = captions_by_image_id.get(image_id_raw, [])
+
+        if len(samples_to_create) >= SAMPLE_BATCH_SIZE:
+            created_samples_batch, paths_not_inserted = _create_batch_samples(
+                session=session, samples=samples_to_create
+            )
+            created_sample_ids.extend(s.sample_id for s in created_samples_batch)
+            logging_context.update_example_paths(paths_not_inserted)
+            _process_batch_captions(
+                session=session,
+                dataset_id=dataset_id,
+                stored_samples=created_samples_batch,
+                image_path_to_captions=image_path_to_captions,
+                captions_to_create=captions_to_create,
+            )
+            samples_to_create.clear()
+            image_path_to_captions.clear()
+
+    if samples_to_create:
+        created_samples_batch, paths_not_inserted = _create_batch_samples(
+            session=session, samples=samples_to_create
+        )
+        created_sample_ids.extend(s.sample_id for s in created_samples_batch)
+        logging_context.update_example_paths(paths_not_inserted)
+        _process_batch_captions(
+            session=session,
+            dataset_id=dataset_id,
+            stored_samples=created_samples_batch,
+            image_path_to_captions=image_path_to_captions,
+            captions_to_create=captions_to_create,
+        )
+
+    if captions_to_create:
+        caption_resolver.create_many(session=session, captions=captions_to_create)
+
+    _log_loading_results(session=session, dataset_id=dataset_id, logging_context=logging_context)
+
+    return created_sample_ids
+
+
 def _log_loading_results(
     session: Session, dataset_id: UUID, logging_context: _LoadingLoggingContext
 ) -> None:
@@ -372,3 +481,32 @@ def _process_batch_annotations(  # noqa: PLR0913
         if len(annotations_to_create) >= ANNOTATION_BATCH_SIZE:
             annotation_resolver.create_many(session=session, annotations=annotations_to_create)
             annotations_to_create.clear()
+
+
+def _process_batch_captions(
+    session: Session,
+    dataset_id: UUID,
+    stored_samples: list[SampleTable],
+    image_path_to_captions: dict[str, list[str]],
+    captions_to_create: list[CaptionCreate],
+) -> None:
+    """Process captions for a batch of samples."""
+    if not stored_samples:
+        return
+
+    for stored_sample in stored_samples:
+        captions = image_path_to_captions[stored_sample.file_path_abs]
+        if not captions:
+            continue
+
+        for caption_text in captions:
+            caption = CaptionCreate(
+                dataset_id=dataset_id,
+                sample_id=stored_sample.sample_id,
+                text=caption_text,
+            )
+            captions_to_create.append(caption)
+
+        if len(captions_to_create) >= ANNOTATION_BATCH_SIZE:
+            caption_resolver.create_many(session=session, captions=captions_to_create)
+            captions_to_create.clear()