Refactor video writer to use imageio instead of skvideo

environment.yml

@@ -12,6 +12,7 @@ dependencies:
     # Packages SLEAP uses directly
     - conda-forge::attrs >=21.2.0  #,<=21.4.0
     - conda-forge::cattrs ==1.1.1
+    - conda-forge::imageio-ffmpeg # Required for imageio to read/write videos with ffmpeg
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::networkx

environment_mac.yml

@@ -12,6 +12,7 @@ dependencies:
     - conda-forge::importlib-metadata <7.1.0
     - conda-forge::cattrs ==1.1.1
     - conda-forge::h5py
+    - conda-forge::imageio-ffmpeg # Required for imageio to read/write videos with ffmpeg
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::keras <2.10.0,>=2.9.0rc0  # Required by tensorflow-macos

environment_no_cuda.yml

@@ -13,6 +13,7 @@ dependencies:
     # Packages SLEAP uses directly
     - conda-forge::attrs >=21.2.0  #,<=21.4.0
     - conda-forge::cattrs ==1.1.1
+    - conda-forge::imageio-ffmpeg # Required for imageio to read/write videos with ffmpeg
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::networkx

sleap/gui/commands.py

@@ -1329,12 +1329,10 @@ def ask(context: CommandContext, params: dict) -> bool:
         # makes mp4's that most programs can't open (VLC can).
         default_out_filename = context.state["filename"] + ".avi"
 
-        # But if we can write mpegs using sci-kit video, use .mp4
-        # since it has trouble writing .avi files.
-        if VideoWriter.can_use_skvideo():
+        if VideoWriter.can_use_ffmpeg():
             default_out_filename = context.state["filename"] + ".mp4"
 
