ENH: add Brain.save_movie()

Author: christianbrodbeck

surfer/utils.py

@@ -5,6 +5,7 @@
 from os import path as op
 import inspect
 from functools import wraps
+import subprocess
 
 import numpy as np
 import nibabel as nib
@@ -241,7 +242,6 @@ def set_log_level(verbose=None, return_old_level=False):
         if verbose not in logging_types:
             raise ValueError('verbose must be of a valid type')
         verbose = logging_types[verbose]
-    logger = logging.getLogger('surfer')
     old_verbose = logger.level
     logger.setLevel(verbose)
     return (old_verbose if return_old_level else None)
@@ -273,7 +273,6 @@ def set_log_file(fname=None, output_format='%(message)s', overwrite=None):
         but additionally raises a warning to notify the user that log
         entries will be appended.
     """
-    logger = logging.getLogger('surfer')
     handlers = logger.handlers
     for h in handlers:
         if isinstance(h, logging.FileHandler):
@@ -634,3 +633,96 @@ def has_fsaverage(subjects_dir=None):
 
 requires_fsaverage = np.testing.dec.skipif(not has_fsaverage(),
                                            'Requires fsaverage subject data')
+
+def has_ffmpeg():
+    """Test whether the FFmpeg is available in a subprocess
+
+    Returns
+    -------
+    ffmpeg_exists : bool
+        True if FFmpeg can be successfully called, False otherwise.
+    """
+    rcode = subprocess.call(["type", "ffmpeg"], stdout=subprocess.PIPE,
+                            stderr=subprocess.PIPE)
+    ffmpeg_exists = rcode == 0
+    return ffmpeg_exists
+
+
+def ffmpeg(dst, frame_path, framerate=10, codec='mpeg4', opt="", inopt="",
+           outopt=""):
+    """Run FFmpeg in a subprocess to convert an image sequence into a movie
+
+    Parameters
+    ----------
+    dst : str
+        Destination path. If the extension is not ".mov" or ".avi", ".mov" is
+        added. If the file already exists it is overwritten.
+    frame_path : str
+        Path to the source frames (with a frame number field like '%04d').
+    framerate : int
+        Framerate of the movie (frames per second).
+    codec : str
+        Codec to use (default 'mpeg4').
+    opt, inopt, outopt : str
+        FFmpeg options, infile options and outfile options (e.g., "-o1 value
+        -o2 value", see FFmpeg help for possible options).
+
+    Notes
+    -----
+    Requires FFmpeg to be in the path. FFmpeg can be downlaoded from `here
+    <http://ffmpeg.org/download.html>`_. Stdout and stderr are written to the
+    logger. If the movie file is not created, a RuntimeError is raised.
+    """
+    # make sure FFmpeg is available
+    if not has_ffmpeg():
+        err = ("FFmpeg is not in the path and is needed for saving "
+               "movies. Install FFmpeg and try again. It can be "
+               "downlaoded from http://ffmpeg.org/download.html.")
+        raise RuntimeError(err)
+
+    # find target path
+    dst = os.path.expanduser(dst)
+    dst = os.path.abspath(dst)
+    root, ext = os.path.splitext(dst)
+    dirname = os.path.dirname(dst)
+    if ext not in ['.mov', '.avi']:
+        dst += '.mov'
+
+    if os.path.exists(dst):
+        os.remove(dst)
+    elif not os.path.exists(dirname):
+        os.mkdir(dirname)
+
+    frame_dir, frame_fmt = os.path.split(frame_path)
+
+    # make the movie
+    cmd = ['ffmpeg']
+    if opt:
+        cmd.extend(opt.split())
+    if inopt:
+        cmd.extend(inopt.split())
+    cmd.extend(('-i', frame_fmt))
+    if outopt:
+        cmd.extend(outopt.split())
+    if '-r ' not in outopt:
+        cmd.extend(('-r', str(framerate)))
+    if not ('-c' in cmd or '-codec' in cmd):
+        cmd.extend(('-c', codec))
+    cmd.append(dst)
+
+    logger.info("Running FFmpeg with command: %s", ' '.join(cmd))
+    sp = subprocess.Popen(cmd, cwd=frame_dir, stdout=subprocess.PIPE,
+                          stderr=subprocess.PIPE)
+
+    # log stdout and stderr
+    stdout, stderr = sp.communicate()
+    std_info = os.linesep.join(("FFmpeg stdout", '=' * 25, stdout))
+    logger.info(std_info)
+    if stderr.strip():
+        err_info = os.linesep.join(("FFmpeg stderr", '=' * 27, stderr))
+        logger.error(err_info)
+
+    # check that movie file is created
+    if not os.path.exists(dst):
+        err = ("FFmpeg failed, no file created; see log for more more information")
+        raise RuntimeError(err)

surfer/viz.py

@@ -1,5 +1,6 @@
 import os
 from os.path import join as pjoin
+from tempfile import mkdtemp
 from warnings import warn
 
 import numpy as np
@@ -16,7 +17,7 @@
 from . import utils, io
 from .config import config
 from .utils import (Surface, verbose, create_color_lut, _get_subjects_dir,
-                    string_types)
+                    string_types, has_ffmpeg, ffmpeg)
 
 
 import logging
@@ -1623,13 +1624,20 @@ def set_data_smoothing_steps(self, smoothing_steps, verbose=None):
                 data["smoothing_steps"] = smoothing_steps
         self._toggle_render(True, views)
 
-    def set_time(self, time):
-        """Set the data time index to the time point closest to time
+    def index_for_time(self, time, rounding='closest'):
+        """Find the data time index closest to a specific time point
 
         Parameters
         ----------
         time : scalar
             Time.
+        rounding : 'closest' | 'up' | 'down
+            How to round if the exact time point is not an index.
+
+        Returns
+        -------
+        index : int
+            Data time index closest to time.
         """
         if self.n_times is None:
             raise RuntimeError("Brain has no time axis")
@@ -1644,7 +1652,27 @@ def set_time(self, time):
                    "[%s, %s]" % (time, tmin, tmax))
             raise ValueError(err)
 
