Improve video grid loading speed (#183)

Author: horatiualmasan

lightly_studio/src/lightly_studio/api/routes/api/video.py

@@ -10,17 +10,16 @@
 from lightly_studio.db_manager import SessionDep
 from lightly_studio.models.video import VideoTable, VideoView, VideoViewsWithCount
 from lightly_studio.resolvers import video_resolver
-from lightly_studio.resolvers.video_resolver import VideosWithCount
 
 video_router = APIRouter(prefix="/datasets/{dataset_id}/video", tags=["video"])
 
 
-@video_router.get("/", response_model=VideoViewsWithCount)
+@video_router.get("/")
 def get_all_videos(
     session: SessionDep,
     dataset_id: Annotated[UUID, Path(title="Dataset Id")],
     pagination: Annotated[PaginatedWithCursor, Depends()],
-) -> VideosWithCount:
+) -> VideoViewsWithCount:
     """Retrieve a list of all videos for a given dataset ID with pagination.
 
     Args:

lightly_studio/src/lightly_studio/examples/example_video_annotations.py

@@ -92,7 +92,9 @@ def load_annotations(session: Session, dataset_id: UUID, annotations_path: Path)
     Temporarily use internal add_samples API until labelformat supports videos natively.
     """
     print("Loading video annotations...")
-    videos = video_resolver.get_all_by_dataset_id(session=session, dataset_id=dataset_id).samples
+    videos = video_resolver.get_all_by_dataset_id_with_frames(
+        session=session, dataset_id=dataset_id
+    )
     video_name_to_video = {video.file_name: video for video in videos}
     yvis_input = YouTubeVISObjectDetectionInput(input_file=annotations_path)
     label_map = add_samples._create_label_map(  # noqa: SLF001

lightly_studio/src/lightly_studio/models/video.py

@@ -60,13 +60,13 @@ class VideoView(SQLModel):
 
     width: int
     height: int
-    duration_s: float
+    duration_s: Optional[float] = None
     fps: float
     file_name: str
     file_path_abs: str
     sample_id: UUID
     sample: SampleView
-    frames: List["FrameView"] = []
+    frame: Optional["FrameView"] = None
 
 
 class VideoViewsWithCount(BaseModel):

lightly_studio/src/lightly_studio/resolvers/video_resolver/__init__.py

@@ -3,15 +3,15 @@
 from lightly_studio.resolvers.video_resolver.create_many import create_many
 from lightly_studio.resolvers.video_resolver.filter_new_paths import filter_new_paths
 from lightly_studio.resolvers.video_resolver.get_all_by_dataset_id import (
-    VideosWithCount,
     get_all_by_dataset_id,
+    get_all_by_dataset_id_with_frames,
 )
 from lightly_studio.resolvers.video_resolver.get_by_id import get_by_id
 
 __all__ = [
-    "VideosWithCount",
     "create_many",
     "filter_new_paths",
     "get_all_by_dataset_id",
+    "get_all_by_dataset_id_with_frames",
     "get_by_id",
 ]

lightly_studio/src/lightly_studio/resolvers/video_resolver/get_all_by_dataset_id.py

@@ -5,32 +5,71 @@
 from collections.abc import Sequence
 from uuid import UUID
 
-from pydantic import BaseModel
+from sqlalchemy import and_
+from sqlalchemy.orm import joinedload, selectinload
 from sqlmodel import Session, col, func, select
 
 from lightly_studio.api.routes.api.validators import Paginated
-from lightly_studio.models.sample import SampleTable
-from lightly_studio.models.video import VideoTable
-
-
-class VideosWithCount(BaseModel):
-    """Result of getting all samples."""
-
-    samples: Sequence[VideoTable]
-    total_count: int
-    next_cursor: int | None = None
+from lightly_studio.models.sample import SampleTable, SampleView
+from lightly_studio.models.video import (
+    FrameView,
+    VideoFrameTable,
+    VideoTable,
+    VideoView,
+    VideoViewsWithCount,
+)
 
 
 def get_all_by_dataset_id(
     session: Session,
     dataset_id: UUID,
     pagination: Paginated | None = None,
     sample_ids: list[UUID] | None = None,
-) -> VideosWithCount:
+) -> VideoViewsWithCount:
     """Retrieve samples for a specific dataset with optional filtering."""
