stackimages

#!/usr/bin/env python
"""
Produces the median of a set of input images after shifting them so that
the largest bright feature is coaligned.

Usage:
  stackimages [-b] [-e P] OUTFN INFNS...
  stackimages [-a] [-c] [-e P] -r RMEAN ROUT INFNS...
  stackimages (-h | --help)
  stackimages (-v | --version)
  
Arguments:
  OUTFN            Filename for the output
  RMEAN            Filename for the output average image
  ROUT             Filename for the output outliers image
  INFNS            Input filenames

Options:
  -b --bkgd        (NOT YET IMPLEMENTED) Align the background instead of the
                   largest feature.
  -e P --exif=P    Look for Exif metadata in the files matching filename
                   pattern P instead of in INFNS.  Useful when INFNS is a
                   set of TIFFs remapped by hugin from a set of JPEGs matching
                   P.  Use quoting to stop the shell from expanding P.
  -r --routliers   Instead of aligning the images, produce robust estimates of
                   their average and outliers.  Alignment can be done ahead of
                   time with hugin (select write remapped images in the
                   stitcher).
  -a --align       Force aligning to the largest feature even when using -r.
                   Not recommended; use hugin or its align_image_stack instead.
  -c --colors      Weight each color seperately when using -r.  This often
                   produces strange color casts.
  -h --help        Show this and exit.
  -v --version     Print version info and exit.
"""
from __future__ import absolute_import, print_function
from glob import glob
import numpy as np
import pyexiv2
from scipy import misc
from scipy import ndimage as ndi
import sys
import ttystatus


def getProgressBar(n, title=' Progress ', period=0.1, counter='x'):
    ts = ttystatus.TerminalStatus(period=period)
    ts.add(ttystatus.ElapsedTime())
    ts.add(ttystatus.Literal(title))
    ts.clear()
    ts.add(ttystatus.ElapsedTime())
    ts.add(ttystatus.Literal(' '))
    ts.add(ttystatus.Counter(counter))
    ts.add(ttystatus.Literal(' of %d ' % n))
    ts.add(ttystatus.PercentDone('done', 'total', decimals=2))
    ts.add(ttystatus.Literal(' done '))
    ts.add(ttystatus.RemainingTime('done', 'total'))
    ts.add(ttystatus.Literal(' '))
    ts.add(ttystatus.ProgressBar('done', 'total'))
    ts['done'] = 0
    ts['total'] = n
    return ts


def read_imstack(fnlist):
    """
    Given a list of image filenames, return a list of numpy arrays for
    the images.
    """
    return [misc.imread(fn) for fn in fnlist]


def cent_of_main_feature(img, medrad=(5,5,1), thresh=32):
    """
    Return the center of the main feature in img.
    """
    med = ndi.median_filter(img.astype(float), size=medrad, mode='constant')
    med = 1.0 / (1.0 + np.exp(thresh - med))
    return ndi.center_of_mass(med)


def align_imstack(instack):
    """
    Given a list of images of the same size and dimensionality, return
    a list of images that have been cropped and shifted to the nearest pixel.
    """
    nim = len(instack)
    ts = getProgressBar(nim, ' Getting centroids: ')
    wi, hi = instack[0].shape[:2]
    cms = []
    for i in range(nim):
        ts['x'] = i
        cms.append(cent_of_main_feature(instack[i]))
        ts['done'] += 1
    cms = np.array(cms)
    print("\ncenters of mass:")
    print(cms)
    cms -= np.mean(cms, axis=0)
    #TODO: Account for color shifts.
    print("Nominal shifts:")
    print(cms)
    cms = np.round(cms).astype(np.int16)
    mincx = min(cms[:, 0])
    mincy = min(cms[:, 1])
    wo = wi - max(cms[:, 0]) + mincx
    ho = hi - max(cms[:, 1]) + mincy
    outstack = np.zeros((nim, wo, ho, 3), instack[0].dtype)
    for i in range(nim):
        d = cms[i, :2]
        xstart = d[0] - mincx
        xstop = xstart + wo
        ystart = d[1] - mincy
        ystop = ystart + ho
        inim = instack[i]
        outstack[i] = inim[xstart:xstop, ystart:ystop, :]
    return outstack


