utilsVisualizationLS.py

"""Convenience functions for easy visualization of 3D volumes
"""
import sys
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.image as mpimg
import matplotlib as mp

def progressbar(it, prefix="", size=60, file=sys.stdout):
    count = len(it)
    def show(j):
        x = int(size*j/count)
        file.write("%s[%s%s] %i/%i\r" % (prefix, "#"*x, "."*(size-x), j, count))
        file.flush()
    show(0)
    for i, item in enumerate(it):
        yield item
        show(i+1)
    file.write("\n")
    file.flush()



def plot_slice_from_axis( volume, slice_idx=0, axis=0, origin='upper', cmap='gray' ):
    """Plot a single slice from a 3D volume

    Parameters
    ----------
    volume : 3D array_like
             3D volume to plot slice from
    slice_idx : int, optional
                Which slice to plot
    axis : int, optional
           Axis perpendicular to image plane. `slice_idx` is along this axis.
    origin : {'lower', 'upper'}, optional
             If 'lower' display row 0 at bottom.
    cmap : str or Colormap instance, optional
           A colormap specification. See
           https://matplotlib.org/gallery/color/colormap_reference.html
    Returns
    -------
    None

    See also
    --------
    view_from_axis
    matplotlib.pyplot.imshow
    """
    if axis == 0:
        im = volume[slice_idx,:,:]
    elif axis == 1:
        im = volume[:,slice_idx,:]
    else:
        im = volume[:,:,slice_idx]

    plt.figure()
    plt.imshow( im, cmap=cmap, origin=origin )
    plt.show()


def plot_histogram( volume, nbin=100, output=False):
    """Plot a histogram of the data based on the full dataset

    Parameters
    ----------
    volume : 3D array_like
             3D volume to plot slice from
    nbin   : int, optional
                Number of bins for the histogram

    Returns
    -------
    n    : 2D array_like
           histogram count
    bins : 2D array_like
           bin edges

    See also
    --------
    matplotlib.pyplot.hist
    """
    plt.figure(figsize=(5, 5))
    Vflat = volume.flatten()
    n, bins, _ = plt.hist(Vflat,nbin)
    plt.title('Histogram')

    if output == True:
        return n, bins

    plt.show()


def plot_orthogonal_slices( volume, z, y, x, origin='lower', cmap='gray', layout=(1,3), show=True ):
    """Alias for subplot_slice_from_axis with slightly different parameter order"""
    return subplot_slice_from_axis(volume, (x, y, z), origin=origin, cmap=cmap, layout=layout, show=show)

def subplot_slice_from_axis( volume, slice_idx, origin='lower', cmap='gray', layout=(1,3), show=True):
    """Plot slices from the three standard orthorgonal axes of a 3D volume

    Parameters
    ----------
    volume : 3D array_like
             3D volume to plot slice from
    slice_idx : int or 3-tuple, optional
                Which slices to plot
    origin : {'lower', 'upper'}, optional
             If 'lower' display row 0 at bottom.
    cmap : str or Colormap instance, optional
           A colormap specification. See
           https://matplotlib.org/gallery/color/colormap_reference.html
    layout : {(2,2), (3,1), (1,3)}, optional
             Number of rows and columns in the layout.
    Returns
    -------
    Figure if show=False else None

    See also
    --------
    plot_orthogonal_slices
    plot_slice_from_axis
    plot_all_slices
    plot_vol

    by bepi@dtu.dk (2019.08.08)
    """

    if type(slice_idx) is int:
        slice_idx = (slice_idx,slice_idx,slice_idx)

    fig = plt.figure(figsize=(layout[1]*5, layout[0]*3.5))

    plt.subplot(layout[0], layout[1], 1)
    plt.imshow( volume[slice_idx[2],:,:], cmap=cmap, origin=origin )
    plt.xlabel('x', fontsize=10)
    plt.ylabel('y', fontsize=10)
    plt.title('z_axis view (slice = ' + str(slice_idx[2]) + ')')
    if volume.dtype != 'bool':
        plt.colorbar()

    plt.subplot(layout[0], layout[1], 2)
    plt.imshow( volume[:,slice_idx[1],:], cmap=cmap, origin=origin )
    plt.xlabel('x', fontsize=10)
    plt.ylabel('z', fontsize=10)
    plt.title('y_axis view (slice = ' + str(slice_idx[1]) + ')')
    if volume.dtype != 'bool':
        plt.colorbar()

    plt.subplot(layout[0], layout[1], 3)
    plt.imshow( volume[:,:,slice_idx[0]], cmap=cmap, origin=origin )
    plt.xlabel('y', fontsize=10)
    plt.ylabel('z', fontsize=10)
    plt.title('x_axis view (slice = ' + str(slice_idx[0]) + ')')
    if volume.dtype != 'bool':
        plt.colorbar()

    if show:
        plt.show()
    else:
        return fig


