lightly-studio 0.3.3__tar.gz → 0.4.0__tar.gz

{lightly_studio-0.3.3 → lightly_studio-0.4.0}/.gitignore

@@ -16,3 +16,6 @@ datasets
 
 # Generated API schema
 openapi.json
+
+# docs build
+/docs/site

lightly_studio-0.4.0/PKG-INFO

@@ -0,0 +1,78 @@
+Metadata-Version: 2.4
+Name: lightly-studio
+Version: 0.4.0
+Summary: LightlyStudio is a lightweight, fast, and easy-to-use data exploration tool for data scientists and engineers.
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: <3.13,>=3.8
+Requires-Dist: annotated-types==0.7.0
+Requires-Dist: duckdb-engine<0.17,>=0.15.0
+Requires-Dist: duckdb<1.3,>=1.2.2
+Requires-Dist: environs<12.0.0
+Requires-Dist: eval-type-backport>=0.2.2
+Requires-Dist: fastapi>=0.115.5
+Requires-Dist: faster-coco-eval>=1.6.5
+Requires-Dist: fsspec>=2023.1.0
+Requires-Dist: labelformat>=0.1.7
+Requires-Dist: lightly-mundig==0.1.7
+Requires-Dist: open-clip-torch>=2.20.0
+Requires-Dist: pyarrow>=17.0.0
+Requires-Dist: python-multipart>=0.0.20
+Requires-Dist: scikit-learn==1.3.2
+Requires-Dist: sqlmodel>=0.0.22
+Requires-Dist: tqdm>=4.65.0
+Requires-Dist: typing-extensions>=4.12.2
+Requires-Dist: uvicorn>=0.32.1
+Requires-Dist: xxhash>=3.5.0
+Provides-Extra: cloud-storage
+Requires-Dist: adlfs>=2023.1.0; extra == 'cloud-storage'
+Requires-Dist: gcsfs>=2023.1.0; extra == 'cloud-storage'
+Requires-Dist: s3fs>=2023.1.0; extra == 'cloud-storage'
+Provides-Extra: lightly-edge
+Requires-Dist: lightly-edge-sdk>=1.0.1b2; extra == 'lightly-edge'
+Requires-Dist: opencv-python; extra == 'lightly-edge'
+Description-Content-Type: text/markdown
+
+<p align="center">
+  <a href="https://lightly.ai/lightly-studio"> 
+    <img src="https://cdn.prod.website-files.com/62cd5ce03261cba217188442/66dac501a8e9a90495970876_Logo%20dark-short-p-800.png" height="50" alt="LightlyStudio logo" />
+  </a>
+</p>
+<p align="center"><strong>Curate, Annotate, and Manage Your Data in LightlyStudio.</strong></p>
+<p align="center">
+  <a href="https://pypi.org/project/lightly-studio">
+    <img src="https://img.shields.io/pypi/pyversions/lightly-studio" alt="PyPI python" />
+  </a>
+  <a href="https://pypi.org/project/lightly-studio">
+    <img src="https://badge.fury.io/py/lightly-studio.svg" alt="PyPI version" />
+  </a>
+  <a href="LICENSE">
+    <img src="https://img.shields.io/badge/License-Apache%202.0-blue.svg" alt="License" />
+  </a>
+  <a href="https://docs.lightly.ai/studio">
+    <img src="https://img.shields.io/badge/Docs-blue" alt="Docs" />
+  </a>
+</p>
+
+---
+
+# Welcome to LightlyStudio!
+
+We at **[Lightly](https://lightly.ai)** created **[LightlyStudio](https://www.lightly.ai/lightly-studio)**, an open-source tool designed to unify your data workflows from curation, annotation and management in a single tool. Since we're big fans of Rust we used it to speed things up. You can work with COCO and ImageNet on a Macbook Pro with M1 and 16GB of memory!
+
+Explore [LightlyStudio on GitHub](https://github.com/lightly-ai/lightly-studio).
+
+Learn more in our [documentation](https://docs.lightly.ai/studio).
+
+<p align="center">
+  <img alt="LightlyStudio Overview" src="https://storage.googleapis.com/lightly-public/studio/studio_overview.gif" width="70%">
+  <br/>
+  <em>Curate, Annotate, and Manage Your Data in LightlyStudio.</em>
+</p>

lightly_studio-0.4.0/USAGE.md

@@ -0,0 +1,36 @@
+<p align="center">
+  <a href="https://lightly.ai/lightly-studio"> 
+    <img src="https://cdn.prod.website-files.com/62cd5ce03261cba217188442/66dac501a8e9a90495970876_Logo%20dark-short-p-800.png" height="50" alt="LightlyStudio logo" />
+  </a>
+</p>
+<p align="center"><strong>Curate, Annotate, and Manage Your Data in LightlyStudio.</strong></p>
+<p align="center">
+  <a href="https://pypi.org/project/lightly-studio">
+    <img src="https://img.shields.io/pypi/pyversions/lightly-studio" alt="PyPI python" />
+  </a>
+  <a href="https://pypi.org/project/lightly-studio">
+    <img src="https://badge.fury.io/py/lightly-studio.svg" alt="PyPI version" />
+  </a>
+  <a href="LICENSE">
+    <img src="https://img.shields.io/badge/License-Apache%202.0-blue.svg" alt="License" />
+  </a>
+  <a href="https://docs.lightly.ai/studio">
+    <img src="https://img.shields.io/badge/Docs-blue" alt="Docs" />
+  </a>
+</p>
+
+---
+
+# Welcome to LightlyStudio!
+
+We at **[Lightly](https://lightly.ai)** created **[LightlyStudio](https://www.lightly.ai/lightly-studio)**, an open-source tool designed to unify your data workflows from curation, annotation and management in a single tool. Since we're big fans of Rust we used it to speed things up. You can work with COCO and ImageNet on a Macbook Pro with M1 and 16GB of memory!
+
+Explore [LightlyStudio on GitHub](https://github.com/lightly-ai/lightly-studio).
+
+Learn more in our [documentation](https://docs.lightly.ai/studio).
+
+<p align="center">
+  <img alt="LightlyStudio Overview" src="https://storage.googleapis.com/lightly-public/studio/studio_overview.gif" width="70%">
+  <br/>
+  <em>Curate, Annotate, and Manage Your Data in LightlyStudio.</em>
+</p>

{lightly_studio-0.3.3 → lightly_studio-0.4.0}/pyproject.toml

@@ -1,15 +1,16 @@
 [project]
 name = "lightly-studio"
-version = "0.3.3"
+version = "0.4.0"
 description = "LightlyStudio is a lightweight, fast, and easy-to-use data exploration tool for data scientists and engineers."
 readme = "USAGE.md"
-requires-python = ">=3.8"
+requires-python = ">=3.8,<3.13"
 classifiers = [
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
     "Operating System :: POSIX :: Linux",
     "Operating System :: MacOS :: MacOS X",
     "Operating System :: Microsoft :: Windows",
@@ -43,7 +44,7 @@ dependencies = [
     #   lightly-mundig==$version
     # Then lock the version. The uv.lock file is updated accordingly.
     # uv lock --upgrade-package lightly-mundig==$version
-    "lightly-mundig==0.1.4",
+    "lightly-mundig==0.1.7",
     "pyarrow>=17.0.0",
 ]
 
@@ -74,6 +75,15 @@ dev = [
     "pytest-xdist>=3.6.1",
 ]
 
+docs = [
+    "mkdocs>=1.5.0",
+    "mkdocs-material>=9.4.0",
+    "mkdocs-git-revision-date-localized-plugin>=1.2.0",
+    "mkdocstrings[python]>=0.24.0",
+    "mkdocstrings-python>=1.8.0",
+    "pymdown-extensions>=10.0.0",
+]
+
 
 [tool.mypy]
 strict = true
@@ -131,3 +141,8 @@ explicit = true
 artifacts = [
     "src/lightly_studio/**/*",
 ]
+
+[tool.ruff]
+extend-exclude = [
+    "docs",
+]

{lightly_studio-0.3.3 → lightly_studio-0.4.0}/src/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-0.4.0/src/lightly_studio/api/features.py

@@ -0,0 +1,5 @@
+"""Global list of the active features."""
+
+from __future__ import annotations
+
+lightly_studio_active_features: list[str] = []

lightly_studio-0.4.0/src/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-0.3.3 → lightly_studio-0.4.0}/src/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-0.3.3 → lightly_studio-0.4.0}/src/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-0.4.0/src/lightly_studio/api/routes/api/metadata.py

@@ -0,0 +1,91 @@
+"""This module contains the API routes for managing datasets."""
+
+from __future__ import annotations
+
+from typing import List
+from uuid import UUID
+
+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,
+)
+
+metadata_router = APIRouter(prefix="/datasets/{dataset_id}", tags=["metadata"])
+
+
+@metadata_router.get("/metadata/info", response_model=List[MetadataInfoView])
+def get_metadata_info(
+    session: SessionDep,
+    dataset_id: Annotated[UUID, Path(title="Dataset Id")],
+) -> list[MetadataInfoView]:
+    """Get all metadata keys and their schema for a dataset.
+
+    Args:
+        session: The database session.
+        dataset_id: The ID of the dataset.
+
+    Returns:
+        List of metadata info objects with name, type, and optionally min/max values
+        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-0.3.3 → lightly_studio-0.4.0}/src/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()