diffusions.py

"""
diffusions: a module providing segmentation algorithm based on diffusion.

The algorithms of this module are based on the "random walker" algorithm.
"""

# Author: Emmanuelle Gouillart <emmanuelle.gouillart@normalesup.org>
# Copyright (c) 2009-2010, Emmanuelle Gouillart
# License: BSD

import warnings

import numpy as np
from scipy import sparse, ndimage
from scipy.sparse.linalg.eigen.arpack import eigen_symmetric
try: 
    from scipy.sparse.linalg.dsolve import umfpack
    u = umfpack.UmfpackContext()
except:
    warnings.warn("""Scipy was built without UMFPACK. Consider rebuilding 
    Scipy with UMFPACK, this will greatly speed up the random walker 
    functions. You may also install pyamg and run the random walker function 
    in amg mode (see the docstrings)
    """)
try:
    from pyamg import smoothed_aggregation_solver
    amg_loaded = True
except ImportError:
    amg_loaded = False 




def my_in1d(ar1, ar2, assume_unique=False):
    if not assume_unique:
        ar1, rev_idx = unique(ar1, return_inverse=True)
        ar2 = np.unique(ar2)
    ar = np.concatenate( (ar1, ar2) )
    # We need this to be a stable sort, so always use 'mergesort'
    # here. The values from the first array should always come before
    # the values from the second array.
    order = ar.argsort(kind='mergesort')
    sar = ar[order]
    equal_adj = (sar[1:] == sar[:-1])
    flag = np.concatenate( (equal_adj, [False] ) )
    indx = order.argsort(kind='mergesort')[:len( ar1 )]

    if assume_unique:
        return flag[indx]
    else:
        return flag[indx][rev_idx]

def unique(ar, return_index=False, return_inverse=False):
    try:
        ar = ar.flatten()
    except AttributeError:
        if not return_inverse and not return_index:
            items = sorted(set(ar))
            return np.asarray(items)
        else:
            ar = np.asanyarray(ar).flatten()
    
    if ar.size == 0:
        if return_inverse and return_index:
            return ar, np.empty(0, np.bool), np.empty(0, np.bool)
        elif return_inverse or return_index:
            return ar, np.empty(0, np.bool)
        else: 
            return ar
        
    if return_inverse or return_index:
        perm = ar.argsort()
        aux = ar[perm]
        flag = np.concatenate(([True], aux[1:] != aux[:-1]))
        if return_inverse:
            iflag = np.cumsum(flag) - 1
            iperm = perm.argsort()
            if return_index:
                return aux[flag], perm[flag], iflag[iperm]
            else:
                return aux[flag], iflag[iperm]
        else:
            return aux[flag], perm[flag]

    else:
        ar.sort()
        flag = np.concatenate(([True], ar[1:] != ar[:-1]))
        return ar[flag]


def test_my_in1d():
    a = np.arange(10)
    b = a[a%2 == 0]
    assert my_in1d(a, b).sum() == 5


if np.__version__ >= '1.4':
    from numpy import in1d
else:
    in1d = my_in1d


#--------- Synthetic data ---------------

def make_2d_syntheticdata(lx, ly=None):
    if ly is None:
        ly = lx
    data = np.zeros((lx, ly)) + 0.1*np.random.randn(lx, ly)
    small_l = int(lx / 5)
    data[lx/2 - small_l:lx/2+small_l, ly/2-small_l:ly/2+small_l] = 1
    data[lx/2 - small_l+1:lx/2+small_l-1, \
         ly/2-small_l+1:ly/2+small_l-1] = \
                        0.1 * np.random.randn(2*small_l-2, 2*small_l-2)
    data[lx/2-small_l, ly/2-small_l/8:ly/2+small_l/8] = 0
    seeds = np.zeros_like(data)
    seeds[lx/5, ly/5] = 1
    seeds[lx/2 + small_l/4, ly/2 - small_l/4] = 2
    return data, seeds