-        idx = np.argmin(np.abs(times - time))
+        if rounding == 'closest':
+            idx = np.argmin(np.abs(times - time))
+        elif rounding == 'up':
+            idx = np.nonzero(times >= time)[0][0]
+        elif rounding == 'down':
+            idx = np.nonzero(times <= time)[0][-1]
+        else:
+            err = "Invalid rounding parameter: %s" % repr(rounding)
+            raise ValueError(err)
+
+        return idx
+
+    def set_time(self, time):
+        """Set the data time index to the time point closest to time
+
+        Parameters
+        ----------
+        time : scalar
+            Time.
+        """
+        idx = self.index_for_time(time)
         self.set_data_time_index(idx)
 
     def _get_colorbars(self, row, col):
@@ -1845,7 +1873,7 @@ def screenshot_single(self, mode='rgb', antialiased=False, row=-1, col=-1):
         brain = self.brain_matrix[row, col]
         return mlab.screenshot(brain._f, mode, antialiased)
 
-    def save_imageset(self, prefix, views,  filetype='png', colorbar='auto',
+    def save_imageset(self, prefix, views, filetype='png', colorbar='auto',
                       row=-1, col=-1):
         """Convenience wrapper for save_image
 
@@ -2034,6 +2062,96 @@ def save_montage(self, filename, order=['lat', 'ven', 'med'],
                 cb.visible = colorbars_visibility[cb]
         return out
 
+    def save_movie(self, dst, tstart=None, tstop=None, step=None,
+                   time_idx=None, montage='current', orientation='h',
+                   border_size=15, colorbar='auto', framerate=10,
+                   codec='mpeg4', row=-1, col=-1, movie_tool='ffmpeg'):
+        """Save a movie (for data with a time axis)
+
+        Parameters
+        ----------
+        dst : str
+            Path at which to sae the movie.
+        tstart : None | float
+            First time point to include (default: all data).
+        tstop : None | float
+            Time point at which to stop the movie (exclusive; default: all
+            data).
+        step : None | int
+            Number of data frames to step forward between movie frames
+            (default 1).
+        time_idx : None | array
+            Index that selects time points form the time axis from which to
+            make the movie. If time_idx is specified, neither of tstart, tstop
+            or tstep should be specified.
+        montage: 'current' | 'single' | list
+            Views to include in the images: 'current' (default) uses the
+            currently displayed image; 'single' uses a single view, specified
+            by the ``row`` and ``col`` parameters; a list can be used to
+            specify a complete montage (see :meth:`save_montage`).
+        orientation: {'h' | 'v'}
+            montage image orientation (horizontal of vertical alignment; only
+            applies if ``montage`` is a flat list)
+        border_size: int
+            Size of image border (more or less space between images)
+        colorbar: None | 'auto' | [int], optional
+            if None no colorbar is visible. If 'auto' is given the colorbar
+            is only shown in the middle view. Otherwise on the listed
+            views when a list of int is passed.
+        framerate : int
+            Framerate of the movie (frames per second).
+        codec : str
+            Codec to use (default 'mpeg4').
+        row : int
+            row index of the brain to use
+        col : int
+            column index of the brain to use
+        movie_tool : 'ffmpeg'
+            Tool to use to convert image sequence into a movie (default:
+            'ffmpeg').
+        """
+        if movie_tool.lower() == 'ffmpeg':
+            if not has_ffmpeg():
+                err = ("FFmpeg is not in the path and is needed for saving "
+                       "movies. Install FFmpeg and try again. It can be "
+                       "downlaoded from http://ffmpeg.org/download.html.")
+                raise RuntimeError(err)
+        else:
+            err = "Currently the only possible movie tool is FFmpeg"
+            raise ValueError(err)
+
+        if tstart is not None:
+            start = self.index_for_time(tstart, rounding='up')
+        else:
+            start = 0
+
+        if tstop is not None:
+            stop = self.index_for_time(tstop, rounding='up')
+        else:
+            stop = self.n_times
+
+        if all(x is None for x in (tstart, tstop, step)):
+            if time_idx is None:
+                time_idx = np.arange(self.n_times)
+        elif time_idx is None:
+            time_idx = np.arange(start, stop, step)
+        else:
+            err = ("Both slice parameters (tstart, tstop, step) and "
+                   "time_idx can not be specified at the same time.")
+            raise TypeError(err)
+
+        n_times = len(time_idx)
+        if n_times == 0:
+            raise ValueError("No time points selected")
+
+        tempdir = mkdtemp()
+        frame_pattern = 'frame%%0%id.png' % (np.floor(np.log10(n_times)) + 1)
+        fname_pattern = os.path.join(tempdir, frame_pattern)
+        self.save_image_sequence(time_idx, fname_pattern, False, row,
+                                 col, montage, orientation, border_size,
+                                 colorbar)
+        ffmpeg(dst, fname_pattern, framerate, codec)
+
     def animate(self, views, n_steps=180., fname=None, use_cache=False,
                 row=-1, col=-1):
         """Animate a rotation.