Add keyframe subsampler to FrameSubsampler

- Fix encode_formats bug iejMac#311
- Fix bug `dictionary keys changed during iteration` in FrameSubsampler

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)")