Multimodal: Link caption to Sample instead of Image (#36)

Author: michal-lightly

lightly_studio/src/lightly_studio/api/routes/api/sample.py

@@ -79,7 +79,7 @@ def read_samples(
                 sample_id=image.sample_id,
                 dataset_id=image.dataset_id,
                 annotations=image.annotations,
-                captions=image.captions,
+                captions=image.sample.captions,
                 tags=image.sample.tags,
                 metadata_dict=image.sample.metadata_dict,
                 width=image.width,
@@ -128,7 +128,7 @@ def read_sample(
         sample_id=image.sample_id,
         dataset_id=image.dataset_id,
         annotations=image.annotations,
-        captions=image.captions,
+        captions=image.sample.captions,
         tags=image.sample.tags,
         metadata_dict=image.sample.metadata_dict,
         width=image.width,

lightly_studio/src/lightly_studio/models/caption.py

@@ -10,7 +10,7 @@
 from sqlmodel import Field, Relationship, SQLModel
 
 if TYPE_CHECKING:
-    from lightly_studio.models.image import ImageTable
+    from lightly_studio.models.sample import SampleTable
 
 
 class CaptionTable(SQLModel, table=True):
@@ -22,9 +22,9 @@ class CaptionTable(SQLModel, table=True):
 
     caption_id: UUID = Field(default_factory=uuid4, primary_key=True)
     dataset_id: UUID = Field(foreign_key="dataset.dataset_id")
-    sample_id: UUID = Field(foreign_key="image.sample_id")
+    sample_id: UUID = Field(foreign_key="sample.sample_id")
 
-    sample: Mapped[Optional["ImageTable"]] = Relationship(
+    sample: Mapped["SampleTable"] = Relationship(
         back_populates="captions",
         sa_relationship_kwargs={"lazy": "select"},
     )
@@ -40,11 +40,10 @@ class CaptionCreate(SQLModel):
     text: str
 
 
-class CaptionImageView(SQLModel):
+class CaptionSampleView(SQLModel):
     """Sample class for caption view."""
 
-    file_path_abs: str
-    file_name: str
+    # TODO(Michal, 10/2025): Remove this class and use CaptionView instead.
     dataset_id: UUID
     sample_id: UUID
 
@@ -61,7 +60,7 @@ class CaptionView(SQLModel):
 class CaptionDetailsView(CaptionView):
     """Response model for caption."""
 
-    sample: CaptionImageView
+    sample: CaptionSampleView
 
 
 class CaptionsListView(BaseModel):

lightly_studio/src/lightly_studio/models/image.py

@@ -16,16 +16,12 @@
     from lightly_studio.models.annotation.annotation_base import (
         AnnotationBaseTable,
     )
-    from lightly_studio.models.caption import CaptionTable
     from lightly_studio.models.metadata import (
-        SampleMetadataTable,
         SampleMetadataView,
     )
     from lightly_studio.models.sample import SampleTable
 else:
     AnnotationBaseTable = object
-    CaptionTable = object
-    SampleMetadataTable = object
     SampleTable = object
     SampleMetadataView = object
 
@@ -65,9 +61,6 @@ class ImageTable(ImageBase, table=True):
     annotations: Mapped[List["AnnotationBaseTable"]] = Relationship(
         back_populates="sample",
     )
-    captions: Mapped[List["CaptionTable"]] = Relationship(
-        back_populates="sample",
-    )
 
     sample: Mapped["SampleTable"] = Relationship()
 
@@ -96,13 +89,13 @@ class ImageViewTag(SQLModel):
     sample_id: UUID
     dataset_id: UUID
     annotations: List["AnnotationView"]
-    captions: List[CaptionView] = []
     width: int
     height: int
 
     # TODO(Michal, 10/2025): Add SampleView to ImageView, don't expose these fields directly.
     tags: List[ImageViewTag]
     metadata_dict: Optional["SampleMetadataView"] = None
+    captions: List[CaptionView] = []
 
 
 class ImageViewsWithCount(BaseModel):

lightly_studio/src/lightly_studio/models/sample.py

@@ -10,6 +10,7 @@
 from lightly_studio.resolvers import metadata_resolver
 
 if TYPE_CHECKING:
+    from lightly_studio.models.caption import CaptionTable, CaptionView
     from lightly_studio.models.metadata import (
         SampleMetadataTable,
         SampleMetadataView,
@@ -21,6 +22,8 @@
     SampleEmbeddingTable = object
     SampleMetadataTable = object
     SampleMetadataView = object
+    CaptionTable = object
+    CaptionView = object
 
 
 class SampleTagLinkTable(SQLModel, table=True):
@@ -58,6 +61,7 @@ class SampleTable(SampleBase, table=True):
     )
     embeddings: Mapped[List["SampleEmbeddingTable"]] = Relationship(back_populates="sample")
     metadata_dict: "SampleMetadataTable" = Relationship(back_populates="sample")
+    captions: Mapped[List["CaptionTable"]] = Relationship(back_populates="sample")
 
     # TODO(Michal, 9/2025): Remove this function in favour of Sample.metadata.
     def __getitem__(self, key: str) -> Any:
@@ -118,3 +122,4 @@ class SampleView(SampleBase):
 
     tags: List["TagTable"] = []
     metadata_dict: Optional["SampleMetadataView"] = None
+    captions: List[CaptionView] = []

lightly_studio/src/lightly_studio/resolvers/image_resolver.py

@@ -156,8 +156,8 @@ def get_all_by_dataset_id(  # noqa: PLR0913
                 joinedload(SampleTable.tags),
                 # Ignore type checker error below as it's a false positive caused by TYPE_CHECKING.
                 joinedload(SampleTable.metadata_dict),  # type: ignore[arg-type]
+                selectinload(SampleTable.captions),
             ),
-            selectinload(ImageTable.captions),
         )
         .where(ImageTable.dataset_id == dataset_id)
     )

lightly_studio/tests/core/test_add_samples.py

@@ -13,7 +13,7 @@
 from sqlmodel import Session
 
 from lightly_studio.core import add_samples
-from lightly_studio.models.image import ImageTable
+from lightly_studio.models.sample import SampleTable
 from lightly_studio.resolvers import caption_resolver, image_resolver
 from tests.helpers_resolvers import create_dataset
 
@@ -123,13 +123,13 @@ def test_load_into_dataset_from_coco_captions(db_session: Session, tmp_path: Pat
     assert captions_result.next_cursor is None
     # Collect all the filename x caption pairs and assert they are as expected
     assert {
-        (c.sample.file_name, c.text)
+        (c.sample.sample_id, c.text)
         for c in captions_result.captions
-        if isinstance(c.sample, ImageTable)
+        if isinstance(c.sample, SampleTable)
     } == {
-        ("image1.jpg", "Caption 1 of image 1"),
-        ("image1.jpg", "Caption 2 of image 1"),
-        ("image2.jpg", "Caption 1 of image 2"),
+        (samples[0].sample_id, "Caption 1 of image 1"),
+        (samples[0].sample_id, "Caption 2 of image 1"),
+        (samples[1].sample_id, "Caption 1 of image 2"),
     }
 
 

lightly_studio/tests/core/test_dataset__coco_caption.py

@@ -7,7 +7,6 @@
 from PIL import Image
 
 from lightly_studio import Dataset
-from lightly_studio.models.image import ImageTable
 from lightly_studio.resolvers import caption_resolver
 
 
@@ -32,7 +31,6 @@ def test_add_samples_from_coco_caption__details_valid(
         )
         assert dataset.name == "test_dataset"
         samples = dataset._inner.get_samples()
-        samples = sorted(samples, key=lambda sample: sample.file_path_abs)
 
         assert len(samples) == 2
         assert {s.file_name for s in samples} == {"image1.jpg", "image2.jpg"}
@@ -45,11 +43,11 @@ def test_add_samples_from_coco_caption__details_valid(
         assert len(captions_result.captions) == 3
         assert captions_result.total_count == 3
         assert captions_result.next_cursor is None
+
         # Collect all the filename x caption pairs and assert they are as expected
+        sample_id_to_file_path = {s.sample.sample_id: s.file_name for s in samples}
         assert {
-            (c.sample.file_name, c.text)
-            for c in captions_result.captions
-            if isinstance(c.sample, ImageTable)
+            (sample_id_to_file_path[c.sample.sample_id], c.text) for c in captions_result.captions
         } == {
             ("image1.jpg", "Caption 1 of image 1"),
             ("image1.jpg", "Caption 2 of image 1"),