diff --git a/environment.yml b/environment.yml
index 9c5758c13..2aba3c7d2 100644
--- a/environment.yml
+++ b/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
diff --git a/environment_mac.yml b/environment_mac.yml
index 42d6e028c..9ab10a1b8 100644
--- a/environment_mac.yml
+++ b/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
diff --git a/environment_no_cuda.yml b/environment_no_cuda.yml
index fc13f839a..2adee7a89 100644
--- a/environment_no_cuda.yml
+++ b/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
diff --git a/sleap/gui/commands.py b/sleap/gui/commands.py
index e3ef8522d..692f19c78 100644
--- a/sleap/gui/commands.py
+++ b/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...",
diff --git a/sleap/gui/dialogs/export_clip.py b/sleap/gui/dialogs/export_clip.py
index 312f9a807..f84766d18 100644
--- a/sleap/gui/dialogs/export_clip.py
+++ b/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 = (
"MP4 file will be encoded using "
- "system ffmpeg via scikit-video (preferred option)."
+ "system ffmpeg via imageio (preferred option)."
)
else:
message = (
- "Unable to use ffpmeg via scikit-video. "
+ "Unable to use ffpmeg via imageio. "
"AVI file will be encoding using OpenCV."
)
diff --git a/sleap/io/videowriter.py b/sleap/io/videowriter.py
index 510fad739..cd710c9d5 100644
--- a/sleap/io/videowriter.py
+++ b/sleap/io/videowriter.py
@@ -12,6 +12,7 @@
from abc import ABC, abstractmethod
import cv2
import numpy as np
+import imageio.v2 as iio
class VideoWriter(ABC):
@@ -32,22 +33,26 @@ 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():
+ 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 +73,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 +90,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()
diff --git a/tests/io/test_videowriter.py b/tests/io/test_videowriter.py
index dea193117..35d9bc6df 100644
--- a/tests/io/test_videowriter.py
+++ b/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"