diff --git a/.gitignore b/.gitignore index def3077933..bc4996e208 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,4 @@ dataset_examples_studio .claude/settings.local.json .claude/settings.json +/video_frame_curation_frames diff --git a/lightly_studio/src/lightly_studio/core/image/add_frames_from_video.py b/lightly_studio/src/lightly_studio/core/image/add_frames_from_video.py new file mode 100644 index 0000000000..7b53d47160 --- /dev/null +++ b/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 diff --git a/lightly_studio/src/lightly_studio/core/image/image_dataset.py b/lightly_studio/src/lightly_studio/core/image/image_dataset.py index a38356259f..314725f3ff 100644 --- a/lightly_studio/src/lightly_studio/core/image/image_dataset.py +++ b/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, diff --git a/lightly_studio/tests/core/image/test_image_dataset__frames_from_videos.py b/lightly_studio/tests/core/image/test_image_dataset__frames_from_videos.py new file mode 100644 index 0000000000..3c6df05383 --- /dev/null +++ b/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) diff --git a/lightly_studio_view/e2e/general/samples-grid.e2e-test.ts b/lightly_studio_view/e2e/general/samples-grid.e2e-test.ts index 01441c81ec..6af0fe27ee 100644 --- a/lightly_studio_view/e2e/general/samples-grid.e2e-test.ts +++ b/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'); } }); diff --git a/lightly_studio_view/e2e/pages/samples-page.ts b/lightly_studio_view/e2e/pages/samples-page.ts index 6dd6ac9479..5d61dde22a 100644 --- a/lightly_studio_view/e2e/pages/samples-page.ts +++ b/lightly_studio_view/e2e/pages/samples-page.ts @@ -155,14 +155,14 @@ export class SamplesPage { await this.page.getByTestId(`similarity-query-tag-${queryTagId}`).click(); } - const nSamplesInput = this.page.getByTestId('sampling-dialog-n-samples-input'); + const nSamplesInput = this.page.getByTestId('selection-dialog-n-samples-input'); await nSamplesInput.clear(); await nSamplesInput.fill(nSamples.toString()); - const tagNameInput = this.page.getByTestId('sampling-dialog-tag-name-input'); + const tagNameInput = this.page.getByTestId('selection-dialog-tag-name-input'); await tagNameInput.fill(tagName); - await pressButton(this.page, 'sampling-dialog-submit'); + await pressButton(this.page, 'selection-dialog-submit'); } async createDiversitySampling(nSamples: number, tagName: string): Promise { diff --git a/lightly_studio_view/src/lib/components/Sampling/AddStrategyButton.svelte b/lightly_studio_view/src/lib/components/Sampling/AddStrategyButton.svelte index 952d8ea57c..054483aa9b 100644 --- a/lightly_studio_view/src/lib/components/Sampling/AddStrategyButton.svelte +++ b/lightly_studio_view/src/lib/components/Sampling/AddStrategyButton.svelte @@ -6,6 +6,7 @@ interface Props { diversityDisabledReason?: string; + deduplicationDisabledReason?: string; similarityDisabledReason?: string; metadataWeightingDisabledReason?: string; classBalancingDisabledReason?: string; @@ -13,6 +14,7 @@ } let { diversityDisabledReason, + deduplicationDisabledReason, similarityDisabledReason, metadataWeightingDisabledReason, classBalancingDisabledReason, @@ -28,6 +30,7 @@ function getDisabledReason(type: StrategyType): string | undefined { if (type === 'diversity') return diversityDisabledReason; + if (type === 'deduplication') return deduplicationDisabledReason; if (type === 'similarity') return similarityDisabledReason; if (type === 'metadata_weighting') return metadataWeightingDisabledReason; if (type === 'class_balancing') return classBalancingDisabledReason; diff --git a/lightly_studio_view/src/lib/components/Sampling/SamplingCombinationDialog.svelte b/lightly_studio_view/src/lib/components/Sampling/SamplingCombinationDialog.svelte index 07b64de6c1..c9c5d6a12b 100644 --- a/lightly_studio_view/src/lib/components/Sampling/SamplingCombinationDialog.svelte +++ b/lightly_studio_view/src/lib/components/Sampling/SamplingCombinationDialog.svelte @@ -37,6 +37,7 @@ // TODO(Leonardo, 06/2026): Update once there are multiple embedding models - currently only one diversity // strategy is supported since all samples share a single embedding space. const hasDiversity = $derived($instances.some((i) => i.type === 'diversity')); + const hasDeduplication = $derived($instances.some((i) => i.type === 'deduplication')); const { tags, @@ -83,6 +84,9 @@ diversityDisabledReason={hasDiversity ? 'Only one diversity strategy can be added per selection.' : undefined} + deduplicationDisabledReason={hasDeduplication + ? 'Only one deduplication strategy can be added per selection.' + : undefined} similarityDisabledReason={isVideoCollection ? 'Not available for video collections. Similarity selection requires image embeddings.' : $tags.length === 0 @@ -108,7 +112,8 @@ annotationLabels={strategyOptions.annotationLabels} annotationSourceOptions={strategyOptions.annotationSourceOptions} metadataFieldNames={strategyOptions.metadataFieldNames} - isDuplicateDisabled={instance.type === 'diversity'} + isDuplicateDisabled={instance.type === 'diversity' || + instance.type === 'deduplication'} onRemove={() => removeStrategy(instance.id)} onDuplicate={() => duplicateStrategy(instance.id)} onUpdate={(params) => updateParams(instance.id, params)} diff --git a/lightly_studio_view/src/lib/components/Sampling/StrategyCard/StrategyCard.svelte b/lightly_studio_view/src/lib/components/Sampling/StrategyCard/StrategyCard.svelte index b9b5f40ebf..96d4ece3d2 100644 --- a/lightly_studio_view/src/lib/components/Sampling/StrategyCard/StrategyCard.svelte +++ b/lightly_studio_view/src/lib/components/Sampling/StrategyCard/StrategyCard.svelte @@ -9,6 +9,7 @@ import { STRATEGY_LABELS, + type DeduplicationParams, type MetadataWeightingParams, type SimilarityParams, type StrategyInstance, @@ -16,6 +17,7 @@ type StrategySummaryTag, type ClassBalancingParams } from '$lib/hooks/useStrategyBuilder'; + import DeduplicationForm from '../forms/DeduplicationForm/DeduplicationForm.svelte'; import MetadataWeightingForm from '../forms/MetadataWeightingForm/MetadataWeightingForm.svelte'; import SimilarityForm from '../forms/SimilarityForm/SimilarityForm.svelte'; import ClassBalancingForm from '../forms/ClassBalancingForm/ClassBalancingForm.svelte'; @@ -110,6 +112,11 @@ min={0} onUpdate={(strength) => onUpdate({ strength })} /> + {:else if instance.type === 'deduplication'} + {:else if instance.type === 'typicality'} { }); }); + describe('deduplication', () => { + it('renders the minimum distance and strength fields when expanded', () => { + const instance: StrategyInstance = { + id: 'abc', + type: 'deduplication', + params: { strength: 1, stopping_condition_minimum_distance: 0.1 }, + isExpanded: true + }; + + render(StrategyCard, { props: { ...defaultProps, instance } }); + + expect( + screen.getByTestId('strategy-deduplication-min-distance-input') + ).toBeInTheDocument(); + expect( + screen.getByTestId('strategy-deduplication-strength-input') + ).toBeInTheDocument(); + }); + + it('calls onUpdate with a new minimum distance', async () => { + const onUpdate = vi.fn(); + const instance: StrategyInstance = { + id: 'abc', + type: 'deduplication', + params: { strength: 1, stopping_condition_minimum_distance: 0.1 }, + isExpanded: true + }; + + render(StrategyCard, { props: { ...defaultProps, onUpdate, instance } }); + + await fireEvent.input(screen.getByTestId('strategy-deduplication-min-distance-input'), { + target: { value: '0.5' } + }); + + expect(onUpdate).toHaveBeenCalledWith({ stopping_condition_minimum_distance: 0.5 }); + }); + }); + describe('typicality', () => { it('renders the strength field when expanded', () => { const instance: StrategyInstance = { diff --git a/lightly_studio_view/src/lib/components/Sampling/forms/DeduplicationForm/DeduplicationForm.svelte b/lightly_studio_view/src/lib/components/Sampling/forms/DeduplicationForm/DeduplicationForm.svelte new file mode 100644 index 0000000000..405e1ba27f --- /dev/null +++ b/lightly_studio_view/src/lib/components/Sampling/forms/DeduplicationForm/DeduplicationForm.svelte @@ -0,0 +1,47 @@ + + +
+
+
+ + +
+ { + const raw = (event.currentTarget as HTMLInputElement).value; + const parsed = Number(raw); + if (raw !== '' && !Number.isNaN(parsed) && parsed > 0) { + onUpdate({ stopping_condition_minimum_distance: parsed }); + } + }} + data-testid="strategy-deduplication-min-distance-input" + class="[appearance:textfield] [&::-webkit-inner-spin-button]:appearance-none [&::-webkit-outer-spin-button]:appearance-none" + /> +
+ onUpdate({ strength })} + /> +
diff --git a/lightly_studio_view/src/lib/hooks/useStrategyBuilder/index.ts b/lightly_studio_view/src/lib/hooks/useStrategyBuilder/index.ts index 8697d8f606..59072eba93 100644 --- a/lightly_studio_view/src/lib/hooks/useStrategyBuilder/index.ts +++ b/lightly_studio_view/src/lib/hooks/useStrategyBuilder/index.ts @@ -8,6 +8,7 @@ export { type StrategyInstance, STRATEGY_DEFAULTS, STRATEGY_LABELS, + type DeduplicationParams, type SimilarityParams, type StrategySummaryTag, type MetadataWeightingParams diff --git a/lightly_studio_view/src/lib/hooks/useStrategyBuilder/strategyValidation.test.ts b/lightly_studio_view/src/lib/hooks/useStrategyBuilder/strategyValidation.test.ts index faf733ac32..1bfe21a295 100644 --- a/lightly_studio_view/src/lib/hooks/useStrategyBuilder/strategyValidation.test.ts +++ b/lightly_studio_view/src/lib/hooks/useStrategyBuilder/strategyValidation.test.ts @@ -8,6 +8,13 @@ const defaultDiversity: StrategyInstance = { isExpanded: true }; +const defaultDeduplication: StrategyInstance = { + id: '1', + type: 'deduplication', + params: { strength: 1, stopping_condition_minimum_distance: 0.1 }, + isExpanded: true +}; + const defaultTypicality: StrategyInstance = { id: '1', type: 'typicality', @@ -95,6 +102,42 @@ describe('isStrategyInstanceValid', () => { }); }); + describe('deduplication', () => { + it('returns true when minimum distance is positive', () => { + expect(isStrategyInstanceValid(defaultDeduplication)).toBe(true); + }); + + it('returns false when minimum distance is 0', () => { + const instance: StrategyInstance = { + ...defaultDeduplication, + params: { ...defaultDeduplication.params, stopping_condition_minimum_distance: 0 } + }; + + expect(isStrategyInstanceValid(instance)).toBe(false); + }); + + it('returns false when minimum distance is negative', () => { + const instance: StrategyInstance = { + ...defaultDeduplication, + params: { ...defaultDeduplication.params, stopping_condition_minimum_distance: -0.1 } + }; + + expect(isStrategyInstanceValid(instance)).toBe(false); + }); + + it('returns false when minimum distance is NaN', () => { + const instance: StrategyInstance = { + ...defaultDeduplication, + params: { + ...defaultDeduplication.params, + stopping_condition_minimum_distance: Number.NaN + } + }; + + expect(isStrategyInstanceValid(instance)).toBe(false); + }); + }); + describe('typicality', () => { it('returns true when strength is valid', () => { expect(isStrategyInstanceValid(defaultTypicality)).toBe(true); diff --git a/lightly_studio_view/src/lib/hooks/useStrategyBuilder/strategyValidation.ts b/lightly_studio_view/src/lib/hooks/useStrategyBuilder/strategyValidation.ts index 8644ab3624..a9f0af04b0 100644 --- a/lightly_studio_view/src/lib/hooks/useStrategyBuilder/strategyValidation.ts +++ b/lightly_studio_view/src/lib/hooks/useStrategyBuilder/strategyValidation.ts @@ -13,6 +13,10 @@ export function isStrategyInstanceValid(instance: StrategyInstance): boolean { return false; } + if (instance.type === 'deduplication') { + return isPositiveNumber(instance.params.stopping_condition_minimum_distance); + } + if (instance.type === 'similarity') { return instance.params.query_tag_id.trim().length > 0; } diff --git a/lightly_studio_view/src/lib/hooks/useStrategyBuilder/types.ts b/lightly_studio_view/src/lib/hooks/useStrategyBuilder/types.ts index 6016cb879e..089dca6ec8 100644 --- a/lightly_studio_view/src/lib/hooks/useStrategyBuilder/types.ts +++ b/lightly_studio_view/src/lib/hooks/useStrategyBuilder/types.ts @@ -3,6 +3,11 @@ export interface DiversityParams { strength: number; } +export interface DeduplicationParams { + strength: number; + stopping_condition_minimum_distance: number; +} + export interface TypicalityParams { strength: number; } @@ -33,6 +38,7 @@ export interface ClassBalancingParams { export interface StrategyParamsByType { diversity: DiversityParams; + deduplication: DeduplicationParams; typicality: TypicalityParams; similarity: SimilarityParams; metadata_weighting: MetadataWeightingParams; @@ -61,6 +67,12 @@ export const STRATEGY_OPTIONS: { type: StrategyType; label: string; description: description: 'Selects samples spread across the embedding space. Use to reduce redundancy and build varied training sets.' }, + { + type: 'deduplication', + label: 'Deduplication', + description: + 'Removes near-duplicates by keeping only samples that are at least a minimum distance apart in embedding space. May select fewer than the requested number of samples.' + }, { type: 'typicality', label: 'Typicality', @@ -93,6 +105,7 @@ export const STRATEGY_LABELS: Record = Object.fromEntries( export const STRATEGY_DEFAULTS: { [K in StrategyType]: StrategyParamsByType[K] } = { diversity: { strength: 1 }, + deduplication: { strength: 1, stopping_condition_minimum_distance: 0.1 }, typicality: { strength: 1 }, similarity: { query_tag_id: '', strength: 1 }, metadata_weighting: { metadata_key: '', strength: 1 }, diff --git a/lightly_studio_view/src/lib/hooks/useSubmitCombinationSelection/strategyApiMapping.test.ts b/lightly_studio_view/src/lib/hooks/useSubmitCombinationSelection/strategyApiMapping.test.ts index 90243dda88..abe7e0e2f1 100644 --- a/lightly_studio_view/src/lib/hooks/useSubmitCombinationSelection/strategyApiMapping.test.ts +++ b/lightly_studio_view/src/lib/hooks/useSubmitCombinationSelection/strategyApiMapping.test.ts @@ -9,6 +9,13 @@ const defaultDiversity: StrategyInstance = { isExpanded: true }; +const defaultDeduplication: StrategyInstance = { + id: 'dedup-1', + type: 'deduplication', + params: { strength: 1, stopping_condition_minimum_distance: 0.1 }, + isExpanded: true +}; + const defaultTypicality: StrategyInstance = { id: 'typ-1', type: 'typicality', @@ -59,6 +66,10 @@ describe('getMetadataKey', () => { expect(getMetadataKey(defaultDiversity)).toBe(''); }); + it('returns empty string for deduplication instances', () => { + expect(getMetadataKey(defaultDeduplication)).toBe(''); + }); + it('returns empty string for class_balancing instances', () => { expect(getMetadataKey(defaultClassBalancing)).toBe(''); }); @@ -73,6 +84,20 @@ describe('toApiStrategy', () => { }); }); + it('maps deduplication to deduplication strategy with null embedding model, strength and min distance', () => { + expect( + toApiStrategy({ + ...defaultDeduplication, + params: { strength: 2, stopping_condition_minimum_distance: 0.25 } + }) + ).toEqual({ + strategy_name: 'deduplication', + embedding_model_name: null, + strength: 2, + stopping_condition_minimum_distance: 0.25 + }); + }); + it('maps typicality to weights strategy using typicality- as metadata key', () => { expect(toApiStrategy({ ...defaultTypicality, params: { strength: 1.5 } })).toEqual({ strategy_name: 'weights', diff --git a/lightly_studio_view/src/lib/hooks/useSubmitCombinationSelection/strategyApiMapping.ts b/lightly_studio_view/src/lib/hooks/useSubmitCombinationSelection/strategyApiMapping.ts index 923de1cc06..23b8c193d1 100644 --- a/lightly_studio_view/src/lib/hooks/useSubmitCombinationSelection/strategyApiMapping.ts +++ b/lightly_studio_view/src/lib/hooks/useSubmitCombinationSelection/strategyApiMapping.ts @@ -17,6 +17,16 @@ export function toApiStrategy(instance: StrategyInstance): SamplingRequest['stra }; } + if (instance.type === 'deduplication') { + return { + strategy_name: 'deduplication', + embedding_model_name: null, + strength: instance.params.strength, + stopping_condition_minimum_distance: + instance.params.stopping_condition_minimum_distance + }; + } + if (instance.type === 'typicality' || instance.type === 'similarity') { return { strategy_name: 'weights', diff --git a/video_frame_curation.duckdb b/video_frame_curation.duckdb new file mode 100644 index 0000000000..717099e142 Binary files /dev/null and b/video_frame_curation.duckdb differ diff --git a/video_frame_curation.duckdb.backup b/video_frame_curation.duckdb.backup new file mode 100644 index 0000000000..f125b243b1 Binary files /dev/null and b/video_frame_curation.duckdb.backup differ diff --git a/video_frame_curation.duckdb.wal b/video_frame_curation.duckdb.wal new file mode 100644 index 0000000000..cececc1ac3 Binary files /dev/null and b/video_frame_curation.duckdb.wal differ diff --git a/video_frame_curation/README.md b/video_frame_curation/README.md new file mode 100644 index 0000000000..6ca008c738 --- /dev/null +++ b/video_frame_curation/README.md @@ -0,0 +1,95 @@ +# Automated video → curated frames pipeline + +This example shows how to run videos through LightlyStudio fully automatically and write +**curated extracted frames** back to storage (e.g. S3). It targets the workflow where new +videos arrive continuously, old ones are deleted by a retention policy, and the deliverable +is a deduplicated, content-rich set of frames per video. + +It treats the frames as a normal **image dataset** (no video dataset, no GUI needed). + +## What it does + +For every new video in the input location, the pipeline: + +1. **Extracts frames** from every new video sub-sampled to a target frame rate, via the + `ImageDataset.add_frames_from_videos(...)` API. +2. **Drops low-content / corrupt frames** using a row-uniformity metric + (`row_uniformity.compute_uniform_row_ratio`, ported from LightlyOne's `uniformRowRatio`). + Higher score ⇒ more uniform / blank, so only frames with score ≤ `uniform_row_ratio_max` + are kept. +3. **Deduplicates** the remaining frames with embedding-based deduplication sampling + (diversity selection with a minimum-distance stopping condition). The selected frames are + marked with the `deduplicated` tag (visible in the GUI), which sampling creates. The tag + is re-created each run, so it reflects the most recent run's selection. +4. **Uploads** the selected frames to `output_uri//`. + +## Incremental processing (only new videos) + +The pipeline connects to a **persistent** LightlyStudio database. Because +`add_frames_from_videos` always tags a video's frames with the video file stem and skips +videos whose stem tag already exists, re-running the pipeline only processes videos added +since the last run. A video that produced zero selected frames still has its stem tag and is +therefore never reprocessed. + +> **Important:** the database must live on **durable storage**. Use a DuckDB file on a +> persistent volume, or point LightlyStudio at Postgres for a production deployment +> (set `LIGHTLY_STUDIO_DATABASE_URL=postgresql://user:pass@host:5432/db` and remove the +> `--db-file` argument). If the database is lost, all videos look new again. + +## Run it locally + +```bash +cd video_frame_curation + +python pipeline.py \ + --input-uri s3://my-input-bucket/videos/ \ + --output-uri s3://my-output-bucket/curated-frames/ \ + --db-file ./video_frame_curation.duckdb \ + --extract-dir ./video_frame_curation_frames \ + --fps 1 \ + --uniform-row-ratio-max 0.025 \ + --dedup-min-distance 0.2 +``` + +Inspect the results in the GUI (point it at the same database). Frames are kept in +`--extract-dir`, and the selected frames carry the `deduplicated` tag. + +Equivalently, configure it via environment variables: + +| Variable | Meaning | +| --- | --- | +| `VIDEO_CURATION_INPUT_URI` | Location of the input videos (e.g. `s3://bucket/videos/`). | +| `VIDEO_CURATION_OUTPUT_URI` | Location the curated frames are written to. | +| `VIDEO_CURATION_DB_FILE` | Path to the persistent DuckDB file. | +| `VIDEO_CURATION_EXTRACT_DIR` | Persistent local directory for extracted frames. | + +S3 access uses the standard AWS credential resolution (`AWS_ACCESS_KEY_ID`, +`AWS_SECRET_ACCESS_KEY`, `AWS_REGION`, or an instance/role profile). + +## Schedule it with Airflow + +Copy `pipeline.py`, `row_uniformity.py`, and `airflow_dag.py` into your Airflow `dags/` +folder, install Airflow in the worker environment, and set the `VIDEO_CURATION_*` +environment variables. The DAG (`video_frame_curation`) runs `@weekly` and simply executes +`pipeline.py` (via a `BashOperator`). Airflow is intentionally **not** a dependency of +LightlyStudio. + +## Curation knobs + +| Argument | Default | Effect | +| --- | --- | --- | +| `--fps` | `1` | Frames per second to extract from each video. | +| `--uniform-row-ratio-max` | `0.025` | Max row-uniformity score to keep a frame (higher ⇒ more uniform/corrupt). | +| `--dedup-min-distance` | `0.2` | Deduplication stopping distance; larger ⇒ fewer, more diverse frames. | + +## Known limitations / TODOs + +- **Local disk usage:** frames are materialized to a persistent local directory + (`--extract-dir`) and kept there, since the database references those files. This assumes + local disk is available and accumulates over time. A future improvement is a disk-free / + streaming path (decode → embed → discard). +- **Database growth:** the persistent database accumulates frame samples over time. A future + improvement is to prune the non-selected frame samples after upload while keeping the + per-video stem tag as the durable processed-marker. +- **Deduplication scope:** all new frames in a run are deduplicated together (across the + videos added since the last run), but not against frames from previous runs. diff --git a/video_frame_curation/airflow_dag.py b/video_frame_curation/airflow_dag.py new file mode 100644 index 0000000000..0cfd1cc412 --- /dev/null +++ b/video_frame_curation/airflow_dag.py @@ -0,0 +1,36 @@ +"""Example Airflow DAG that runs the video-frame curation pipeline on a schedule. + +Drop this file (together with ``pipeline.py`` and ``row_uniformity.py``) into your Airflow +``dags/`` folder. The DAG runs weekly and processes only the videos that are new since the +previous run. Airflow is not a dependency of LightlyStudio; install it separately in the +environment that runs the scheduler/workers. + +``pipeline.py`` is a plain script, so the DAG just runs it. Configure it via environment +variables (see ``README.md``): +- ``VIDEO_CURATION_INPUT_URI`` - location of the input videos, e.g. ``s3://bucket/videos/`` +- ``VIDEO_CURATION_OUTPUT_URI`` - location for the curated frames +- ``VIDEO_CURATION_DB_FILE`` - path to the persistent DuckDB file (durable storage!) +- ``VIDEO_CURATION_EXTRACT_DIR`` - persistent dir for extracted frames (durable storage!) +""" + +from __future__ import annotations + +import os +from datetime import datetime + +from airflow import DAG +from airflow.operators.bash import BashOperator + +PIPELINE_SCRIPT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pipeline.py") + +with DAG( + dag_id="video_frame_curation", + description="Extract, filter and deduplicate frames from new videos into curated frames.", + start_date=datetime(2026, 1, 1), + schedule="@weekly", + catchup=False, +) as dag: + BashOperator( + task_id="curate_video_frames", + bash_command=f"python {PIPELINE_SCRIPT}", + ) diff --git a/video_frame_curation/pipeline.py b/video_frame_curation/pipeline.py new file mode 100644 index 0000000000..38dcfa345c --- /dev/null +++ b/video_frame_curation/pipeline.py @@ -0,0 +1,225 @@ +"""Automated video -> curated frames pipeline. + +For the new videos in an input location this script: +1. extracts frames (sub-sampled to a target fps), +2. drops low-content / corrupt frames using a row-uniformity metric, +3. deduplicates the remaining frames by embedding distance, and +4. uploads the selected frames to an output location. + +It is designed to run repeatedly (e.g. weekly via cron / Airflow). It connects to a +**persistent** LightlyStudio database and only processes new videos: because +``ImageDataset.add_frames_from_videos`` tags every video's frames with the video stem and +skips videos whose stem tag already exists, re-running only processes videos added since the +last run. A video that produced zero selected frames still has its stem tag, so it is never +reprocessed. + +Run it directly: + + python pipeline.py \\ + --input-uri s3://my-input-bucket/videos/ \\ + --output-uri s3://my-output-bucket/curated-frames/ \\ + --db-file ./video_frame_curation.duckdb + +or configure it via the ``VIDEO_CURATION_*`` environment variables (see ``README.md``). +""" + +from __future__ import annotations + +import argparse +import logging +import os +from pathlib import Path + +import fsspec +from PIL import Image + +import lightly_studio as ls +from lightly_studio.core.dataset_query.image_sample_field import ImageSampleField +from lightly_studio.core.image.image_sample import ImageSample +from lightly_studio.database import db_manager +from lightly_studio.resolvers import tag_resolver +from lightly_studio.sampling.sample import Sampling +from row_uniformity import compute_uniform_row_ratio + +logger = logging.getLogger(__name__) + +DATASET_NAME = "video_frame_curation" +UNIFORM_ROW_RATIO_KEY = "uniform_row_ratio" +# GUI-visible tag holding the frames selected by deduplication. Sampling creates it; it is +# re-created on each run (sampling requires an unused tag name), so it reflects the latest run. +SELECTED_TAG = "deduplicated" + + +def parse_args() -> argparse.Namespace: + """Parse the pipeline configuration, defaulting to the ``VIDEO_CURATION_*`` env vars.""" + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--input-uri", + default=os.environ.get("VIDEO_CURATION_INPUT_URI"), + help="Location holding the input videos (e.g. s3://bucket/videos/).", + ) + parser.add_argument( + "--output-uri", + default=os.environ.get("VIDEO_CURATION_OUTPUT_URI"), + help="Location the selected frames are written to.", + ) + parser.add_argument( + "--db-file", + default=os.environ.get("VIDEO_CURATION_DB_FILE", "./video_frame_curation.duckdb"), + help="Path to the persistent DuckDB file (must be on durable storage).", + ) + parser.add_argument( + "--extract-dir", + default=os.environ.get("VIDEO_CURATION_EXTRACT_DIR", "./video_frame_curation_frames"), + help=( + "Persistent local directory for extracted frames. The database references these " + "files, so they must stay on disk for the GUI to display them." + ), + ) + parser.add_argument( + "--fps", type=float, default=1.0, help="Frames per second to extract from each video." + ) + parser.add_argument( + "--uniform-row-ratio-max", + type=float, + default=0.025, + help="Keep only frames whose row-uniformity score is <= this (higher = more uniform).", + ) + parser.add_argument( + "--dedup-min-distance", + type=float, + default=0.2, + help="Deduplication stopping distance; larger => fewer, more diverse frames.", + ) + args = parser.parse_args() + if not args.input_uri or not args.output_uri: + parser.error( + "input and output locations are required " + "(--input-uri/--output-uri or VIDEO_CURATION_INPUT_URI/VIDEO_CURATION_OUTPUT_URI)." + ) + return args + + +def run_pipeline( + input_uri: str, + output_uri: str, + db_file: str, + extract_dir: str, + fps: float, + uniform_row_ratio_max: float, + dedup_min_distance: float, +) -> None: + """Process all not-yet-processed videos from ``input_uri`` into ``output_uri``.""" + # Reset any engine from a previous invocation so repeated runs in a long-lived worker + # (e.g. an Airflow worker) reconnect cleanly to the persistent database. + db_manager.close() + db_manager.connect(db_file=db_file) + dataset = ls.ImageDataset.load_or_create(name=DATASET_NAME) + + # Extract frames from all new videos in one pass. Already-processed videos (whose stem + # tag exists) are skipped, and every frame is tagged with its video stem. + new_ids = dataset.add_frames_from_videos( + videos_path=input_uri, extract_dir=Path(extract_dir), fps=fps, embed=True + ) + if not new_ids: + logger.info("No new videos to process in %s.", input_uri) + return + logger.info("Extracted %d frames from new videos.", len(new_ids)) + + # Drop low-content / corrupt frames using the row-uniformity metric. + kept_ids = _add_row_uniformity_and_filter( + dataset=dataset, sample_ids=new_ids, uniform_row_ratio_max=uniform_row_ratio_max + ) + logger.info("Kept %d/%d frames after row-uniformity filter.", len(kept_ids), len(new_ids)) + if not kept_ids: + return + + # Deduplicate the kept frames together and upload the selection. + selected = _deduplicate( + dataset=dataset, kept_ids=kept_ids, dedup_min_distance=dedup_min_distance + ) + logger.info("Selected %d frames after deduplication.", len(selected)) + _upload_frames(samples=selected, output_uri=output_uri) + + +def _add_row_uniformity_and_filter( + dataset: ls.ImageDataset, + sample_ids: list, + uniform_row_ratio_max: float, +) -> list: + """Store the row-uniformity metric per frame and return the ids of content-rich frames.""" + sample_metadata = [] + kept_ids = [] + for sample_id in sample_ids: + sample = dataset.get_sample(sample_id) + with Image.open(sample.file_path_abs) as image: + ratio = compute_uniform_row_ratio(image_greyscale=image.convert("L")) + sample_metadata.append((sample_id, {UNIFORM_ROW_RATIO_KEY: ratio})) + if ratio <= uniform_row_ratio_max: + kept_ids.append(sample_id) + + dataset.update_metadata(sample_metadata) + return kept_ids + + +def _deduplicate( + dataset: ls.ImageDataset, kept_ids: list, dedup_min_distance: float +) -> list[ImageSample]: + """Deduplicate this run's kept frames; the selection is tagged ``deduplicated``. + + Sampling runs directly on ``kept_ids`` and creates the result tag itself. The tag is + cleared first because sampling requires an unused tag name, so it reflects the latest run. + """ + _delete_tag(dataset=dataset, tag_name=SELECTED_TAG) + Sampling( + dataset_id=dataset.collection_id, + session=dataset.session, + input_sample_ids=kept_ids, + ).deduplicate( + n_samples_to_select=len(kept_ids), + sampling_result_tag_name=SELECTED_TAG, + stopping_condition_minimum_distance=dedup_min_distance, + ) + return dataset.query().match(ImageSampleField.tags.contains(SELECTED_TAG)).to_list() + + +def _upload_frames(samples: list[ImageSample], output_uri: str) -> None: + """Upload the given frames to ``output_uri//`` via fsspec. + + The video stem is recovered from the frame's parent directory (frames are extracted into + one subdirectory per video). + """ + out_fs, out_base = fsspec.core.url_to_fs(output_uri) + base = out_base.rstrip("/") + created_dirs: set[str] = set() + for sample in samples: + video_stem = Path(sample.file_path_abs).parent.name + dest_dir = f"{base}/{video_stem}" + if dest_dir not in created_dirs: + out_fs.makedirs(dest_dir, exist_ok=True) + created_dirs.add(dest_dir) + out_fs.put_file(sample.file_path_abs, f"{dest_dir}/{sample.file_name}") + + +def _delete_tag(dataset: ls.ImageDataset, tag_name: str) -> None: + tag = tag_resolver.get_by_name( + session=dataset.session, tag_name=tag_name, collection_id=dataset.collection_id + ) + if tag is not None: + tag_resolver.delete(session=dataset.session, tag_id=tag.tag_id) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + cli_args = parse_args() + run_pipeline( + input_uri=cli_args.input_uri, + output_uri=cli_args.output_uri, + db_file=cli_args.db_file, + extract_dir=cli_args.extract_dir, + fps=cli_args.fps, + uniform_row_ratio_max=cli_args.uniform_row_ratio_max, + dedup_min_distance=cli_args.dedup_min_distance, + ) diff --git a/video_frame_curation/row_uniformity.py b/video_frame_curation/row_uniformity.py new file mode 100644 index 0000000000..e8e91d7625 --- /dev/null +++ b/video_frame_curation/row_uniformity.py @@ -0,0 +1,93 @@ +"""Row-uniformity image metric. + +Computes the fraction of "uniform" rows in an image. A higher score means the image is +more likely to be blank / corrupt / a transition frame (e.g. a fade to black), so a low +score indicates a content-rich frame. + +This is a port of the ``uniformRowRatio`` metric used in LightlyOne +(``lightly_worker/inspect_rgb/rgb_properties/row_uniformity.py``), kept close to the +original so behavior matches what existing users are used to. +""" + +from __future__ import annotations + +import torch +from PIL.Image import Image as PILImage +from torch import Tensor +from torch.backends import mkldnn +from torch.nn import functional as nn_functional +from torchvision.transforms import functional as tr_functional + +RESIZE_MAX_HEIGHT = 400 +UNIFORM_PIXEL_LAPLACE_THRESHOLD = 5 / 255 +UNIFORM_ROW_PIXEL_RATIO = 0.97 + +# Laplacian filter kernel used to detect uniform (near-constant) pixel neighborhoods. +LAPLACIAN_KERNEL = torch.tensor( + [[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype=torch.float32 +).view(1, 1, 3, 3) + + +def compute_uniform_row_ratio(image_greyscale: PILImage) -> float: + """Compute the row-uniformity score of a greyscale image. + + The result is a float in ``[0, 1]``. Higher means the image is more likely to be + blank / corrupt: it is the proportion of rows in which almost all pixels are uniform. + + Returns ``0.0`` if the image has a dimension smaller than 3 pixels (the Laplacian + filter cannot be applied in that case). + + Algorithm: + 1. Resize to a max height of ``RESIZE_MAX_HEIGHT`` pixels. + 2. Apply the Laplacian filter. + 3. Detect uniform pixels (filter response below a threshold). + 4. Compute the ratio of uniform pixels per row. + 5. Return the proportion of rows whose ratio exceeds ``UNIFORM_ROW_PIXEL_RATIO``. + + Args: + image_greyscale: A single-channel (greyscale) PIL image. + """ + orig_width, orig_height = image_greyscale.size + if orig_width < 3 or orig_height < 3: + return 0.0 + + tensor_image = tr_functional.to_tensor(pic=image_greyscale) + + # Resize to a max height of RESIZE_MAX_HEIGHT pixels. + shrink_factor = RESIZE_MAX_HEIGHT / orig_height + if shrink_factor < 1: + target_height = int(orig_height * shrink_factor) + # The target width must stay at least 3, otherwise the Laplacian filter fails. + target_width = max(3, int(orig_width * shrink_factor)) + tensor_image = tr_functional.resize( + img=tensor_image, + size=(target_height, target_width), + antialias=False, + ) + + uniform_pixels = _detect_uniform_pixels(tensor_image=tensor_image) + + # Detect uniform rows from the ratio of uniform pixels per row. + row_ratios = uniform_pixels.mean(dim=2) + uniform_rows = (row_ratios > UNIFORM_ROW_PIXEL_RATIO).float() + + return torch.mean(uniform_rows).item() + + +def _detect_uniform_pixels(tensor_image: Tensor) -> Tensor: + """Detect uniform pixels via a Laplacian filter. + + Returns a float tensor that is 1 where the filter response is near zero (uniform pixel) + and 0 otherwise. + """ + # Disable the MKL-DNN backend: it switches conv2d algorithms based on image size, which + # can cause large memory spikes when processing many differently sized frames. + with mkldnn.flags(enabled=False): + tensor_image = nn_functional.conv2d( + input=tensor_image.unsqueeze(0), weight=LAPLACIAN_KERNEL + ).squeeze(0) + + return ( + (tensor_image > -UNIFORM_PIXEL_LAPLACE_THRESHOLD) + & (tensor_image < UNIFORM_PIXEL_LAPLACE_THRESHOLD) + ).float()