Add image_resolver.get_for_export with collection filter support

lightly_studio/src/lightly_studio/resolvers/image_resolver/__init__.py

@@ -11,6 +11,7 @@
 )
 from lightly_studio.resolvers.image_resolver.get_by_id import get_by_id
 from lightly_studio.resolvers.image_resolver.get_dimension_bounds import get_dimension_bounds
+from lightly_studio.resolvers.image_resolver.get_for_export import get_for_export
 from lightly_studio.resolvers.image_resolver.get_many_by_id import get_many_by_id
 from lightly_studio.resolvers.image_resolver.get_sample_ids import (
     build_sample_ids_query,
@@ -30,6 +31,7 @@
     "get_all_by_collection_id",
     "get_by_id",
     "get_dimension_bounds",
+    "get_for_export",
     "get_many_by_id",
     "get_sample_ids",
     "get_sample_ids_by_paths",

lightly_studio/src/lightly_studio/resolvers/image_resolver/get_all_by_collection_id.py

@@ -100,13 +100,20 @@ def get_all_by_collection_id(  # noqa: PLR0913
     """Retrieve samples for a specific collection with optional filtering."""
     # Resolve any embedding-plot region selection to concrete sample ids on the filter before the
     # query is built (the point-in-polygon test needs the session, which `apply` lacks).
-    sample_filter = filters.sample_filter if filters is not None else None
-    if sample_filter is not None and sample_filter.embedding_region is not None:
-        sample_filter.region_sample_ids = embedding_region_resolver.get_sample_ids_in_region(
+    if (
+        filters is not None
+        and filters.sample_filter is not None
+        and filters.sample_filter.embedding_region is not None
+    ):
+        region_sample_ids = embedding_region_resolver.get_sample_ids_in_region(
             session=session,
             collection_id=collection_id,
-            region=sample_filter.embedding_region,
+            region=filters.sample_filter.embedding_region,
         )
+        resolved_sample_filter = filters.sample_filter.model_copy(
+            update={"region_sample_ids": region_sample_ids}
+        )
+        filters = filters.model_copy(update={"sample_filter": resolved_sample_filter})
 
     embedding_model_id, distance_expr = get_distance_expression(
         session=session,

lightly_studio/src/lightly_studio/resolvers/image_resolver/get_for_export.py

@@ -0,0 +1,56 @@
+"""Implementation of get_for_export function for images."""
+
+from __future__ import annotations
+
+from collections.abc import Generator
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.core.image.image_sample import ImageSample
+from lightly_studio.models.image import ImageTable
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.resolvers import embedding_region_resolver
+from lightly_studio.resolvers.image_filter import ImageFilter
+
+
+def get_for_export(
+    session: Session,
+    collection_id: UUID,
+    collection_filter: ImageFilter | None,
+) -> Generator[ImageSample, None, None]:
+    """Return all images in a collection as a lazy generator of ImageSamples.
+
+    If ``collection_filter`` is provided, only images matching the filter are
+    returned. Embedding-region filters are resolved to concrete sample IDs
+    before the query is executed.
+
+    Args:
+        session: Database session.
+        collection_id: ID of the collection to export.
+        collection_filter: Optional filter to restrict which images are returned.
+
+    Returns:
+        Generator of ImageSamples for the matching images.
+    """
+    query = (
+        select(ImageTable).join(ImageTable.sample).where(SampleTable.collection_id == collection_id)
+    )
+    if collection_filter is not None:
+        sample_filter = collection_filter.sample_filter
+        if sample_filter is not None and sample_filter.embedding_region is not None:
+            region_sample_ids = embedding_region_resolver.get_sample_ids_in_region(
+                session=session,
+                collection_id=collection_id,
+                region=sample_filter.embedding_region,
+            )
+            resolved_sample_filter = sample_filter.model_copy(
+                update={"region_sample_ids": region_sample_ids}
+            )
+            collection_filter = collection_filter.model_copy(
+                update={"sample_filter": resolved_sample_filter}
+            )
+        query = query.where(
+            col(SampleTable.sample_id).in_(collection_filter.build_sample_ids_query(collection_id))
+        )
+    return (ImageSample(row) for row in session.exec(query))

lightly_studio/tests/resolvers/image_resolver/test_get_for_export.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+from sqlmodel import Session
+
+from lightly_studio.models.embedding_region import EmbeddingRegion, Point2D
+from lightly_studio.models.two_dim_embedding import TwoDimEmbeddingTable
+from lightly_studio.resolvers import image_resolver, sample_embedding_resolver
+from lightly_studio.resolvers.image_filter import FilterDimensions, ImageFilter
+from lightly_studio.resolvers.sample_resolver.sample_filter import SampleFilter
+from tests.helpers_resolvers import (
+    create_collection,
+    create_embedding_model,
+    create_image,
+    create_sample_embedding,
+)
+
+
+def test_get_for_export__no_filter(db_session: Session) -> None:
+    collection = create_collection(session=db_session)
+    other_collection = create_collection(session=db_session)
+
+    image1 = create_image(
+        session=db_session,
+        collection_id=collection.collection_id,
+        file_path_abs="/data/img1.jpg",
+    )
+    image2 = create_image(
+        session=db_session,
+        collection_id=collection.collection_id,
+        file_path_abs="/data/img2.jpg",
+    )
+    create_image(
+        session=db_session,
+        collection_id=other_collection.collection_id,
+        file_path_abs="/data/other.jpg",
+    )
+
+    result = list(
+        image_resolver.get_for_export(
+            session=db_session,
+            collection_id=collection.collection_id,
+            collection_filter=None,
+        )
+    )
+
+    assert {s.sample_id for s in result} == {image1.sample_id, image2.sample_id}
+
+
+def test_get_for_export__with_image_filter(db_session: Session) -> None:
+    collection = create_collection(session=db_session)
+
+    small_image = create_image(
+        session=db_session,
+        collection_id=collection.collection_id,
+        file_path_abs="/data/small.jpg",
+        width=100,
+        height=100,
+    )
+    large_image = create_image(
+        session=db_session,
+        collection_id=collection.collection_id,
+        file_path_abs="/data/large.jpg",
+        width=800,
+        height=800,
+    )
+
+    result = list(
+        image_resolver.get_for_export(
+            session=db_session,
+            collection_id=collection.collection_id,
+            collection_filter=ImageFilter(width=FilterDimensions(min=500)),
+        )
+    )
+
+    assert {s.sample_id for s in result} == {large_image.sample_id}
+
+
+def test_get_for_export__with_embedding_region_filter(db_session: Session) -> None:
+    collection = create_collection(session=db_session)
+    embedding_model = create_embedding_model(
+        session=db_session,
+        collection_id=collection.collection_id,
+        embedding_dimension=3,
+    )
+
+    image_inside1 = create_image(
+        session=db_session,
+        collection_id=collection.collection_id,
+        file_path_abs="inside1.jpg",
+    )
+    create_sample_embedding(
+        session=db_session,
+        sample_id=image_inside1.sample_id,
+        embedding_model_id=embedding_model.embedding_model_id,
+        embedding=[1.0, 0.2, 0.3],
+    )
+    image_inside2 = create_image(
+        session=db_session,
+        collection_id=collection.collection_id,
+        file_path_abs="inside2.jpg",
+    )
+    create_sample_embedding(
+        session=db_session,
+        sample_id=image_inside2.sample_id,
+        embedding_model_id=embedding_model.embedding_model_id,
+        embedding=[2.0, 0.2, 0.3],
+    )
+    image_outside = create_image(
+        session=db_session,
+        collection_id=collection.collection_id,
+        file_path_abs="outside.jpg",
+    )
+    create_sample_embedding(
+        session=db_session,
+        sample_id=image_outside.sample_id,
+        embedding_model_id=embedding_model.embedding_model_id,
+        embedding=[3.0, 0.2, 0.3],
+    )
+
+    # Seed 2D coordinates: inside1 and inside2 within (0,0)-(10,10), outside at (100, 100).
+    cache_key, sample_ids_in_order = sample_embedding_resolver.get_hash_by_collection_id(
+        session=db_session,
+        collection_id=collection.collection_id,
+        embedding_model_id=embedding_model.embedding_model_id,
+    )
+    inside1_i = sample_ids_in_order.index(image_inside1.sample_id)
+    inside2_i = sample_ids_in_order.index(image_inside2.sample_id)
+    outside_i = sample_ids_in_order.index(image_outside.sample_id)
+    x = [0.0, 0.0, 0.0]
+    y = [0.0, 0.0, 0.0]
+    x[inside1_i] = 1.0
+    y[inside1_i] = 1.0
+    x[inside2_i] = 5.0
+    y[inside2_i] = 5.0
+    x[outside_i] = 100.0
+    y[outside_i] = 100.0
+    db_session.add(TwoDimEmbeddingTable(hash=cache_key, x=x, y=y))
+    db_session.commit()
+
+    region = EmbeddingRegion(
+        polygon=[
+            Point2D(x=0, y=0),
+            Point2D(x=10, y=0),
+            Point2D(x=10, y=10),
+            Point2D(x=0, y=10),
+        ]
+    )
+    collection_filter = ImageFilter(sample_filter=SampleFilter(embedding_region=region))
+
+    result = list(
+        image_resolver.get_for_export(
+            session=db_session,
+            collection_id=collection.collection_id,
+            collection_filter=collection_filter,
+        )
+    )
+
+    assert {s.sample_id for s in result} == {image_inside1.sample_id, image_inside2.sample_id}