imtools.py

# imtools.py - Simple but often used load/save/show routines for images
# 
# Author: Stefan Fuertinger [stefan.fuertinger@gmx.at]
# Created: June 13 2012
# Last modified: <2017-09-21 11:15:41>

from __future__ import division

import numpy as np
import matplotlib.pyplot as plt
import os
import datetime
import shutil
from glob import glob
from string import join

##########################################################################################
def imwrite(figobj,fstr,dpi=None):
    """
    Save a Matplotlib figure camera-ready using a "tight" bounding box

    Parameters
    ----------
    figobj : Matplotlib figure
        Matplotlib figure object to be saved
    fstr : string
        String holding the filename to be used to save the figure. If 
        a specific file format is wanted, provide it with `fstr`, e.g., 
        `fstr = 'output.tiff'`. If `fstr` does not contain a filename extension
        the Matplotlib default (png) will be used. 
    dpi : integer
        The wanted resolution of the output in dots per inch. If `None` the 
        Matplotlib default will be used. 
       
    Returns
    -------
    Nothing : None

    Notes
    -----
    This is a humble attempt to get rid of the huge white areas around plots 
    that are generated by Matplotlib's `savefig` when saving a figure as an 
    image using default values. It tries to mimic 
    `export_fig for MATLAB <http://www.mathworks.com/matlabcentral/fileexchange/23629-export-fig>`_. 
    The result, however, is not perfect yet...

    See also
    --------
    savefig : in the `Matplotlib documentation <http://matplotlib.org/api/pyplot_api.html#matplotlib.pyplot.savefig>`_
    """

    # Check if `figobj` is really a figure
    if type(figobj).__name__ != "Figure":
        raise TypeError("figobj has to be a valid matplotlib Figure object!")

    # Make sure `fstr` doesn't contain weirdness and points to an existing place
    if not isinstance(fstr,(str,unicode)):
        raise TypeError("Output filename has to be a string!")
    fstr = str(fstr)
    if fstr.find("~") == 0:
        fstr = os.path.expanduser('~') + fstr[1:]
    slash = fstr.rfind(os.sep)
    if slash >= 0 and not os.path.isdir(fstr[:slash]):
        raise ValueError('Invalid path for output file: '+fstr+'!')

    # Check if filename extension has been provided
    dt = fstr.rfind('.')
    if dt == -1:
        fname = fstr+'.png'
        ext   = 'png'
    elif len(fstr[dt+1:]) < 2: 
        print "Invalid filename extension: "+fstr[dt:]+" Defaulting to png..."
        fname = fstr[:dt]+'.png'
        ext   = 'png'
    else: 
        fname = fstr
        ext   = fstr[dt+1:]

    # Save the figure using "tight" options for the bounding box
    figobj.savefig(fname,bbox_inches="tight",ppad_inches=0,dpi=dpi,format=ext)

    return

##########################################################################################
def normalize(I,a=0,b=1):
    """
    Re-scale a NumPy ndarray

    Parameters
    ----------
    I: NumPy ndarray
        Array to be normalized
    a : float
        The lower normalization bound. 
        By default `a = 0` (`a` has to satisfy `a < b`)
    b : float
        The upper normalization bound. 
        By default `b = 1` (`b` has to satisfy `b < a`)
       
    Returns
    -------
    In : NumPy ndarray
        Scaled version of the input array `I`, such that `a = In.min()` and 
        `b = In.max()`

    Examples
    --------
    >>> I = array([[-1,.2],[100,0]])
    >>> In = normalize(I,a=-10,b=12)
    >>> In 
        array([[-10.        ,  -9.73861386],
               [ 12.        , -10.        ]])
    """

    # Ensure that `I` is a NumPy-ndarray
    try: tmp = I.size == 1
    except TypeError: raise TypeError('I has to be a NumPy ndarray!')
    if (tmp): raise ValueError('I has to be a NumPy ndarray of size > 1!')

    # If normalization bounds are user specified, check them
    try: tmp = b <= a
    except TypeError: raise TypeError('a and b have to be scalars satisfying a < b!')
    if (tmp):
        raise ValueError('a has to be strictly smaller than b!')
    if np.absolute(a - b) < np.finfo(float).eps:
        raise ValueError('|a-b|<eps, no normalization possible')

    # Get min and max of I
    Imin   = I.min()
    Imax   = I.max()

    # If min and max values of I are identical do nothing, if they differ close to machine precision abort
    if Imin == Imax:
        return I
    elif np.absolute(Imin - Imax) < np.finfo(float).eps:
        raise ValueError('|Imin-Imax|<eps, no normalization possible')

    # Make a local copy of I
    I = I.copy()

    # Here the normalization is done
    I = (I - Imin)*(b - a)/(Imax - Imin) + a

    # Return normalized array
    return I

