Lig 9937 video frame curation

.gitignore

@@ -15,3 +15,4 @@ dataset_examples_studio
 
 .claude/settings.local.json
 .claude/settings.json
+/video_frame_curation_frames

lightly_studio/src/lightly_studio/core/image/add_frames_from_video.py

@@ -0,0 +1,115 @@
+"""Extract frames from a video to image files on disk.
+
+Used by ``ImageDataset.add_frames_from_videos`` to support image-frame workflows where the
+frames originate from videos. Frames are decoded with PyAV, optionally sub-sampled to a
+target frame rate, and written as JPEG files so they can be treated as a normal image dataset.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import fsspec
+from av import container
+
+from lightly_studio.core.video.add_videos import (
+    VIDEO_EXTENSIONS,
+    _configure_stream_threading,
+    _get_frame_rotation_deg,
+)
+
+__all__ = ["VIDEO_EXTENSIONS", "extract_frames_to_dir"]
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_VIDEO_CHANNEL = 0
+# JPEG quality for the extracted frames.
+JPEG_QUALITY = 95
+# Guards against dropping a frame that sits exactly on the capture boundary due to
+# floating point rounding of the decoded timestamp.
+_TIMESTAMP_EPSILON_S = 1e-9
+
+
+def extract_frames_to_dir(
+    video_path: str,
+    extract_dir: Path,
+    fps: float | None = None,
+    video_channel: int = DEFAULT_VIDEO_CHANNEL,
+    num_decode_threads: int | None = None,
+) -> list[str]:
+    """Decode a video and write (sub-sampled) frames as JPEG images to a directory.
+
+    Args:
+        video_path: Local or remote (e.g. ``s3://``) path to the video. Opened via fsspec,
+            so the same cloud protocols as the rest of the SDK are supported.
+        extract_dir: Local directory the frames are written to. Created if it does not exist.
+        fps: Target frames per second. If ``None`` or ``<= 0``, every decoded frame is kept.
+            Otherwise a frame is kept whenever at least ``1 / fps`` seconds have passed since
+            the previously kept frame (timestamp based, so it also works for variable frame
+            rate videos).
+        video_channel: Index of the video stream to decode.
+        num_decode_threads: Optional override for the number of FFmpeg decode threads. If
+            omitted, the available CPU cores - 1 (max 16) are used.
+
+    Returns:
+        The list of written frame file paths, in decode order. Frames are named
+        ``{video_stem}_{frame_number:06d}.jpg`` where ``frame_number`` is the decode index.
+
+    # TODO(malte, 06/2026): Frames are materialized to local disk, which assumes disk space
+    #   is available. Add a disk-free / streaming path (decode -> embed -> discard) so the
+    #   workflow needs no extra storage.
+    """
+    extract_dir = Path(extract_dir)
+    extract_dir.mkdir(parents=True, exist_ok=True)
+    video_stem = Path(video_path).stem
+
+    capture_interval_s = 1.0 / fps if fps and fps > 0 else None
+    next_capture_time_s = 0.0
+    written_paths: list[str] = []
+
+    fs, fs_path = fsspec.core.url_to_fs(url=video_path)
+    with fs.open(path=fs_path, mode="rb") as video_file:
+        video_container = container.open(file=video_file)
+        try:
+            video_stream = video_container.streams.video[video_channel]
+            _configure_stream_threading(
+                video_stream=video_stream, num_decode_threads=num_decode_threads
+            )
+            time_base = video_stream.time_base
+
+            for frame_number, frame in enumerate(video_container.decode(video_stream)):
+                if frame.pts is not None and time_base is not None:
+                    frame_timestamp_s = float(frame.pts * time_base)
+                elif frame.time is not None:
+                    # Fallback to frame.time when pts/time_base are unavailable.
+                    frame_timestamp_s = frame.time
+                else:
+                    frame_timestamp_s = float(frame_number)
+
+                if (
+                    capture_interval_s is not None
+                    and frame_timestamp_s < next_capture_time_s - _TIMESTAMP_EPSILON_S
+                ):
+                    continue
+                if capture_interval_s is not None:
+                    next_capture_time_s = frame_timestamp_s + capture_interval_s
+
+                image = frame.to_image()
+                rotation_deg = _get_frame_rotation_deg(frame=frame)
+                if rotation_deg:
+                    # PIL rotates counter-clockwise by the given angle, which matches the
+                    # rotation the frame-serving code applies (see video_frames_media.py).
+                    image = image.rotate(rotation_deg, expand=True)
+
+                frame_path = extract_dir / f"{video_stem}_{frame_number:06d}.jpg"
+                image.save(frame_path, quality=JPEG_QUALITY)
+                written_paths.append(str(frame_path))
+        finally:
+            video_container.close()
+
+    # Logged at debug level so it does not interleave with a progress bar over many videos.
+    logger.debug(
+        "Extracted %d frames from %s into %s.", len(written_paths), video_path, extract_dir
+    )
+    return written_paths

lightly_studio/src/lightly_studio/core/image/image_dataset.py

@@ -23,10 +23,11 @@
     ObjectDetectionInput,
 )
 from sqlmodel import Session
+from tqdm import tqdm
 
 from lightly_studio.core.dataset import BaseSampleDataset
 from lightly_studio.core.dataset_query.dataset_query import DatasetQuery
-from lightly_studio.core.image import add_annotations, add_images
+from lightly_studio.core.image import add_annotations, add_frames_from_video, add_images
 from lightly_studio.core.image.image_sample import ImageSample
 from lightly_studio.dataset import fsspec_lister
 from lightly_studio.dataset.embedding_manager import EmbeddingManagerProvider
@@ -56,6 +57,7 @@ class ImageDataset(BaseSampleDataset[ImageSample]):
     Samples can be added to the dataset using various methods:
     ```python
     dataset.add_images_from_path(...)
+    dataset.add_frames_from_videos(...)
     dataset.add_samples_from_yolo(...)
     dataset.add_samples_from_coco(...)
     dataset.add_samples_from_coco_caption(...)
@@ -179,6 +181,98 @@ def add_images_from_path(
                 sample_ids=created_sample_ids,
             )
 
+    def add_frames_from_videos(
+        self,
+        videos_path: PathLike,
+        extract_dir: PathLike,
+        fps: float | None = None,
+        embed: bool = True,
+    ) -> list[UUID]:
+        """Extract frames from all videos under a path and add them as images.
+
+        This supports image-frame workflows where the frames originate from videos: the
+        videos under ``videos_path`` are decoded, their frames written to ``extract_dir`` as
+        JPEG images, and added to the dataset as a normal image dataset (no video dataset is
+        created). Each video's frames are always tagged with the video file stem, which
+        preserves the frame-to-video relationship.
+
+        Videos whose stem tag already exists in the dataset are skipped. This makes repeated
+        runs against a persistent dataset incremental and idempotent (mirroring how video
+        loading skips already-loaded paths), so it can be called periodically as new videos
+        arrive without re-extracting old ones.
+
+        Args:
+            videos_path: Local or remote (e.g. ``s3://``) directory, glob, or single video
+                file. All contained video files are processed.
+            extract_dir: Local directory the extracted frames are written to (one
+                subdirectory per video). Created if it does not exist. This assumes local
+                disk space is available.
+            fps: Target frames per second. If ``None`` or ``<= 0``, every decoded frame is
+                added. Otherwise frames are sub-sampled to roughly ``fps`` frames per second.
+            embed: If True, generate embeddings for the newly added frames.
+
+        Returns:
+            A list of UUIDs of the created frame samples across all processed videos.
+        """
+        video_paths = list(
+            fsspec_lister.iter_files_from_path(
+                path=str(videos_path),
+                allowed_extensions=add_frames_from_video.VIDEO_EXTENSIONS,
+            )
+        )
+        logger.info(f"Found {len(video_paths)} videos in {videos_path}.")
+
+        extract_dir = Path(extract_dir)
+        created_sample_ids: list[UUID] = []
+        for video_path in tqdm(video_paths, desc="Extracting frames from videos", unit=" video"):
+            video_stem = Path(video_path).stem
+
+            # Skip videos that have already been processed into this dataset.
+            if (
+                tag_resolver.get_by_name(
+                    session=self.session,
+                    tag_name=video_stem,
+                    collection_id=self.collection_id,
+                )
+                is not None
+            ):
+                logger.info(f"Skipping already-processed video {video_path}.")
+                continue
+
+            frame_paths = add_frames_from_video.extract_frames_to_dir(
+                video_path=video_path,
+                extract_dir=extract_dir / video_stem,
+                fps=fps,
+            )
+
+            video_sample_ids = add_images.load_into_dataset_from_paths(
+                session=self.session,
+                root_collection_id=self.collection_id,
+                image_paths=frame_paths,
+                show_progress=False,
+            )
+            if video_sample_ids:
+                video_tag = tag_resolver.get_or_create_sample_tag_by_name(
+                    session=self.session,
+                    collection_id=self.collection_id,
+                    tag_name=video_stem,
+                )
+                tag_resolver.add_sample_ids_to_tag_id(
+                    session=self.session,
+                    tag_id=video_tag.tag_id,
+                    sample_ids=video_sample_ids,
+                )
+            created_sample_ids.extend(video_sample_ids)
+
+        if embed:
+            _generate_embeddings_image(
+                session=self.session,
+                collection_id=self.collection_id,
+                sample_ids=created_sample_ids,
+            )
+
+        return created_sample_ids
+
     def add_annotations_from_labelformat(
         self,
         input_labels: ObjectDetectionInput | InstanceSegmentationInput,

lightly_studio/tests/core/image/test_image_dataset__frames_from_videos.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from lightly_studio import ImageDataset
+from tests.resolvers.video.helpers import create_video_file
+
+
+def test_add_frames_from_videos__all_frames(
+    patch_collection: None,  # noqa: ARG001
+    tmp_path: Path,
+) -> None:
+    videos_dir = tmp_path / "videos"
+    create_video_file(
+        output_path=videos_dir / "clip.mp4", width=64, height=64, num_frames=30, fps=10
+    )
+    extract_dir = tmp_path / "frames"
+
+    dataset = ImageDataset.create(name="test_dataset")
+    created_ids = dataset.add_frames_from_videos(
+        videos_path=videos_dir, extract_dir=extract_dir, fps=None, embed=False
+    )
+
+    # All decoded frames are extracted and added.
+    assert len(created_ids) == 30
+    assert len(dataset.query().to_list()) == 30
+    # Frames are written to disk under a per-video subdirectory.
+    assert len(list((extract_dir / "clip").glob("*.jpg"))) == 30
+
+
+def test_add_frames_from_videos__fps_subsamples(
+    patch_collection: None,  # noqa: ARG001
+    tmp_path: Path,
+) -> None:
+    # 30 frames at 10 fps == 3 seconds of footage.
+    videos_dir = tmp_path / "videos"
+    create_video_file(
+        output_path=videos_dir / "clip.mp4", width=64, height=64, num_frames=30, fps=10
+    )
+
+    dataset_low = ImageDataset.create(name="low_fps")
+    low_ids = dataset_low.add_frames_from_videos(
+        videos_path=videos_dir, extract_dir=tmp_path / "low", fps=2, embed=False
+    )
+
+    dataset_high = ImageDataset.create(name="high_fps")
+    high_ids = dataset_high.add_frames_from_videos(
+        videos_path=videos_dir, extract_dir=tmp_path / "high", fps=5, embed=False
+    )
+
+    # Sub-sampling keeps fewer frames, and a higher fps keeps more of them.
+    assert len(low_ids) < len(high_ids) <= 30
+    # ~2 fps over 3 seconds -> roughly 6 frames (allow rounding slack).
+    assert 5 <= len(low_ids) <= 7
+
+
+def test_add_frames_from_videos__tags_with_video_stem(
+    patch_collection: None,  # noqa: ARG001
+    tmp_path: Path,
+) -> None:
+    videos_dir = tmp_path / "videos"
+    create_video_file(
+        output_path=videos_dir / "my_clip.mp4", width=64, height=64, num_frames=4, fps=2
+    )
+
+    dataset = ImageDataset.create(name="test_dataset")
+    dataset.add_frames_from_videos(
+        videos_path=videos_dir, extract_dir=tmp_path / "frames", fps=None, embed=False
+    )
+
+    samples = dataset.query().to_list()
+    assert samples
+    assert all("my_clip" in sample.tags for sample in samples)
+
+
+def test_add_frames_from_videos__multiple_videos(
+    patch_collection: None,  # noqa: ARG001
+    tmp_path: Path,
+) -> None:
+    videos_dir = tmp_path / "videos"
+    create_video_file(
+        output_path=videos_dir / "clip_a.mp4", width=64, height=64, num_frames=4, fps=2
+    )
+    create_video_file(
+        output_path=videos_dir / "clip_b.mp4", width=64, height=64, num_frames=6, fps=2
+    )
+
+    dataset = ImageDataset.create(name="test_dataset")
+    created_ids = dataset.add_frames_from_videos(
+        videos_path=videos_dir, extract_dir=tmp_path / "frames", fps=None, embed=False
+    )
+
+    # Both videos are processed in a single call.
+    assert len(created_ids) == 4 + 6
+    samples = dataset.query().to_list()
+    # Each video's frames are tagged with that video's stem.
+    assert len([s for s in samples if "clip_a" in s.tags]) == 4
+    assert len([s for s in samples if "clip_b" in s.tags]) == 6
+
+
+def test_add_frames_from_videos__skips_already_processed(
+    patch_collection: None,  # noqa: ARG001
+    tmp_path: Path,
+) -> None:
+    videos_dir = tmp_path / "videos"
+    create_video_file(output_path=videos_dir / "clip.mp4", width=64, height=64, num_frames=4, fps=2)
+
+    dataset = ImageDataset.create(name="test_dataset")
+    first_ids = dataset.add_frames_from_videos(
+        videos_path=videos_dir, extract_dir=tmp_path / "frames", fps=None, embed=False
+    )
+    assert len(first_ids) == 4
+
+    # The video already has its stem tag, so a second run adds nothing.
+    second_ids = dataset.add_frames_from_videos(
+        videos_path=videos_dir, extract_dir=tmp_path / "frames_2", fps=None, embed=False
+    )
+    assert second_ids == []
+    assert len(dataset.query().to_list()) == 4
+
+
+def test_add_frames_from_videos__embeds(
+    patch_collection: None,  # noqa: ARG001
+    tmp_path: Path,
+) -> None:
+    videos_dir = tmp_path / "videos"
+    create_video_file(output_path=videos_dir / "clip.mp4", width=64, height=64, num_frames=4, fps=2)
+
+    dataset = ImageDataset.create(name="test_dataset")
+    dataset.add_frames_from_videos(
+        videos_path=videos_dir, extract_dir=tmp_path / "frames", fps=None, embed=True
+    )
+
+    samples = dataset.query().to_list()
+    assert samples
+    assert all(len(sample.sample_table.embeddings) == 1 for sample in samples)

lightly_studio_view/e2e/general/samples-grid.e2e-test.ts

@@ -300,7 +300,7 @@ test('Sampling shows error toast when tag already exists', async ({ page, sample
         ).toBeVisible({ timeout: 10000 });
 
         // Close the sampling dialog to reset for the next iteration.
-        await pressButton(page, 'sampling-dialog-cancel');
+        await pressButton(page, 'selection-dialog-cancel');
     }
 });