def med_filt_fnlist(fnlist):
    return np.median(align_imstack(read_imstack(fnlist)), axis=0)


def robust_mean_and_outliers(aligned_stack, center_of_stack=None, wexp=4,
                             scale=32.0, softener=0.5, bindcolors=True):
    """
    Calculates a weighted mean and antimean of aligned_stack.

    Parameters
    ----------
    aligned_stack: list of images
        The images can be a mix of grayscale m x n arrays or RGB m x n x 3 arrays,
        but they must be aligned.
    center_of_stack: m x n or m x n x 3 array or None
        An array representing the nominal center of the stack.
        If None np.median(aligned_stack, axis=0) will be used.
    wexp: float
        Exponent to use for the weighting of the mean and antimean.
        The higher it is the more the outliers will be dominated by the intensities
        farthest from center_of_stack's, and the closer the mean will be to a simple
        average of intensities in the neighborhood of center_of_stack.
    scale: float
        The scale of the intensity difference separating outliers from the main
        distribution.
    softener: float > 0
        This reduces the noise amplification in routliers by mixing some of rmean into
        it, especially when the outlier amplitude is < scale.  It also avoids a 0/0 in
        the estimate when a pixel's value is the same for all images, which can happen
        since the pixel values are integers.
    bindcolors: bool
        If False, colors are weighted separately, possibly producing odd results.

    Returns
    -------
    rmean: np.average(aligned_stack, weights=1.0/(1.0 + z**wexp)),
        where z = |im - center| / scale
    routliers: np.average(aligned_stack,
                          weights=(softener**2 + (z**wexp)**2)**0.5/(1.0 + z**wexp))
    """
    ni = len(aligned_stack)
    nx, ny, nc = aligned_stack[0].shape

    if not nc % 2:  # Assume the last channel is alpha, i.e. RGBA
        nc -= 1
        mask_by_alpha = True
    else:
        mask_by_alpha = False

    # Memory problems...memmap helps but is still thrashy when the stack is
    # processed in toto.  The problem requires handling the ni axis in one go,
    # but nx, ny, and/or nc can be split.
    
    # stack = np.memmap('stack.mmap', dtype=np.uint8, mode='w+',
    #                   shape=(ni, ny, nc))
    # del stack
    # stack = np.memmap('stack.mmap', dtype=np.uint8, mode='r+',
    #                   shape=(ni, ny, nc))

    ts = getProgressBar(nx, ' Initializing arrays: ')

    #cos = np.empty((ny, nc))
    rmean = np.empty((nx, ny, nc), dtype=np.uint8)
    routliers = np.empty((nx, ny, nc), dtype=np.uint8)
    stsl = np.empty((ni, ny, nc), dtype=np.uint8)
    zwexp = np.empty((ni, ny, nc))
    if mask_by_alpha:
        #rmean = np.ma.masked_array(rmean)
        #routliers = np.ma.masked_array(routliers)
        stsl = np.ma.masked_array(stsl)
        stsl.mask = np.zeros_like(stsl)
        zwexp = np.ma.masked_array(zwexp)
        avg = np.ma.average
    else:
        avg = np.average
        
    s2 = softener**2
    for xi in range(nx):
        ts['x'] = xi
        for ind in range(ni):
            # The :nc is needed when masking by alpha, and harmless otherwise.
            stsl[ind] = aligned_stack[ind][xi, :, :nc]
        if center_of_stack is None:
            if mask_by_alpha:
                for ind in range(ni):
                    stsl.mask[ind, :, 0] = aligned_stack[ind][xi, :, nc] == 0
                    for ci in range(1, nc): # Broadcast to the other colors
                        stsl.mask[ind, :, ci] = stsl.mask[ind, :, 0]
                cos = np.ma.median(stsl, axis=0) # Normal median does NOT work.
            else:
                cos = np.median(stsl, axis=0)
        else:
            cos = center_of_stack[xi]

        for ind in range(ni):
            zwexp[ind] = ((stsl[ind] - cos) / scale)**wexp
        if bindcolors:  
            avz = np.mean(zwexp, axis=-1)
            # This aggravatingly more than doubles the runtime, but unfortunately it's
            # probably inevitable.  See
            # http://www.socouldanyone.com/2013/03/converting-grayscale-to-rgb-with-numpy.html
            #zwexp[:,:,:] = avz[:, :, np.newaxis]
            for ci in range(nc):
                zwexp[..., ci] = avz
        
        w = 1.0/(1.0 + zwexp)
        #sw = np.sum(w, axis=0)
        #w[sw == 0] = 1.0
        rmean[xi] = np.round(avg(stsl, weights=w, axis=0)).astype(np.uint8)
        w = (zwexp**2 + s2)**0.5
        routliers[xi] = np.round(avg(stsl, weights=w, axis=0)).astype(np.uint8)
        ts['done'] += 1
    print()  # The progress bar needs a \n at the end.
    return rmean, routliers