def make_3d_syntheticdata(lx, ly=None, lz=None):
    if ly is None:
        ly = lx
    if lz is None:
        lz = lx
    data = np.zeros((lx, ly, lz)) + 0.1*np.random.randn(lx, ly, lz)
    small_l = int(lx/5)
    data[lx/2-small_l:lx/2+small_l,\
         ly/2-small_l:ly/2+small_l,\
         lz/2-small_l:lz/2+small_l] = 1
    data[lx/2-small_l+1:lx/2+small_l-1,\
         ly/2-small_l+1:ly/2+small_l-1,
         lz/2-small_l+1:lz/2+small_l-1] = 0
    # make a hole
    hole_size = np.max([1, small_l/8])
    data[lx/2-small_l,\
            ly/2-hole_size:ly/2+hole_size,\
            lz/2-hole_size:lz/2+hole_size] = 0
    seeds = np.zeros_like(data)
    seeds[lx/5, ly/5, lz/5] = 1
    seeds[lx/2+small_l/4, ly/2-small_l/4, lz/2-small_l/4] = 2
    return data, seeds


#-----------Laplacian--------------------

def _make_edges_3d(lx, ly=None, lz=None):
    """ Returns a list of edges for a 3D image.
    
        Parameters
        ===========
        lx: integer
            The size of the grid in the x direction.
        ly: integer, optinal
            The size of the grid in the y direction, defaults
            to lx.
        lz: integer, optinal
            The size of the grid in the z direction, defaults
            to lx.
    """
    if ly is None:
        ly = lx
    if lz is None:
        lz = lx
    vertices = np.arange(lx*ly*lz).reshape((lx, ly, lz))
    edges_deep = np.vstack((vertices[:, :, :-1].ravel(),
                            vertices[:, :, 1:].ravel()))
    edges_right = np.vstack((vertices[:, :-1].ravel(), vertices[:, 1:].ravel()))
    edges_down = np.vstack((vertices[:-1].ravel(), vertices[1:].ravel()))
    edges = np.hstack((edges_deep, edges_right, edges_down))
    return edges

def _make_weights_3d(edges, data, beta=130, eps=1.e-6):
    lx, ly, lz = data.shape
    gradients = _make_distances_3d(edges, data)**2 
    weights = np.exp(- beta*gradients / (10*data.std())) + eps
    return weights

def _make_distances_3d(edges, data):
    lx, ly, lz = data.shape
    gradients = np.abs(data[edges[0]/(ly*lz), \
                                 (edges[0] % (ly*lz))/lz, \
                                 (edges[0] % (ly*lz))%lz] - \
                            data[edges[1]/(ly*lz), \
                                 (edges[1] % (ly*lz))/lz, \
                                 (edges[1] % (ly*lz)) % lz])
    return gradients

def _make_adaptive_weights(edges, data):
    print "adaptive"
    gradients = _make_distances_3d(edges, data)
    i_indices = np.hstack((edges[0], edges[1]))
    w = np.hstack((gradients, gradients))
    nb = np.bincount(i_indices).astype(np.float)
    total_weight = np.bincount(i_indices, weights=w)
    sigmas = total_weight / nb
    sigma_of_edges = np.array([sigmas[edges[0]], sigmas[edges[1]]])
    return _make_weights_adaptative(gradients, sigma_of_edges)

def _make_weights_adaptative(gradients, sigma_of_edges, eps=1.e-10):
    sigma_i, sigma_j = sigma_of_edges
    weights = np.exp(- gradients**2 / (sigma_i * sigma_j)) + eps
    return weights 

def _make_laplacian_sparse(edges, weights):
    """
    Sparse implementation
    """
    pixel_nb = len(np.unique(edges.ravel()))
    diag = np.arange(pixel_nb)
    i_indices = np.hstack((edges[0], edges[1]))
    j_indices = np.hstack((edges[1], edges[0]))
    data = np.hstack((-weights, -weights))
    lap = sparse.coo_matrix((data, (i_indices, j_indices)), 
                            shape=(pixel_nb, pixel_nb))
    connect = - np.ravel(lap.sum(axis=1)) 
    lap = sparse.coo_matrix((np.hstack((data, connect)),
                (np.hstack((i_indices,diag)), np.hstack((j_indices, diag)))), 
                shape=(pixel_nb, pixel_nb))
    return lap.tocsc()

