Source code for ants.plotting.plot_grid

"""
Functions for plotting ants images
"""


__all__ = [
    "plot_grid"
]

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


[docs] def plot_grid( images, slices=None, axes=2, # general figure arguments figsize=1.0, rpad=0, cpad=0, vmin=None, vmax=None, colorbar=True, cmap="Greys_r", # title arguments title=None, tfontsize=20, title_dx=0, title_dy=0, # row arguments rlabels=None, rfontsize=14, rfontcolor="white", rfacecolor="black", # column arguments clabels=None, cfontsize=14, cfontcolor="white", cfacecolor="black", # save arguments filename=None, dpi=400, transparent=True, # other args **kwargs ): """ Plot a collection of images in an arbitrarily-defined grid Matplotlib named colors: https://matplotlib.org/examples/color/named_colors.html Arguments --------- images : list of ANTsImage types image(s) to plot. if one image, this image will be used for all grid locations. if multiple images, they should be arrange in a list the same shape as the `gridsize` argument. slices : integer or list of integers slice indices to plot if one integer, this slice index will be used for all images if multiple integers, they should be arranged in a list the same shape as the `gridsize` argument axes : integer or list of integers axis or axes along which to plot image slices if one integer, this axis will be used for all images if multiple integers, they should be arranged in a list the same shape as the `gridsize` argument Example ------- >>> import ants >>> import numpy as np >>> mni1 = ants.image_read(ants.get_data('mni')) >>> mni2 = mni1.smooth_image(1.) >>> mni3 = mni1.smooth_image(2.) >>> mni4 = mni1.smooth_image(3.) >>> images = np.asarray([[mni1, mni2], ... [mni3, mni4]]) >>> slices = np.asarray([[100, 100], ... [100, 100]]) >>> ants.plot_grid(images=images, slices=slices, title='2x2 Grid') >>> images2d = np.asarray([[mni1.slice_image(2,100), mni2.slice_image(2,100)], ... [mni3.slice_image(2,100), mni4.slice_image(2,100)]]) >>> ants.plot_grid(images=images2d, title='2x2 Grid Pre-Sliced') >>> ants.plot_grid(images.reshape(1,4), slices.reshape(1,4), title='1x4 Grid') >>> ants.plot_grid(images.reshape(4,1), slices.reshape(4,1), title='4x1 Grid') >>> # Padding between rows and/or columns >>> ants.plot_grid(images, slices, cpad=0.02, title='Col Padding') >>> ants.plot_grid(images, slices, rpad=0.02, title='Row Padding') >>> ants.plot_grid(images, slices, rpad=0.02, cpad=0.02, title='Row and Col Padding') >>> # Adding plain row and/or column labels >>> ants.plot_grid(images, slices, title='Adding Row Labels', rlabels=['Row #1', 'Row #2']) >>> ants.plot_grid(images, slices, title='Adding Col Labels', clabels=['Col #1', 'Col #2']) >>> ants.plot_grid(images, slices, title='Row and Col Labels', rlabels=['Row 1', 'Row 2'], clabels=['Col 1', 'Col 2']) >>> # Making a publication-quality image >>> images = np.asarray([[mni1, mni2, mni2], ... [mni3, mni4, mni4]]) >>> slices = np.asarray([[100, 100, 100], ... [100, 100, 100]]) >>> axes = np.asarray([[0, 1, 2], [0, 1, 2]]) >>> ants.plot_grid(images, slices, axes, title='Publication Figures with ANTsPy', tfontsize=20, title_dy=0.03, title_dx=-0.04, rlabels=['Row 1', 'Row 2'], clabels=['Col 1', 'Col 2', 'Col 3'], rfontsize=16, cfontsize=16) """ def mirror_matrix(x): return x[::-1, :] def rotate270_matrix(x): return mirror_matrix(x.T) def rotate180_matrix(x): return x[::-1, ::-1] def rotate90_matrix(x): return mirror_matrix(x).T def flip_matrix(x): return mirror_matrix(rotate180_matrix(x)) def reorient_slice(x, axis): if axis != 1: x = rotate90_matrix(x) if axis == 1: x = rotate90_matrix(x) x = mirror_matrix(x) return x def slice_image(img, axis, idx): if axis == 0: return img[idx, :, :].numpy() elif axis == 1: return img[:, idx, :].numpy() elif axis == 2: return img[:, :, idx].numpy() elif axis == -1: return img[:, :, idx].numpy() elif axis == -2: return img[:, idx, :].numpy() elif axis == -3: return img[idx, :, :].numpy() else: raise ValueError("axis %i not valid" % axis) if isinstance(images, np.ndarray): images = images.tolist() if not isinstance(images, list): raise ValueError("images argument must be of type list") if not isinstance(images[0], list): images = [images] if slices is None: one_slice = True if isinstance(slices, int): one_slice = True if isinstance(slices, np.ndarray): slices = slices.tolist() if isinstance(slices, list): one_slice = False if not isinstance(slices[0], list): slices = [slices] nslicerow = len(slices) nslicecol = len(slices[0]) nrow = len(images) ncol = len(images[0]) if rlabels is None: rlabels = [None] * nrow if clabels is None: clabels = [None] * ncol if not one_slice: if (nrow != nslicerow) or (ncol != nslicecol): raise ValueError( "`images` arg shape (%i,%i) must equal `slices` arg shape (%i,%i)!" % (nrow, ncol, nslicerow, nslicecol) ) fig = plt.figure(figsize=((ncol + 1) * 2.5 * figsize, (nrow + 1) * 2.5 * figsize)) if title is not None: basex = 0.5 basey = 0.9 if clabels[0] is None else 0.95 fig.suptitle(title, fontsize=tfontsize, x=basex + title_dx, y=basey + title_dy) if (cpad > 0) and (rpad > 0): bothgridpad = max(cpad, rpad) cpad = 0 rpad = 0 else: bothgridpad = 0.0 gs = gridspec.GridSpec( nrow, ncol, wspace=bothgridpad, hspace=0.0, top=1.0 - 0.5 / (nrow + 1), bottom=0.5 / (nrow + 1) + cpad, left=0.5 / (ncol + 1) + rpad, right=1 - 0.5 / (ncol + 1), ) if isinstance(vmin, (int, float)): vmins = [vmin] * nrow elif vmin is None: vmins = [None] * nrow else: vmins = vmin if isinstance(vmax, (int, float)): vmaxs = [vmax] * nrow elif vmax is None: vmaxs = [None] * nrow else: vmaxs = vmax if isinstance(cmap, str): cmaps = [cmap] * nrow elif cmap is None: cmaps = [None] * nrow else: cmaps = cmap for rowidx, rvmin, rvmax, rcmap in zip(range(nrow), vmins, vmaxs, cmaps): for colidx in range(ncol): ax = plt.subplot(gs[rowidx, colidx]) if colidx == 0: if rlabels[rowidx] is not None: bottom, height = 0.25, 0.5 top = bottom + height # add label text ax.text( -0.07, 0.5 * (bottom + top), rlabels[rowidx], horizontalalignment="right", verticalalignment="center", rotation="vertical", transform=ax.transAxes, color=rfontcolor, fontsize=rfontsize, ) # add label background extra = 0.3 if rowidx == 0 else 0.0 rect = patches.Rectangle( (-0.3, 0), 0.3, 1.0 + extra, facecolor=rfacecolor, alpha=1.0, transform=ax.transAxes, clip_on=False, ) ax.add_patch(rect) if rowidx == 0: if clabels[colidx] is not None: bottom, height = 0.25, 0.5 left, width = 0.25, 0.5 right = left + width top = bottom + height ax.text( 0.5 * (left + right), 0.09 + top + bottom, clabels[colidx], horizontalalignment="center", verticalalignment="center", rotation="horizontal", transform=ax.transAxes, color=cfontcolor, fontsize=cfontsize, ) # add label background rect = patches.Rectangle( (0, 1.0), 1.0, 0.3, facecolor=cfacecolor, alpha=1.0, transform=ax.transAxes, clip_on=False, ) ax.add_patch(rect) tmpimg = images[rowidx][colidx] if isinstance(axes, int): tmpaxis = axes else: tmpaxis = axes[rowidx][colidx] if tmpimg.dimension == 2: tmpslice = tmpimg.numpy() tmpslice = reorient_slice(tmpslice, tmpaxis) else: sliceidx = slices[rowidx][colidx] if not one_slice else slices if sliceidx is None: sliceidx = math.ceil(tmpimg.shape[tmpaxis] / 2) tmpslice = slice_image(tmpimg, tmpaxis, sliceidx) tmpslice = reorient_slice(tmpslice, tmpaxis) im = ax.imshow(tmpslice, cmap=rcmap, aspect="auto", vmin=rvmin, vmax=rvmax) ax.axis("off") # A colorbar solution with make_axes_locatable will not allow y-scaling of the colorbar. # from mpl_toolkits.axes_grid1 import make_axes_locatable # divider = make_axes_locatable(ax) # cax = divider.append_axes('right', size='5%', pad=0.05) if colorbar: axins = inset_axes(ax, width="5%", # width = 5% of parent_bbox width height="90%", # height : 50% loc='center left', bbox_to_anchor=(1.03, 0., 1, 1), bbox_transform=ax.transAxes, borderpad=0, ) fig.colorbar(im, cax=axins, orientation='vertical') if filename is not None: filename = os.path.expanduser(filename) plt.savefig(filename, dpi=dpi, transparent=transparent, bbox_inches="tight") plt.close(fig) else: plt.show()