From 615bf9d02397cd473916aedccc2e012079c486ab Mon Sep 17 00:00:00 2001 From: libeanim Date: Sat, 24 Feb 2024 18:38:01 +0100 Subject: [PATCH 1/2] Add keyframe subsampler to FrameSubsampler - Fix encode_formats bug #311 - Fix bug `dictionary keys changed during iteration` in FrameSubsampler --- video2dataset/subsamplers/frame_subsampler.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/video2dataset/subsamplers/frame_subsampler.py b/video2dataset/subsamplers/frame_subsampler.py index e3c1b572..66daefa3 100644 --- a/video2dataset/subsamplers/frame_subsampler.py +++ b/video2dataset/subsamplers/frame_subsampler.py @@ -36,11 +36,10 @@ def __init__(self, frame_rate, downsample_method="fps", encode_format="mp4"): self.frame_rate = frame_rate self.downsample_method = downsample_method self.output_modality = "video" if downsample_method == "fps" else "jpg" - self.encode_format = encode_format - + self.encode_formats = { "video": encode_format } + def __call__(self, streams, metadata=None): - # TODO: you might not want to pop it (f.e. in case of other subsamplers) - video_bytes = streams.pop("video") + video_bytes = streams["video"] subsampled_bytes, subsampled_metas = [], [] for i, vid_bytes in enumerate(video_bytes): with tempfile.TemporaryDirectory() as tmpdir: @@ -52,6 +51,10 @@ def __call__(self, streams, metadata=None): _ = ffmpeg.input(f"{tmpdir}/input.mp4") _ = _.filter("fps", fps=self.frame_rate) _ = _.output(f"{tmpdir}/output.mp4", reset_timestamps=1).run(capture_stdout=True, quiet=True) + elif self.downsample_method == "keyframe": + _ = ffmpeg.input(f"{tmpdir}/input.mp4", discard='nokey') + # _ = _.filter("select", "eq(pict_type,I)") + _ = _.output(f"{tmpdir}/output.mp4", **{"c:s": 'copy', 'c': 'copy', 'copyts': None}).run(capture_stdout=True, quiet=True) elif "frame" in self.downsample_method: _ = ffmpeg.input(f"{tmpdir}/input.mp4") _ = _.filter("select", "eq(n,0)") From 12d7ae171bf261dbb65a358f49ea97927615a33e Mon Sep 17 00:00:00 2001 From: libeanim Date: Sat, 24 Feb 2024 18:52:32 +0100 Subject: [PATCH 2/2] Format code in frame_subsampler --- video2dataset/subsamplers/frame_subsampler.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/video2dataset/subsamplers/frame_subsampler.py b/video2dataset/subsamplers/frame_subsampler.py index 66daefa3..c1d9ae6d 100644 --- a/video2dataset/subsamplers/frame_subsampler.py +++ b/video2dataset/subsamplers/frame_subsampler.py @@ -1,6 +1,7 @@ """ frame subsampler adjusts the fps of the videos to some constant value """ + import tempfile import os import copy @@ -36,8 +37,8 @@ def __init__(self, frame_rate, downsample_method="fps", encode_format="mp4"): self.frame_rate = frame_rate self.downsample_method = downsample_method self.output_modality = "video" if downsample_method == "fps" else "jpg" - self.encode_formats = { "video": encode_format } - + self.encode_formats = {"video": encode_format} + def __call__(self, streams, metadata=None): video_bytes = streams["video"] subsampled_bytes, subsampled_metas = [], [] @@ -52,9 +53,10 @@ def __call__(self, streams, metadata=None): _ = _.filter("fps", fps=self.frame_rate) _ = _.output(f"{tmpdir}/output.mp4", reset_timestamps=1).run(capture_stdout=True, quiet=True) elif self.downsample_method == "keyframe": - _ = ffmpeg.input(f"{tmpdir}/input.mp4", discard='nokey') + _ = ffmpeg.input(f"{tmpdir}/input.mp4", discard="nokey") # _ = _.filter("select", "eq(pict_type,I)") - _ = _.output(f"{tmpdir}/output.mp4", **{"c:s": 'copy', 'c': 'copy', 'copyts': None}).run(capture_stdout=True, quiet=True) + _ = _.output(f"{tmpdir}/output.mp4", **{"c:s": "copy", "c": "copy", "copyts": None}) + _ = _.run(capture_stdout=True, quiet=True) elif "frame" in self.downsample_method: _ = ffmpeg.input(f"{tmpdir}/input.mp4") _ = _.filter("select", "eq(n,0)")