def _make_normed_laplacian(edges, weights):
    """
    Sparse implementation
    """
    pixel_nb = len(np.unique(edges.ravel()))
    i_indices = np.hstack((edges[0], edges[1]))
    j_indices = np.hstack((edges[1], edges[0]))
    data = np.hstack((-weights, -weights))
    lap = sparse.coo_matrix((data, (i_indices, j_indices)), 
                            shape=(pixel_nb, pixel_nb))
    w = -np.ravel(lap.sum(axis=1))
    print w.min(), w.max()
    data *= 1. / (np.sqrt(w[i_indices]*w[j_indices]))
    #eps = 0
    #data = np.hstack((data, eps*np.ones_like(diag)))
    #i_indices = np.hstack((i_indices, diag))
    #j_indices = np.hstack((j_indices, diag))
    lap = sparse.coo_matrix((-data, (i_indices, j_indices)),
                            shape=(pixel_nb, pixel_nb))
    return lap.tocsc(), w

def _clean_labels_ar(X, labels):
    labels = np.ravel(labels)
    labels[labels == 0] = X
    return labels

def _buildAB(lap_sparse, labels):
    lx, ly, lz = labels.shape
    labels = labels.ravel()
    labels = labels[labels >= 0]
    total_ind = np.arange(labels.size) 
    unmarked = total_ind[labels == 0]
    seeds_indices = np.setdiff1d(total_ind, unmarked)
    B = lap_sparse[unmarked][:, seeds_indices]
    lap_sparse = lap_sparse[unmarked][:, unmarked]
    nlabels = labels.max()
    Bi = sparse.lil_matrix((nlabels, B.shape[0]))
    for lab in range(1, nlabels+1):
        print lab
        fs = sparse.csr_matrix((labels[seeds_indices] == lab)\
                                                [:, np.newaxis])
        Bi[lab-1, :] = (B.tocsr()* fs)
    return lap_sparse, Bi
    
def _trim_edges_weights(edges, weights, mask):
    inds = np.arange(mask.size)
    inds = inds[mask.ravel()]
    ind_mask = np.logical_and(in1d(edges[0], inds),
                          in1d(edges[1], inds))
    edges, weights = edges[:, ind_mask], weights[ind_mask]
    maxval = edges.max()
    order = np.searchsorted(np.unique(edges.ravel()), np.arange(maxval+1))
    edges = order[edges]
    return edges, weights


def _build_laplacian(data, mask=None, normed=False, beta=50):
    lx, ly, lz = data.shape
    edges = _make_edges_3d(lx, ly, lz)
    if beta is None:
        weights = _make_adaptive_weights(edges, data)
    else:
        weights = _make_weights_3d(edges, data, beta=beta, eps=1.e-10)
    if mask is not None:
        edges, weights = _trim_edges_weights(edges, weights, mask)
    if not normed:
        lap =  _make_laplacian_sparse(edges, weights)
        del edges, weights
        return lap
    else:
        lap, w = _make_normed_laplacian(edges, weights)
        del edges, weights
        return lap, w



#----------- Random walker algorithms (with markers or with prior) -------------