def plot_all_slices( volume, axis=0 ):
    """Plot all slices from a 3D volume or image stack

    Parameters
    ----------
    volume : 3D array_like
             3D volume to plot slice from
    axis : int, optional
           Axis perpendicular to image plane
    Returns
    -------
    None

    See also
    --------
    view_from_axis
    matplotlib.pyplot.imshow

    """
    n = volume.shape[axis]

    # Aim for a rectangular grid. Should be configurable
    rows = max(1, round(n**(0.5)))
    cols = n // rows + (n % rows > 0)
    idx = [ range(volume.shape[ax]) for ax in range(axis)]

    for i in range(n):
        plt.subplot(rows, cols, i+1)
        plt.imshow( volume[idx + [i,...]], cmap='gray' )


def ind2rgb( labelimage,cmap = plt.cm.get_cmap('nipy_spectral')):
    """Convert indexed image to and rgb image

    Parameters
    ----------
    volume : 3D array_like
             3D indexed image, e.g. produced with ndimage.measurements.label
    cmap :  matplotlib.colors.LinearSegmentedColormap
    
    Returns
    -------
    rgb : labelled image as rgb

    See also
    --------
    ndimage.measurements.label

    by ctri@dtu.dk modified by bepi@dtu.dk (2019.09.17)
    """

    # define colormap
    norm = colors.Normalize( vmin=labelimage.min(), vmax=labelimage.max())
    
    # convert to rgb
    rgba = cmap(norm(labelimage))

    return rgba

def show_vol( volume, axis=0, origin='upper', cmap='gray' ):
    """Show slices of a volume.

    Parameters
    ----------
    volume : 3D array_like
             A 3D volume to display
    axis : int, optional
           Axis perpendicular to image plane
    origin : {'lower', 'upper'}, optional
             If 'lower' display row 0 at bottom.
    cmap : str or Colormap instance
           A colormap specification. See
           https://matplotlib.org/gallery/color/colormap_reference.html
    Returns
    -------
    None

    See also
    --------
    plot_slice_from_axis
    plot_all_slices
    plot_orthogonal_slices

    by vand@dtu.dk. 
    Modified by ctri@dtu.dk (2019.09.11) 
    Modified (added colormap option) by bepi@dtu.dk (2019.09.14)
    """
    # swap axis depending on axis input
    if axis == 1 :
        volume = np.swapaxes(volume,0,1)
    elif axis == 2:
        volume = np.swapaxes(volume,0,2)

    # find range of data
    v_min = volume.min()
    v_max = volume.max()

    dim = volume.shape
    Z = dim[0]
    z = round(0.5*Z)-1

    fig, ax = plt.subplots()
    ax.imshow(volume[z,:,:], cmap=cmap, origin=origin, vmin=v_min, vmax=v_max)
    plt.title('slice z=%i' %z)

    def update_drawing():
        ax.clear()
        ax.imshow(volume[z,:,:], cmap=cmap, origin=origin, vmin=v_min, vmax=v_max)
        plt.draw()
        plt.title('slice z=%i' %z)

    def key_press(event):
        nonlocal z
        if event.key == "up":
            z = min(z+1,Z-1)
        elif event.key == 'down':
            z = max(z-1,0)
        elif event.key == 'right':
            z = min(z+10,Z-1)
        elif event.key == 'left':
            z = max(z-10,0)
        elif event.key == 'pagedown':
            z = min(z+50,Z+1)
        elif event.key == 'pageup':
            z = max(z-50,0)
        update_drawing()

    fig.canvas.mpl_connect('key_press_event', key_press)

    def scroll(event):
        if event.step > 0:
            event.key = 'up'
        elif event.step < 0:
            event.key = 'down'
        key_press(event)

    fig.canvas.mpl_connect('key_press_event', key_press)
    fig.canvas.mpl_connect('scroll_event', scroll)

    instructions = '''Instructions:
    Controls
    --------
        'up-arrow'    : Next slice in volume
        'down-arrow'  : Previous slice in volume
        'right-arrow' : +10 slices
        'left-arrow'  : -10 slices
        'PgUp'        : +50 slices
        'PgDown'      : -50 slices

        scroll wheel  : previous/next slice
        '''
    print( instructions )
    plt.show()


