import os.path as op
import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
import seaborn as sns

def plot_mask_distribution(fname, ax=None, hist=False):
    """Plot the distribution of voxel coordinates in a mask file.

    Parameters
    ----------
    fname : string
        path to binary mask file

    Returns
    -------
    ax : matplotlib axis object
        axis with plot on it

    """
    if ax is None:
        ax = plt.subplot(111)
    img = nib.load(fname)
    mask = img.get_data()
    aff = img.get_affine()
    vox = np.where(mask)
    vox = np.vstack([vox, np.ones(mask.sum())])
    coords = np.dot(aff, vox)[:-1]
    colors = sns.get_color_list()[:3]
    for axis, data, color in zip(["x", "y", "z"], coords, colors):
        if hist:
            sns.kdeplot(data, hist=True, label=axis, color=color, ax=ax)
        else:
            sns.kdeplot(data, shade=True, label=axis, color=color, ax=ax)
    ax.legend()
    ax.set_title(op.basename(fname))
    return ax

@@ -1,25 +1,4 @@
-import os.path as op%0A
 impo
@@ -120,22 +120,16 @@
 ion(
-fname, ax=None
+mask_img
 , hi
@@ -195,24 +195,33 @@
 s in a mask 
+image or 
 file.%0A%0A    P
@@ -263,16 +263,33 @@
 : string
+ or nibabel image
 %0A       
@@ -313,16 +313,53 @@
 ask file
+ or image object with data and affine
 %0A%0A    Re
@@ -504,28 +504,109 @@
    i
-mg = nib.load(fname)
+f isinstance(mask_img, basestring):%0A        img = nib.load(mask_img)%0A    else:%0A        img = mask_img
 %0A   
@@ -1066,45 +1066,8 @@
 d()%0A
-    ax.set_title(op.basename(fname))%0A
     

Allow plot_mask to take nibabel image