Factor out shared ffmpeg functions

Author: Daniel Flores

test/generate_reference_resources.py

@@ -28,6 +28,40 @@ def convert_image_to_tensor(image_path):
     os.remove(image_path)
 
 
+def get_frame_by_index(video_path, frame, output_path, stream):
+    cmd = [
+        "ffmpeg",
+        "-y",
+        "-i",
+        video_path,
+        "-map",
+        f"0:{stream}",
+        "-vf",
+        f"select=eq(n\\,{frame})",
+        "-vsync",
+        "vfr",
+        "-q:v",
+        "2",
+        output_path,
+    ]
+    subprocess.run(cmd, check=True)
+
+
+def get_frame_by_timestamp(video_path, timestamp, output_path):
+    cmd = [
+        "ffmpeg",
+        "-y",
+        "-ss",
+        str(timestamp),
+        "-i",
+        video_path,
+        "-frames:v",
+        "1",
+        output_path,
+    ]
+    subprocess.run(cmd, check=True)
+
+
 def main():
     SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
     TORCHCODEC_PATH = os.path.dirname(SCRIPT_DIR)
@@ -47,42 +81,15 @@ def main():
             # so that the name of reference frames matches the index when accessing that
             # frame in the Python decoder.
             output_bmp = f"{VIDEO_PATH}.stream{stream}.frame{frame:06d}.bmp"
-            frame_name = f"{frame:06d}"
-            cmd = [
-                "ffmpeg",
-                "-y",
-                "-i",
-                VIDEO_PATH,
-                "-map",
-                f"0:{stream}",
-                "-vf",
-                f"select=eq(n\\,{frame})",
-                "-vsync",
-                "vfr",
-                "-q:v",
-                "2",
-                output_bmp,
-            ]
-            subprocess.run(cmd, check=True)
+            get_frame_by_index(VIDEO_PATH, frame, output_bmp, stream=stream)
             convert_image_to_tensor(output_bmp)
 
     # Extract individual frames at specific timestamps, including the last frame of the video.
     seek_timestamp = [6.0, 6.1, 10.0, 12.979633]
     timestamp_name = [f"{seek_timestamp:06f}" for seek_timestamp in seek_timestamp]
     for timestamp, name in zip(seek_timestamp, timestamp_name):
         output_bmp = f"{VIDEO_PATH}.time{name}.bmp"
-        cmd = [
-            "ffmpeg",
-            "-y",
-            "-ss",
-            str(timestamp),
-            "-i",
-            VIDEO_PATH,
-            "-frames:v",
-            "1",
-            f"{VIDEO_PATH}.time{name}.bmp",
-        ]
-        subprocess.run(cmd, check=True)
+        get_frame_by_timestamp(VIDEO_PATH, timestamp, output_bmp)
         convert_image_to_tensor(output_bmp)
 
     # This video was generated by running the following:
@@ -93,22 +100,8 @@ def main():
     VIDEO_PATH = os.path.join(RESOURCES_DIR, "h265_video.mp4")
     FRAMES = [5]
     for frame in FRAMES:
-        frame_name = f"{frame:06d}"
-        output_bmp = f"{VIDEO_PATH}.stream0.frame{frame_name}.bmp"
-        cmd = [
-            "ffmpeg",
-            "-y",
-            "-i",
-            VIDEO_PATH,
-            "-vf",
-            f"select=eq(n\\,{frame})",
-            "-vsync",
-            "vfr",
-            "-q:v",
-            "2",
-            output_bmp,
-        ]
-        subprocess.run(cmd, check=True)
+        output_bmp = f"{VIDEO_PATH}.stream0.frame{frame:06d}.bmp"
+        get_frame_by_index(VIDEO_PATH, frame, output_bmp, stream=0)
         convert_image_to_tensor(output_bmp)
 
     # This video was generated by running the following:
@@ -118,22 +111,8 @@ def main():
     FRAMES = [10]
 
     for frame in FRAMES:
-        frame_name = f"{frame:06d}"
-        output_bmp = f"{VIDEO_PATH}.stream0.frame{frame_name}.bmp"
-        cmd = [
-            "ffmpeg",
-            "-y",
-            "-i",
-            VIDEO_PATH,
-            "-vf",
-            f"select=eq(n\\,{frame})",
-            "-vsync",
-            "vfr",
-            "-q:v",
-            "2",
-            output_bmp,
-        ]
-        subprocess.run(cmd, check=True)
+        output_bmp = f"{VIDEO_PATH}.stream0.frame{frame:06d}.bmp"
+        get_frame_by_index(VIDEO_PATH, frame, output_bmp, stream=0)
         convert_image_to_tensor(output_bmp)