def rand_cmap(nlabels, type='bright', first_color_black=True, last_color_black=False, verbose=True):
    """
    Creates a random colormap to be used together with matplotlib. Useful for segmentation tasks
    :param nlabels: Number of labels (size of colormap)
    :param type: 'bright' for strong colors, 'soft' for pastel colors
    :param first_color_black: Option to use first color as black, True or False
    :param last_color_black: Option to use last color as black, True or False
    :param verbose: Prints the number of labels and shows the colormap. True or False
    :return: colormap for matplotlib

    https://github.com/delestro/rand_cmap
    bepi@dtu.dk
    """
    from matplotlib.colors import LinearSegmentedColormap
    import colorsys
    import numpy as np


    if type not in ('bright', 'soft'):
        print ('Please choose "bright" or "soft" for type')
        return

    if verbose:
        print('Number of labels: ' + str(nlabels))

    # Generate color map for bright colors, based on hsv
    if type == 'bright':
        randHSVcolors = [(np.random.uniform(low=0.0, high=1),
                          np.random.uniform(low=0.2, high=1),
                          np.random.uniform(low=0.9, high=1)) for i in range(nlabels)]

        # Convert HSV list to RGB
        randRGBcolors = []
        for HSVcolor in randHSVcolors:
            randRGBcolors.append(colorsys.hsv_to_rgb(HSVcolor[0], HSVcolor[1], HSVcolor[2]))

        if first_color_black:
            randRGBcolors[0] = [0, 0, 0]

        if last_color_black:
            randRGBcolors[-1] = [0, 0, 0]

        random_colormap = LinearSegmentedColormap.from_list('new_map', randRGBcolors, N=nlabels)

    # Generate soft pastel colors, by limiting the RGB spectrum
    if type == 'soft':
        low = 0.6
        high = 0.95
        randRGBcolors = [(np.random.uniform(low=low, high=high),
                          np.random.uniform(low=low, high=high),
                          np.random.uniform(low=low, high=high)) for i in range(nlabels)]

        if first_color_black:
            randRGBcolors[0] = [0, 0, 0]

        if last_color_black:
            randRGBcolors[-1] = [0, 0, 0]
        random_colormap = LinearSegmentedColormap.from_list('new_map', randRGBcolors, N=nlabels)

    # Display colorbar
    if verbose:
        from matplotlib import colors, colorbar
        from matplotlib import pyplot as plt
        fig, ax = plt.subplots(1, 1, figsize=(15, 0.5))

        bounds = np.linspace(0, nlabels, nlabels + 1)
        norm = colors.BoundaryNorm(bounds, nlabels)

        cb = colorbar.ColorbarBase(ax, cmap=random_colormap, norm=norm, spacing='proportional', ticks=None,
                                   boundaries=bounds, format='%1i', orientation=u'horizontal')

    return random_colormap

def color_skeleton(B,S,D,fig_show=True):
    """color map the skeleton into 0-black 1-white and the positive distances cool and negative warm

    Parameters
    ----------
    B : 3D array_like
        A 3D binary array
    S : 3D array_like
        A 3D skeleton array
    D : 3D array_like
        A 3D distance array
    fig_show : logical, optional
               show the figure produces by the binary image, skeleton, distances and the colormap
    
    Returns
    -------
    U : 3D array_like
        A 3D distance array
    cmap : matplotlib.colors.LinearSegmentedColormap
    cmbar : matplotlib.cm.ScalarMappable
    cbtick : array_like
             position of the labels of the colorbar
    cbticklabel :  str, 
                   list of labels for colormap
    

    by bepi@dtu.dk (2019.09.01) Modified (2019.09.16)
    """
    U = B.astype(np.uint8)
    ds = D[S]
    minmax = np.array([ds.min(), ds.max()])
    ds[ds>0] = 2+126*(minmax[1]-ds[ds>0])/minmax[1]# flipping for symmetry
    ds[ds<0] = 129+126*ds[ds<0]/minmax[0]
    U[S] = ds
    
    #cmap = [0 0 0; 1 1 1; autumn(127);plt.get_cmap('winter_r')[127])];
    # sample the colormaps that you want to use. Use 128 from each so we get 256 colors in total
    colors2 = plt.cm.autumn(np.linspace(0., 1, 127))
    colors3 = plt.cm.winter_r(np.linspace(0, 1, 127))
    color01 = np.array([[0,0,0,1],[1,1,1,1]])

    # combine them and build a new colormap
    colors23 = np.vstack((colors2, colors3))
    colorsall = np.vstack((color01,colors23))

    cmap = mp.colors.LinearSegmentedColormap.from_list('my_colormap', colorsall)
    cmap.set_bad(color= 'black',alpha = '0')
    
    barmap = mp.colors.LinearSegmentedColormap.from_list('my_colormap_1', colors23)
    
    norm = mp.colors.Normalize(vmin=0, vmax=254)
    cmbar = mp.cm.ScalarMappable(cmap=barmap,norm=norm)
    cmbar.set_array([])
    
    cbtl = np.round(10*np.array([minmax[1],0,-minmax[0]]))/10
    cbticklabel = [str(x) for x in cbtl]
    
    cbtick = np.array([1,126,253])
    
    if fig_show:
        fig = plt.figure(figsize=(5,5))
        plt.imshow(U, cmap=cmap)
        cbar= fig.colorbar(cmbar)
        cbar.set_ticks(cbtick)
        cbar.set_ticklabels(cbticklabel)
        plt.show()

    return(U,cmap,cmbar,cbtick,cbticklabel)