##########################################################################################
def blendedges(Im,chim):
    """
    Superimpose a (binary) edge set on an gray-scale image using Matplotlib's `imshow`

    Parameters
    ----------
    Im: NumPy 2darray
        Grayscale image (has to be a 2D array)
    chim: NumPy 2darray
        Binary edge map (has to be a 2D array). Note that the edge map must be binary, i.e., 
        it must only contain the values 0 and 1
       
    Returns
    -------
    Nothing : None

    See also
    --------
    imshow : in the `Matplotlib documentation <http://matplotlib.org/api/pyplot_api.html#matplotlib.pyplot.imshow>`_
    Stackoverflow : `This submission <http://stackoverflow.com/questions/2495656/variable-alpha-blending-in-pylab>`_ illustrates how to use variable alpha blending in Matplotlib
    """

    # Sanity checks
    if type(Im).__name__ != "ndarray":
        raise TypeError("Im has to be a NumPy ndarray!")
    else:
        if len(Im.shape) > 2: raise ValueError("Im has to be 2-dimensional!")
        try: Im.shape[1]
        except: raise ValueError("Im has to be an image!")
        if np.isnan(Im).max() == True or np.isinf(Im).max() == True:
            raise ValueError("Im must not contain NaNs or Infs!")

    if type(chim).__name__ != "ndarray":
        raise TypeError("chim has to be a NumPy ndarray!")
    else:
        if len(chim.shape) > 2: raise ValueError("chim has to be 2-dimensional!")
        try: chim.shape[1]
        except: raise ValueError("chim has to be an edge map!")
        if np.isnan(chim).max() == True or np.isinf(chim).max() == True:
            raise ValueError("chim must not contain NaNs or Infs!")
        chim = chim.astype(float)
        chiu = np.unique(chim)
        if chiu.size != 2: raise ValueError("chim has to be binary!")
        if chiu.min() != 0 or chiu.max() != 1: raise ValueError("chim has to be a binary edge map!")

    # Now do something
    plt.imshow(Im,cmap="gray",interpolation="nearest")
    plt.hold(True)
    plt.imshow(mycmap(chim))
    plt.axis("off")
    plt.draw()

    return

##########################################################################################
def mycmap(x):
    """
    Generate a custom color map, setting alpha values to one on edge
    points, and to zero otherwise
    
    Notes
    -----
    This code is based on the suggestion found at this 
    `Stackoverflow thread <http://stackoverflow.com/questions/2495656/variable-alpha-blending-in-pylab >`
    """

    # Convert edge map to Matplotlib colormap (shape (N,N,4))
    tmp = plt.cm.hsv(x)

    # Set alpha values to one on edge points
    tmp[:,:,3] = x

    return tmp