+    # Subquery to find the minimum frame_number for each video
+    min_frame_subquery = (
+        select(
+            VideoFrameTable.parent_sample_id,
+            func.min(col(VideoFrameTable.frame_number)).label("min_frame_number"),
+        )
+        .group_by(col(VideoFrameTable.parent_sample_id))
+        .subquery()
+    )
+    # TODO(Horatiu, 11/2025): Check if it is possible to optimize this query.
+    # Query to get videos with their first frame (frame with min frame_number)
+    # First join the subquery to VideoTable, then join VideoFrameTable
     samples_query = (
-        select(VideoTable).join(VideoTable.sample).where(SampleTable.dataset_id == dataset_id)
+        select(VideoTable, VideoFrameTable)
+        .join(VideoTable.sample)
+        .outerjoin(
+            min_frame_subquery,
+            min_frame_subquery.c.parent_sample_id == VideoTable.sample_id,
+        )
+        .outerjoin(
+            VideoFrameTable,
+            and_(
+                col(VideoFrameTable.parent_sample_id) == col(VideoTable.sample_id),
+                col(VideoFrameTable.frame_number) == min_frame_subquery.c.min_frame_number,
+            ),
+        )
+        .where(SampleTable.dataset_id == dataset_id)
+        .options(
+            selectinload(VideoFrameTable.sample).options(
+                joinedload(SampleTable.tags),
+                # Ignore type checker error - false positive from TYPE_CHECKING.
+                joinedload(SampleTable.metadata_dict),  # type: ignore[arg-type]
+                selectinload(SampleTable.captions),
+            ),
+            selectinload(VideoTable.sample).options(
+                joinedload(SampleTable.tags),
+                # Ignore type checker error - false positive from TYPE_CHECKING.
+                joinedload(SampleTable.metadata_dict),  # type: ignore[arg-type]
+                selectinload(SampleTable.captions),
+            ),
+        )
     )
+
     total_count_query = (
         select(func.count())
         .select_from(VideoTable)
@@ -54,8 +93,55 @@ def get_all_by_dataset_id(
     if pagination and pagination.offset + pagination.limit < total_count:
         next_cursor = pagination.offset + pagination.limit
 
-    return VideosWithCount(
-        samples=session.exec(samples_query).all(),
+    # Fetch videos with their first frames and convert to VideoView
+    results = session.exec(samples_query).all()
+    video_views = [
+        _convert_video_table_to_view(video=video, first_frame=first_frame)
+        for video, first_frame in results
+    ]
+
+    return VideoViewsWithCount(
+        samples=video_views,
         total_count=total_count,
         next_cursor=next_cursor,
     )
+
+
+# TODO(Horatiu, 11/2025): This should be deleted when we have proper way of getting all frames for
+# a video.
+def get_all_by_dataset_id_with_frames(
+    session: Session,
+    dataset_id: UUID,
+) -> Sequence[VideoTable]:
+    """Retrieve video table with all the samples."""
+    samples_query = (
+        select(VideoTable).join(VideoTable.sample).where(SampleTable.dataset_id == dataset_id)
+    )
+    samples_query = samples_query.order_by(col(VideoTable.file_path_abs).asc())
+    return session.exec(samples_query).all()
+
+
+def _convert_video_table_to_view(
+    video: VideoTable, first_frame: VideoFrameTable | None
+) -> VideoView:
+    """Convert VideoTable to VideoView with only the first frame."""
+    first_frame_view = None
+    if first_frame:
+        first_frame_view = FrameView(
+            frame_number=first_frame.frame_number,
+            frame_timestamp_s=first_frame.frame_timestamp_s,
+            sample_id=first_frame.sample_id,
+            sample=SampleView.model_validate(first_frame.sample),
+        )
+
+    return VideoView(
+        width=video.width,
+        height=video.height,
+        duration_s=video.duration_s,
+        fps=video.fps,
+        file_name=video.file_name,
+        file_path_abs=video.file_path_abs,
+        sample_id=video.sample_id,
+        sample=SampleView.model_validate(video.sample),
+        frame=first_frame_view,
+    )

lightly_studio/tests/api/routes/api/test_frame.py

@@ -8,8 +8,7 @@
 from lightly_studio.api.routes.api.status import HTTP_STATUS_OK
 from lightly_studio.models.dataset import SampleType
 from tests.helpers_resolvers import create_dataset
-from tests.resolvers.video_frame_resolver.helpers import create_video_with_frames
-from tests.resolvers.video_resolver.helpers import VideoStub
+from tests.resolvers.video.helpers import VideoStub, create_video_with_frames
 
 
 def test_get_all_frames(

lightly_studio/tests/api/routes/api/test_video.py

@@ -7,7 +7,7 @@
 from lightly_studio.models.dataset import SampleType
 from lightly_studio.resolvers import video_resolver
 from tests.helpers_resolvers import create_dataset
-from tests.resolvers.video_resolver.helpers import VideoStub, create_videos
+from tests.resolvers.video.helpers import VideoStub, create_videos
 
 
 def test_get_all_videos(test_client: TestClient, db_session: Session) -> None:

lightly_studio/tests/core/test_add_videos.py

@@ -57,6 +57,8 @@ def test_load_into_dataset_from_paths(db_session: Session, tmp_path: Path) -> No
     video = videos[0]
     assert video.file_name == "test_video_0.mp4"
     assert video.file_path_abs == str(second_video_path)
+    assert video.frame is not None
+    assert video.frame.frame_number == 0
     video = videos[1]
     assert video.file_name == "test_video_1.mp4"
     assert video.file_path_abs == str(first_video_path)

lightly_studio/tests/resolvers/test_annotation_resolver.py

@@ -30,7 +30,7 @@
     create_image,
     create_tag,
 )
-from tests.resolvers.video_resolver.helpers import VideoStub, create_videos
+from tests.resolvers.video.helpers import VideoStub, create_videos
 
 
 @dataclass