Adding functionality to generate_movie().

act/tests/utils/test_io_utils.py

@@ -1,14 +1,17 @@
 import tempfile
 from pathlib import Path
-from os import PathLike, getcwd
+from os import PathLike, getcwd, chdir
 from string import ascii_letters
 import random
 import glob
 import pytest
 import os
+import numpy as np
+import shutil
 
 import act
 from act.tests import sample_files
+from arm_test_data import locate as test_data_locate
 
 try:
     import moviepy.video.io.ImageSequenceClip
@@ -222,5 +225,46 @@ def test_generate_movie():
         'https://github.com/ARM-DOE/ACT/blob/main/act/tests/plotting/baseline/test_contourf.png?raw=true',
         'https://github.com/ARM-DOE/ACT/blob/main/act/tests/plotting/baseline/test_contourf2.png?raw=true'
     ]
-    result = act.utils.generate_movie(files, fps=5)
-    assert str(result) == 'movie.mp4'
+    cwd = Path.cwd()
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        try:
+            chdir(tmpdirname)
+
+            # Test URL path for making movie
+            result = act.utils.generate_movie(files, fps=5)
+            assert Path(result).name == 'movie.mp4'
+
+            # Test list of files for making movie
+            files = ['test_contour.png', 'test_contour2.png', 'test_contourf.png', 'test_contourf2.png']
+            basepath = Path(Path(__file__).parents[1], 'plotting', 'baseline')
+            files = [Path(basepath, fl) for fl in files]
+            write_filename = Path(tmpdirname, 'one', 'two', 'three', 'movie_filename_testing.mp4')
+            result = act.utils.generate_movie(files, write_filename=write_filename)
+            assert result == str(write_filename)
+            assert np.isclose(Path(write_filename).stat().st_size, 173189, 1000)
+
+            # Test PosixPath generator for making movie
+            file_generator = basepath.glob('test_contour[!_]*.png')
+            result = act.utils.generate_movie(file_generator, write_filename=write_filename)
+            assert result == str(write_filename)
+            assert np.isclose(Path(write_filename).stat().st_size, 173189, 1000)
+
+            # Test passing path to directory of images
+            dir_path = Path(tmpdirname, 'images')
+            dir_path.mkdir()
+            for fl in files:
+                shutil.copy(fl, Path(dir_path, Path(fl).name))
+
+            files = dir_path.glob('*.*')
+            result = act.utils.generate_movie(dir_path)
+            assert Path(result).name == 'movie.mp4'
+            assert np.isclose(Path(result).stat().st_size, 173189, 1000)
+
+            # Test converting movie format
+            write_filename = 'movie2.mp4'
+            result = act.utils.generate_movie(result, write_filename=write_filename)
+            assert Path(result).name == write_filename
+            assert np.isclose(Path(result).stat().st_size, 173189, 1000)
+
+        finally:
+            chdir(cwd)

act/utils/io_utils.py

@@ -7,9 +7,12 @@
 import gzip
 import shutil
 import tempfile
+import numpy as np
+import types
 
 try:
     import moviepy.video.io.ImageSequenceClip
+    from moviepy.video.io.VideoFileClip import VideoFileClip
     MOVIEPY_AVAILABLE = True
 except ImportError:
     MOVIEPY_AVAILABLE = False
@@ -284,53 +287,66 @@ def unpack_gzip(filename, write_directory=None, remove=False):
     return str(write_filename)
 
 
-def generate_movie(images, write_directory=None, write_filename=None, fps=10, codec=None, threads=None):
+def generate_movie(images, write_filename=None, fps=10, **kwargs):
     """
-    Creates a movie from a list of images
+    Creates a movie from a list of images or convert movie to different type
 
     ...
 
     Parameters
     ----------
-    images : list
-        List of images in the correct order to make into a movie
-    write_directory : str, pahtlib.Path, list, None
-        Path to directory to place newly created gunzip file.
+    images : list, PosixPath generator, path to a directory, single string/PosixPath to movie
+        List of images in the correct order to make into a movie or a generator from
+        a pathlib.Path.glob() search. If a path to directory will create movie from all files
+        in that directory in alpanumeric order. If a single file to a movie will allow for converting
+        to new format defined by file extension of write_filename.
     write_filename : str, pathlib.Path, None
-        Movie output filename
+        Movie output filename. Default is 'movie.mp4' in current directory. If a path to a directory
+        that does not exist, will create the directory path.
     fps: int
-        Frames per second
-    codec: int
-        Codec to use for image encoding
-    threads: int
-        Number of threads to use for ffmpeg
+        Frames per second. Passed into moviepy->ImageSequenceClip() method
+    **kwargs: dict
+        Optional keywords passed into moviepy->write_videofile() method
 
 
     Returns
     -------
     write_filename : str
-        Full path name of created gunzip file
+        Full path name of created movie file
 
     """
     if not MOVIEPY_AVAILABLE:
         raise ImportError(
             'MoviePy needs to be installed on your system to make movies.'
         )
 
+    # Set default movie name
     if write_filename is None:
-        write_filename = 'movie.mp4'
-
-    if write_directory is not None:
-        write_directory = Path(write_directory)
-        write_directory.mkdir(parents=True, exist_ok=True)
-        write_filename = Path(write_filename).name
-    elif Path(write_filename).parent != Path('.'):
-        write_directory = Path(write_filename).parent
+        write_filename = Path(Path().cwd(), 'movie.mp4')
+
+    # Check if images is pointing to a directory. If so ensure is a string not PosixPath
+    IS_MOVIE = False
+    if isinstance(images, (types.GeneratorType, list, tuple)):
+        images = [str(image) for image in images]
+        images.sort()
+    elif isinstance(images, (PathLike, str)) and Path(images).is_file():
+        IS_MOVIE = True
+        images = str(images)
+    elif isinstance(images, (PathLike, str)) and Path(images).is_dir():
+        images = str(images)
+
+    # Ensure full path to filename exists
+    write_directory = Path(write_filename).parent
+    write_directory.mkdir(parents=True, exist_ok=True)
+
+    if IS_MOVIE:
+        with VideoFileClip(images) as clip:
+            # Not sure why but need to set the duration of the clip with subclip() to write
+            # the full file out.
+            clip = clip.subclip(t_start=clip.start, t_end=clip.end * clip.fps)
+            clip.write_videofile(str(write_filename), fps=fps, **kwargs)
     else:
-        write_directory = Path('.')
-
-    full_path = write_directory / write_filename
-    clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(images, fps=fps)
-    clip.write_videofile(str(full_path), codec=codec, threads=threads)
+        clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(images, fps=fps)
+        clip.write_videofile(str(write_filename), **kwargs)
 
-    return full_path
+    return str(write_filename)