def summarize_exif(fnlist, outfn, comment=""):
    """
    Merge the EXIF tags of the images in fnlist, add a comment saying
    what was done, and write the result to outfn.  outfn must already exist.
    """
    out = pyexiv2.ImageMetadata(outfn)
    out.read()
    mid = pyexiv2.ImageMetadata(fnlist[len(fnlist) // 2])
    mid.read()
    
    # Does not work and will crash ipython!
    #out.update(mid)

    for k, v in mid.items():
        try:
            out[k] = v.value  # Use value instead of raw_value.
        except:
            print(k, "could not be copied.")

    times = []
    isos = []
    exps = []
    fnumbers = []
    foci = []
    for f in fnlist:
        metadata = pyexiv2.ImageMetadata(f)
        metadata.read()
        times.append(metadata['Exif.Photo.DateTimeOriginal'].value)
        isos.append(metadata['Exif.Photo.ISOSpeedRatings'].value)
        exps.append(metadata['Exif.Photo.ExposureTime'].value)
        fnumbers.append(metadata['Exif.Photo.FNumber'].value)
        foci.append(metadata['Exif.Photo.FocalLength'].value)
    if not comment:
        comment = "Median of " + ", ".join(fnlist) + "\n"
    if len(set([t.date() for t in times])) == 1:
        comment += "Date: %s\n" % times[0].strftime("%Y-%m-%d")
        comment += "Times: " + ", ".join([t.strftime("%H:%M:%S")
                                          for t in times]) + "\n"
    else:
        comment += "Times: " + ", ".join([t.strftime("%Y-%m-%d %H:%M:%S")
                                          for t in times]) + "\n"
    es = []
    for e in exps:
        if e.denominator > 1:
            es.append("%d/%d" % (e.numerator, e.denominator))
        else:
            es.append("%d" % e.numerator)
    comment += "Exposures: " + ", ".join(es) + "s\n"
    comment += "ISOs: " + ", ".join(map(str, isos)) + "\n"
    if len(set(fnumbers)) == 1:
        comment += "f stop: %.1f\n" % fnumbers[0]
    else:
        comment += "f stops: " + ", ".join(["%.1f" % fstop
                                            for fstop in fnumbers]) + "\n"
    if len(set(foci)) == 1:
        comment += "Focal length: %.2f mm\n" % foci[0]
    else:
        comment += "Focal lengths: " + ", ".join(["%.2f" % flength
                                                  for flength in foci])
    out['Exif.Photo.UserComment'] = comment
    out.write()
    return out
    
if __name__ == '__main__':
    from docopt import docopt

    args = docopt(__doc__, version="1.3.0")

    if args['--exif']:
        exiffns = []
        for pat in args['--exif'].split():
            exiffns += glob(pat)
    else:
        exiffns = args['INFNS']

    comment = "Made with " + " ".join(sys.argv) + "\n"
    if args['--routliers']:
        stack = read_imstack(args['INFNS'])
        if args['--align']:
            stack = align_imstack(stack)
        rmean, routliers = robust_mean_and_outliers(stack,
                                                    bindcolors=not args['--colors'])
        misc.imsave(args['RMEAN'], rmean)
        misc.imsave(args['ROUT'], routliers)
        summarize_exif(exiffns, args['RMEAN'], comment)
        summarize_exif(exiffns, args['ROUT'], comment)
    else:
        medim = med_filt_fnlist(args['INFNS'])
        misc.imsave(args['OUTFN'], medim)
        summarize_exif(exiffns, args['OUTFN'], comment)