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"