def random_walker(data, labels, beta=130, mode='bf', copy=True):
    """
        Segmentation with random walker algorithm by Leo Grady, 
        given some data and an array of labels (the more labeled 
        pixels, the less unknowns and the faster the resolution)

        Parameters
        ----------

        data : array_like
            Image to be segmented in regions. `data` can be two- or
            three-dimensional.

        labels : array of ints
            Array of seed markers labeled with different integers
            for different phases. Negative labels correspond to inactive
            pixels that do not diffuse (they are removed from the graph).

        beta : float
            Penalization coefficient for the random walker motion
            (the greater `beta`, the more difficult the diffusion).

        mode : {'bf', 'amg'}
            Mode for solving the linear system in the random walker 
            algorithm. `mode` can be either 'bf' (for brute force),
            in which case matrices are directly inverted, or 'amg'
            (for algebraic multigrid solver), in which case a multigrid
            approach is used. The 'amg' mode uses the pyamg module 
            (http://code.google.com/p/pyamg/), which must be installed
            to use this mode.

        copy : bool
            If copy is False, the `labels` array will be overwritten with
            the result of the segmentation. Use copy=False if you want to 
            save on memory.

        Returns
        -------

        output : ndarray of ints
            Array in which each pixel has been attributed the label number
            that reached the pixel first by diffusion.

        Notes
        -----

        The algorithm was first proposed in *Random walks for image 
        segmentation*, Leo Grady, IEEE Trans Pattern Anal Mach Intell. 
        2006 Nov;28(11):1768-83.

        Examples
        --------

        >>> a = np.zeros((10, 10)) + 0.2*np.random.random((10, 10))
        >>> a[5:8, 5:8] += 1
        >>> b = np.zeros_like(a)
        >>> b[3,3] = 1 #Marker for first phase
        >>> b[6,6] = 2 #Marker for second phase
        >>> random_walker(a, b)
        array([[ 1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.],
               [ 1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.],
               [ 1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.],
               [ 1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.],
               [ 1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.],
               [ 1.,  1.,  1.,  1.,  1.,  2.,  2.,  2.,  1.,  1.],
               [ 1.,  1.,  1.,  1.,  1.,  2.,  2.,  2.,  1.,  1.],
               [ 1.,  1.,  1.,  1.,  1.,  2.,  2.,  2.,  1.,  1.],
               [ 1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.],
               [ 1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.]])

    """
    # We work with 3-D arrays
    data = np.atleast_3d(data)
    if copy:
        labels = np.copy(labels)
    labels = labels.astype(np.int)
    # If the array has pruned zones, be sure that no isolated pixels
    # exist between pruned zones (they could not be determined)
    if np.any(labels<0):
        filled = ndimage.binary_propagation(labels>0, mask=labels>=0)
        labels[np.logical_and(np.logical_not(filled), labels == 0)] = -1
        del filled
        marked = ndimage.binary_opening(labels >= 0)
        mask = labels >= 0
        labels[np.logical_and(mask >= 0, np.logical_not(marked))] = -1
    labels = np.atleast_3d(labels)
    if np.any(labels < 0):
        lap_sparse = _build_laplacian(data, mask=labels >= 0, beta=beta)
    else:
        lap_sparse = _build_laplacian(data, beta=beta)
    lap_sparse, B = _buildAB(lap_sparse, labels)
    # We solve the linear system
    # lap_sparse X = B
    # where X[i, j] is the probability that a marker of label i arrives 
    # first at pixel j by diffusion
    if mode == 'bf':
        lap_sparse = lap_sparse.tocsc()
        solver = sparse.linalg.factorized(lap_sparse.astype(np.double))
        X = np.array([solver(np.array((-B[i, :]).todense()).ravel())\
                for i in range(B.shape[0])])
        X= np.argmax(X, axis=0) + 1
        data = np.squeeze(data)
        return (_clean_labels_ar(X, labels)).reshape(data.shape)
    elif mode == 'amg':
        if not amg_loaded:
            print """the pyamg module (http://code.google.com/p/pyamg/)
            must be installed to use the amg mode"""
            raise ImportError
        lap_sparse = lap_sparse.tocsr()
        le = lap_sparse.shape[0]
        mls = smoothed_aggregation_solver(lap_sparse)
        del lap_sparse
        ll = np.zeros(le, dtype=np.int32)
        proba_max = np.zeros(le, dtype=np.float32)
        for i in range(B.shape[0]):
            print i
            x = mls.solve(np.ravel(-B[i, :].todense()).astype(np.float32))
            mask = x > proba_max
            ll[mask] = i
            proba_max[mask] = (x[mask]).astype(np.float32)
            del mask
        del proba_max
        ll = _clean_labels_ar(ll + 1, labels)
        data = np.squeeze(data)
        return ll.reshape(data.shape)





