separate viz functions to be more manageable

ants/viz/__init__.py

@@ -1,14 +1,14 @@
 from .create_tiled_mosaic import create_tiled_mosaic
-from .plot import (
-    plot,
-    movie,
-    plot_hist,
-    plot_grid,
-    plot_ortho,
-    plot_ortho_double,
-    plot_ortho_stack,
-    plot_directory,
-)
+
+from .plot import plot
+from .movie import movie
+from .plot_hist import plot_hist
+from .plot_grid import plot_grid
+from .plot_ortho import plot_ortho
+from .plot_ortho_double import plot_ortho_double
+from .plot_ortho_stack import plot_ortho_stack
+from .plot_directory import plot_directory
+
 from .render_surface_function import render_surface_function
 from .surface import (surf, surf_fold, surf_smooth, get_canonical_views)
 from .volume import (vol, vol_fold)

ants/viz/movie.py

@@ -0,0 +1,91 @@
+"""
+Functions for plotting ants images
+"""
+
+
+__all__ = [
+    "movie"
+]
+
+import fnmatch
+import math
+import os
+import warnings
+
+from matplotlib import gridspec
+import matplotlib.pyplot as plt
+import matplotlib.patheffects as path_effects
+import matplotlib.lines as mlines
+import matplotlib.patches as patches
+import matplotlib.mlab as mlab
+import matplotlib.animation as animation
+from mpl_toolkits.axes_grid1.inset_locator import inset_axes
+
+
+import numpy as np
+
+from .. import registration as reg
+from ..core import ants_image as iio
+from ..core import ants_image_io as iio2
+from ..core import ants_transform as tio
+from ..core import ants_transform_io as tio2
+
+def movie(image, filename=None, writer=None, fps=30):
+    """
+    Create and save a movie - mp4, gif, etc - of the various
+    2D slices of a 3D ants image
+
+    Try this:
+        conda install -c conda-forge ffmpeg
+
+    Example
+    -------
+    >>> import ants
+    >>> mni = ants.image_read(ants.get_data('mni'))
+    >>> ants.movie(mni, filename='~/desktop/movie.mp4')
+    """
+
+    image = image.pad_image()
+    img_arr = image.numpy()
+
+    minidx = max(0, np.where(image > 0)[0][0] - 5)
+    maxidx = max(image.shape[0], np.where(image > 0)[0][-1] + 5)
+
+    # Creare your figure and axes
+    fig, ax = plt.subplots(1)
+
+    im = ax.imshow(
+        img_arr[minidx, :, :],
+        animated=True,
+        cmap="Greys_r",
+        vmin=image.quantile(0.05),
+        vmax=image.quantile(0.95),
+    )
+
+    ax.axis("off")
+
+    def init():
+        fig.axes("off")
+        return (im,)
+
+    def updatefig(frame):
+        im.set_array(img_arr[frame, :, :])
+        return (im,)
+
+    ani = animation.FuncAnimation(
+        fig,
+        updatefig,
+        frames=np.arange(minidx, maxidx),
+        # init_func=init,
+        interval=50,
+        blit=True,
+    )
+
+    if writer is None:
+        writer = animation.FFMpegWriter(fps=fps)
+
+    if filename is not None:
+        filename = os.path.expanduser(filename)
+        ani.save(filename, writer=writer)
+    else:
+        plt.show()