##########################################################################################
def recmovie(figobj=None, movie=None, savedir=None, fps=None):
    """
    Save Matplotlib figures and generate a movie sequences 

    Parameters
    ----------
    figobj : Matplotlib figure
        Figure to base the movie on
    movie : str 
        Filename of the generated movie
    savedir : str
        Output directory for video file or name of directory of/for source image files
    fps : int
        Target frames per second

    Returns
    -------
    Nothing : None

    Examples
    --------
    The command 

    >>> recmovie(figobj) 

    saves the Matplotlib figure-object `figobj` in the default directory `_tmp` as png-image. 
    If the default directory is empty the image will be named `_tmp0000.png`, 
    otherwise the highest number in the png-file-names incremented by one will be
    used as filename. 

    Use 

    >>> recmovie(figobj,savedir="somedir") 

    to save the Matplotlib figure-object `figobj` in the directory defined by the string 
    `savedir`. If the directory does not exist, it will be created. 
    If the directory `savedir` is empty the image will be named `_tmp0000.png`, otherwise the 
    highest number in the png-file-names incremented by one will be used as filename. 

    The command `recmovie()` will attempt to use `mencoder <http://en.wikipedia.org/wiki/MEncoder>`_
    to generate an avi-movie composed of the 
    png-images found in the default directory `_tmp`. The movie's default name will be
    composed of the default prefix `_tmp` and the current date and time. 
    After the movie has been generated the default-directory `_tmp` and its contents 
    will be deleted. 

    Use 

    >>> recmovie(movie="somename") 

    to launch mencoder and generate an avi-video composed of the png-images found in the 
    default directory `_tmp`. 
    The string `movie` will be used as filename for the generated movie (in the above example
    a file `somename.avi` will be created). 
    After the movie has been generated the default-directory `_tmp` and its contents 
    will be deleted. Similarly 

    >>> recmovie(savedir="somedir") 

    will generate an avi-movie composed of the png-images found in the directory specified 
    by the string `savedir`. The movie's name is given by the prefix `_tmp` and the current 
    date and time. 
    After the movie has been generated it will be moved to the directory `savedir`. If 
    a movie-file of the same name exists in `savedir` a *WARNING* is printed and the movie 
    will not be moved. Analogously, 

    >>> recmovie(movie="somename",savedir="somedir") 

    will generate an avi-movie named "somename" composed of the png-images found in the 
    directory specified by the string `savedir`. 
    After the movie has been generated it will be moved to the directory `savedir`. If 
    a movie-file of the same name exists in `savedir` a *WARNING* is printed and the movie 
    will not be moved. 

    **Note:** the command 

    >>> recmovie(figobj,movie="somename",savedir="somedir") 

    will ONLY save the Matplotlib-figure-object `figobj` in the directory defined by the 
    string `savdir`. The optional argument `movie` will be ignored. 

    **Note:** the default-directory, image-format and movie-type can be changed in the source code 
    by editing the variables `prefix`, `imgtype` and `movtype`. 

    See also
    --------
    Matplotlib : a `collection of codes <http://matplotlib.org/examples/animation/index.html>`_ illustrating how to use animation in Matplotlib 
    """

    # Set default values
    prefix    = "_tmp"
    numdigits = "%04d"
    imgtype   = "png"
    movtype   = "avi"
    fpsno     = 25
    
    now         = datetime.datetime.now()
    savedirname = prefix
    moviename   = "movie"+"_"+repr(now.hour)+repr(now.minute)+repr(now.second)

    # Assign defaults
    if movie is None:
        movie = moviename
    if savedir is None:
        savedir = savedirname
    if fps is None:
        fps = fpsno

    # Make sure `figobj` is actually a figure
    if figobj != None:
        if type(figobj).__name__ != "Figure":
            raise TypeError("figobj has to be a valid Matplotlib Figure object!")

    # Check if movie filename makes sense and points to an existing location
    if not isinstance(movie,(str,unicode)):
        raise TypeError("Output filename for movie has to be a string!")
    movie = str(movie)
    if movie.find("~") == 0:
        movie = os.path.expanduser('~') + movie[1:]
    if not os.path.isdir(movie[:movie.rfind(os.sep)]):
        raise ValueError('Invalid path for output filename for movie: '+movie+'!')

    # Make sure `savedir` exists
    if not isinstance(savedir,(str,unicode)):
        raise TypeError('Output filename has to be a string!')
    savedir = str(savedir)
    if savedir.find("~") == 0:
        savedir = os.path.expanduser('~') + savedir[1:]
    if not os.path.isdir(savedir):
        raise ValueError('Invalid path for output file: '+savedir+'!')
        
    # Convert possible float argument to integer, if it does not work raise a TypeError
    try: 
        fps = int(fps) 
    except:
        raise TypeError("fps has to be an integer (see man mencoder for details)!")

    # Check if mencoder is available
    if os.system("which mencoder > /dev/null") != 0:
        print "\n\nWARNING: mencoder was not found on your system - movie generation won't work!!!\n\n"

    # Check if movie already exists, if yes abort
    if len(glob(movie)) != 0:
        errormsg = "Movie %s already exists! Aborting..."%savedir
        raise ValueError(errormsg)

    # If not already existent, create directory savedir
    if len(glob(savedir)) == 0:
        os.mkdir(savedir)

    # If a none-default savedir was chosen, automatically keep images
    # and move movie to this non-standard savedir
    if savedir != savedirname:
        keepimgs  = 1
    else:
        keepimgs = 0

    # List all imgtype-files in directory savedir
    filelist = glob(savedir+os.sep+"*."+imgtype)

    # If we have been called with a figobj save it in savedir
    if figobj != None:

        # If there are already imgtype-files in the directory then filelist!=0
        if len(filelist) != 0:

            # Sort filelist, s.t. the file having the highest no. is the last element
            filelist.sort()
            scounter = filelist[-1]

            # Remove the path and prefix from the last elements filename
            scounter = scounter.replace(savedir+os.sep+prefix,"")

            # Split the name further s.t. it is only number+imgtype
            scounter = scounter.split(".")[0]

            # Convert the string holding the file's no. to an integer
            counter  = int(scounter) + 1

        # No files are present in savedir, start with 0
        else:
            counter = 0

        # Generate the name the file is stored under (prefix+numdigits(counter).imgtype, e.g. _tmp0102.png)
        fname = savedir+os.sep+prefix+numdigits+"."+imgtype
        fname = fname%counter

        # Save the figure using the just generated filename
        figobj.savefig(fname)

    # User wants to generate a movie consisting of imgtyp-files in a directory savedir
    else:

        # Check if there are any files to process in savedir, if not raise an error
        if len(filelist) == 0:
            errormsg = "No %s-images found in directory %s! Aborting..."%(imgtype,savedir)
            raise ValueError(errormsg)

        # This is the standard command used to generate the movie
        command = ["mencoder",
                   "mf://*.png",
                   "-mf",
                   "type=png:w=800:h=600:fps=25",
                   "-ovc",
                   "lavc",
                   "-lavcopts",
                   "vcodec=mpeg4",
                   "-oac",
                   "copy",
                   "-o",
                   "output.avi"]

        # Make necessary changes here (like pointing to the right savedir, imgtype, movie,...)
        command[1]  = "mf://"+savedir+os.sep+"*."+imgtype
        command[3]  = "type="+imgtype+":w=800:h=600:fps="+str(fps)
        command[-1] = movie+"."+movtype

        # Call mencoder to generate movie
        os.system(join(command))

        # If we have been called using the default savedir, erase it (and its contents)
        if keepimgs == 0:
            for f in glob(savedir+os.sep+"*"+imgtype):
                os.unlink(f)
            os.rmdir(savedir)

        # If not don't erase it but (try to) move the generated movie into this savedir
        else:
            try:
                shutil.move(movie+"."+movtype,savedir+os.sep)
            except:
               print "\n\n\nWARNING: Movie %s already exists in directory %s. I won't move it there but keep it here. "%(movie,savedir)