def random_walker_prior(data, prior, mode='bf', gamma=1.e-2):
    """
        Parameters
        ----------

        data : array_like
            Image to be segmented in regions. `data` can be two- or
            three-dimensional.


        prior : array_like
            Array of 1-dimensional probabilities that the pixels 
            belong to the different phases. The size of `prior` is n x s
            where n is the number of phases to be segmented, and s the 
            total number of pixels.

        mode : {'bf', 'amg'}
            Mode for solving the linear system in the random walker 
            algorithm. `mode` can be either 'bf' (for brute force),
            in which case matrices are directly inverted, or 'amg'
            (for algebraic multigrid solver), in which case a multigrid
            approach is used. The 'amg' mode uses the pyamg module 
            (http://code.google.com/p/pyamg/), which must be installed
            to use this mode.

        gamma : float
            gamma determines the absence of confidence into the prior. 
            The smaller gamma, the more the output values will be determined
            according to the prior only. Conversely, the greater gamma, 
            the more continuous the segmented regions will be.

        Returns
        -------

        output : ndarray of ints
            Segmentation of data. The number of phases corresponds to the
            number of lines in prior.

        Notes
        -----

        The algorithm was first proposed in *Multilabel random walker 
        image segmentation using prior models*, L. Grady, IEEE CVPR 2005,
        p. 770 (2005).

        Examples
        --------
        >>> a = np.zeros((40, 40))
        >>> a[10:-10, 10:-10] = 1
        >>> a += 0.7*np.random.random((40, 40))
        >>> p = a.max() - a.ravel()
        >>> q = a.ravel()
        >>> prior = np.array([p, q])
        >>> labs = random_walker_prior(a, prior)
    """
    data = np.atleast_3d(data)
    print "building lap"
    lap_sparse = _build_laplacian(data)
    print "lap ok"
    dia = range(data.size)
    lap_sparse = lap_sparse + sparse.lil_diags(
                            [gamma*prior.sum(axis=0)], [0], lap_sparse.shape)
    del dia
    if mode == 'bf':
        lap_sparse = lap_sparse.tocsc()
        solver = sparse.linalg.factorized(lap_sparse.astype(np.double))
        X = np.array([solver(gamma*label_prior)
                      for label_prior in prior])
    elif mode == 'amg':
        if not amg_loaded:
            print """the pyamg module (http://code.google.com/p/pyamg/)
            must be installed to use the amg mode"""
            raise ImportError
        print "converting"
        lap_sparse = lap_sparse.tocsr()
        print "making mls"
        mls = smoothed_aggregation_solver(lap_sparse)
        del lap_sparse
        print "mls ok"
        X = np.array([mls.solve(gamma*label_prior)
                      for label_prior in prior])
        del mls
    return np.squeeze((np.argmax(X, axis=0)).reshape(data.shape))

def fiedler_vector(data, mask, mode='bf'):
    """
    Compute the second eigenmode of the Laplacian built on data,
    which separates in two main diffusions basins. Useful to separate
    an image in two regions.

    Parameters
    ----------
    data: ndarray of floats
        Array containing two objects stuck together that one 
        wants to separate.

    mask: ndarray of bools
        Mask defining the region of interest to be separated in
        two regions.

    mode : {'bf', 'amg'}
        Mode for computing the eigenmode of the Laplacian. `mode` can 
        be either 'bf' (for brute force) or 'amg' (for algebraic 
        multigrid solver), in which case a multigrid approach is used. 
        The 'amg' mode uses the pyamg module 
        (http://code.google.com/p/pyamg/), which must be installed
        to use this mode.

    Note
    ----
    The algorithm does not work well if more than two objects are
    stuck together, or if mask has more than one connex component.

    References
    ----------
    www.stat.cmu.edu/~cshalizi/350/lectures/15/lecture-15.pdf

    Examples
    --------
    >>> x, y = np.indices((40, 40))
    >>> x1, y1, x2, y2 = 14, 14, 28, 26
    >>> r1, r2 = 11, 10
    >>> mask_circle1 = (x - x1)**2 + (y - y1)**2 < r1**2
    >>> mask_circle2 = (x - x2)**2 + (y - y2)**2 < r2**2
    >>> image = np.logical_or(mask_circle1, mask_circle2)
    >>> v = fiedler_vector(image, image)
    """
    lap, w = _build_laplacian(np.atleast_3d(data), 
                np.atleast_3d(mask), normed=True)
    if mode == 'bf':
        vv = eigen_symmetric(lap, which='LA', k=5)
        print vv[0]
        values = 1. / np.sqrt(w) * vv[1][:, -2]
    if mode == 'amg':
        ml = smoothed_aggregation_solver(lap.tocsr())
        X = np.rand(lap.shape[0], 4)
        X[:, 0] = 1. / np.sqrt(lap.shape[0])
        M = ml.aspreconditioner()
        W, V = sparse.linalg.lobpcg(-lap, X, M=M, tol=1e-8, largest=True)
        print W
        values = V[:, -2]
    result = np.zeros_like(data).astype(np.float)
    result[mask] = values
    return result