-        # Ask where use wants to save video file
+        # Ask where user wants to save video file
         filename, _ = FileDialog.save(
             context.app,
             caption="Save Video As...",

sleap/gui/dialogs/export_clip.py

@@ -11,16 +11,16 @@ def __init__(self):
 
         super().__init__(form_name="labeled_clip_form")
 
-        can_use_skvideo = VideoWriter.can_use_skvideo()
+        can_use_ffmpeg = VideoWriter.can_use_ffmpeg()
 
-        if can_use_skvideo:
+        if can_use_ffmpeg:
             message = (
                 "<i><b>MP4</b> file will be encoded using "
-                "system ffmpeg via scikit-video (preferred option).</i>"
+                "system ffmpeg via imageio (preferred option).</i>"
             )
         else:
             message = (
-                "<i>Unable to use ffpmeg via scikit-video. "
+                "<i>Unable to use ffpmeg via imageio. "
                 "<b>AVI</b> file will be encoding using OpenCV.</i>"
             )
 

sleap/io/videowriter.py

@@ -12,6 +12,8 @@
 from abc import ABC, abstractmethod
 import cv2
 import numpy as np
+import imageio.v2 as iio
+import logging
 
 
 class VideoWriter(ABC):
@@ -32,22 +34,27 @@ def close(self):
     @staticmethod
     def safe_builder(filename, height, width, fps):
         """Builds VideoWriter based on available dependencies."""
-        if VideoWriter.can_use_skvideo():
-            return VideoWriterSkvideo(filename, height, width, fps)
+        if VideoWriter.can_use_ffmpeg():
+            logging.info("USING IMAGEIO")
+            return VideoWriterImageio(filename, height, width, fps)
         else:
             return VideoWriterOpenCV(filename, height, width, fps)
 
     @staticmethod
-    def can_use_skvideo():
-        # See if we can import skvideo
+    def can_use_ffmpeg():
+        """Check if ffmpeg is available for writing videos."""
         try:
-            import skvideo
+            import imageio_ffmpeg as ffmpeg
         except ImportError:
             return False
 
-        # See if skvideo can find FFMPEG
-        if skvideo.getFFmpegVersion() != "0.0.0":
-            return True
+        try:
+            # Try to get the version of the ffmpeg plugin
+            ffmpeg_version = ffmpeg.get_ffmpeg_version()
+            if ffmpeg_version:
+                return True
+        except Exception:
+            return False
 
         return False
 
@@ -68,11 +75,11 @@ def close(self):
         self._writer.release()
 
 
-class VideoWriterSkvideo(VideoWriter):
-    """Writes video using scikit-video as wrapper for ffmpeg.
+class VideoWriterImageio(VideoWriter):
+    """Writes video using imageio as a wrapper for ffmpeg.
 
     Attributes:
-        filename: Path to mp4 file to save to.
+        filename: Path to video file to save to.
         height: Height of movie frames.
         width: Width of movie frames.
         fps: Playback framerate to save at.
@@ -85,28 +92,38 @@ class VideoWriterSkvideo(VideoWriter):
     def __init__(
         self, filename, height, width, fps, crf: int = 21, preset: str = "superfast"
     ):
-        import skvideo.io
-
-        fps = str(fps)
-        self._writer = skvideo.io.FFmpegWriter(
+        self.filename = filename
+        self.height = height
+        self.width = width
+        self.fps = fps
+        self.crf = crf
+        self.preset = preset
+
+        import imageio_ffmpeg as ffmpeg
+
+        # Imageio's ffmpeg writer parameters
+        # https://imageio.readthedocs.io/en/stable/examples.html#writing-videos-with-ffmpeg-and-vaapi
+        # Use `ffmpeg -h encoder=libx264`` to see all options for libx264 output_params
+        # output_params must be a list of strings
+        # iio.help(name='FFMPEG') to test
+        self.writer = iio.get_writer(
             filename,
-            inputdict={
-                "-r": fps,
-            },
-            outputdict={
-                "-c:v": "libx264",
-                "-preset": preset,
-                "-vf": "scale=trunc(iw/2)*2:trunc(ih/2)*2",  # Need even dims for libx264
-                "-framerate": fps,
-                "-crf": str(crf),
-                "-pix_fmt": "yuv420p",
-            },
+            fps=fps,
+            codec="libx264",
+            format="FFMPEG",
+            pixelformat="yuv420p",
+            output_params=[
+                "-preset",
+                preset,
+                "-crf",
+                str(crf),
+            ],
         )
 
     def add_frame(self, img, bgr: bool = False):
         if bgr:
             img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-        self._writer.writeFrame(img)
+        self.writer.append_data(img)
 
     def close(self):
-        self._writer.close()
+        self.writer.close()

tests/io/test_videowriter.py

@@ -1,5 +1,7 @@
 import os
-from sleap.io.videowriter import VideoWriter, VideoWriterOpenCV
+import cv2
+from pathlib import Path
+from sleap.io.videowriter import VideoWriter, VideoWriterOpenCV, VideoWriterImageio
 
 
 def test_video_writer(tmpdir, small_robot_mp4_vid):
@@ -38,3 +40,62 @@ def test_cv_video_writer(tmpdir, small_robot_mp4_vid):
     writer.close()
 
     assert os.path.exists(out_path)
+
+
+def test_imageio_video_writer_avi(tmpdir, small_robot_mp4_vid):
+    out_path = Path(tmpdir) / "clip.avi"
+
+    # Make sure imageio video writer works
+    writer = VideoWriterImageio(
+        out_path,
+        height=small_robot_mp4_vid.height,
+        width=small_robot_mp4_vid.width,
+        fps=small_robot_mp4_vid.fps,
+    )
+
+    writer.add_frame(small_robot_mp4_vid[0][0])
+    writer.add_frame(small_robot_mp4_vid[1][0])
+
+    writer.close()
+
+    assert os.path.exists(out_path)
+    # Check attributes
+    assert writer.height == small_robot_mp4_vid.height
+    assert writer.width == small_robot_mp4_vid.width
+    assert writer.fps == small_robot_mp4_vid.fps
+    assert writer.filename == out_path
+    assert writer.crf == 21
+    assert writer.preset == "superfast"
+
+
+def test_imageio_video_writer_odd_size(tmpdir, movenet_video):
+    out_path = Path(tmpdir) / "clip.mp4"
+
+    # Reduce the size of the video frames by 1 pixel in each dimension
+    reduced_height = movenet_video.height - 1
+    reduced_width = movenet_video.width - 1
+
+    # Initialize the writer with the reduced dimensions
+    writer = VideoWriterImageio(
+        out_path,
+        height=reduced_height,
+        width=reduced_width,
+        fps=movenet_video.fps,
+    )
+
+    # Resize frames and add them to the video
+    for i in range(len(movenet_video) - 1):
+        frame = movenet_video[i][0]  # Access the actual frame object
+        reduced_frame = cv2.resize(frame, (reduced_width, reduced_height))
+        writer.add_frame(reduced_frame)
+
+    writer.close()
+
+    # Assertions to validate the test
+    assert os.path.exists(out_path)
+    assert writer.height == reduced_height
+    assert writer.width == reduced_width
+    assert writer.fps == movenet_video.fps
+    assert writer.filename == out_path
+    assert writer.crf == 21
+    assert writer.preset == "superfast"