#----------- Tests --------------------------------

def test_2d():
    lx = 70
    ly = 100
    data, labels = make_2d_syntheticdata(lx, ly)
    print "making"
    labels_bf = random_walker(data, labels, beta=90)
    data, labels = make_2d_syntheticdata(lx, ly)
    assert (labels_bf.reshape((lx, ly))[25:45, 40:60] == 2).all()
    if amg_loaded:
        labels_amg = random_walker(data, labels, beta=90, mode='amg')
        assert (labels_amg.reshape((lx, ly))[25:45, 40:60] == 2).all()
    return data, labels_bf, labels_amg


def test_2d_inactive():
    lx = 70
    ly = 100
    data, labels = make_2d_syntheticdata(lx, ly)
    labels[10:20, 10:20] = -1
    print "making"
    labels[46:50, 33:38] = -2
    labels = random_walker(data, labels, beta=90)
    assert (labels.reshape((lx, ly))[25:45, 40:60] == 2).all()
    return data, labels

def test_2d_inactive2():
    lx = 70
    ly = 100
    data, labels = make_2d_syntheticdata(lx, ly)
    labels[10:20, 10:20] = -1
    labels[15, 15] = 0
    labels[10, 10] = 0
    labels[11, 11] = 1
    labels = random_walker(data, labels, beta=90)
    assert (labels.reshape((lx, ly))[25:45, 40:60] == 2).all()
    return data, labels


def test_3d():    
    n=30
    lx, ly, lz = n, n, n
    data, labels = make_3d_syntheticdata(lx, ly, lz)
    labels = random_walker(data, labels)
    assert (labels.reshape(data.shape)[13:17,13:17,13:17] == 2).all()
    return data, labels

def test_3d_inactive():
    n=30
    lx, ly, lz = n, n, n
    data, labels = make_3d_syntheticdata(lx, ly, lz)
    old_labels = np.copy(labels)
    labels[5:25, 26:29, 26:29] = -1
    after_labels = np.copy(labels)
    labels = random_walker(data, labels)
    assert (labels.reshape(data.shape)[13:17,13:17,13:17] == 2).all()
    return data, labels, old_labels, after_labels

def test_fiedler():
    x, y = np.indices((40, 40))
    x1, y1, x2, y2 = 14, 14, 28, 26
    r1, r2 = 10, 10
    mask_circle1 = (x - x1)**2 + (y - y1)**2 < r1**2
    mask_circle2 = (x - x2)**2 + (y - y2)**2 < r2**2
    image = np.logical_or(mask_circle1, mask_circle2)
    v = fiedler_vector(image, image)
    vm = v[image]
    # Test that the image is separated in two regions
    # that have almost the same area
    assert np.abs((vm>0).sum() - (vm<0).sum()) <= 2
    return image, v

def test_rw_with_prior():
    a = np.zeros((40, 40))
    a[10:-10, 10:-10] = 1
    a += 0.7*np.random.random((40, 40))
    p = a.max() - a.ravel()
    q = a.ravel()
    prior = np.array([p, q])
    labs = random_walker_prior(a, prior)
    assert (labs[11:-11, 11:-11] == 1).all()
    if amg_loaded:
        labs_amg = random_walker_prior(a, prior, mode='amg')
        assert (labs_amg[11:-11, 11:-11] == 1).all()