# -*- coding: utf-8 -*-
"""
Created on Thu May 18 15:56:20 2017

@author: david
"""

import numpy as np
import scipy as sp
import scipy.misc
import scipy.ndimage
import matplotlib.pyplot as plt
import os

#calculate a 3D image from a 2D binary mask
def binary(infile, sigma=2):
    
    I = sp.misc.imread(infile).astype(np.bool)
    
    #if the image has more than one channel, just keep the first one
    if(I.ndim == 3):
        I = I[:, :, 0]
    
    L = []
    while np.count_nonzero(I) != 0:
        L.append(I)
        I = sp.ndimage.binary_erosion(I)
        
    #create a 3D image representing the new stack    
    S = np.zeros( (I.shape[0], I.shape[1], len(L) * 2 - 1) )
    
    #for each image in the list
    for i in range(0, len(L)):
        if(i == 0):
            S[:, :, len(L) - 1] = L[0]
        else:
            S[:, :, len(L) - 1 + i] = L[i]
            S[:, :, len(L) - 1 - i] = L[i]
            
    S = sp.ndimage.filters.gaussian_filter(S, sigma)
    return S

#generate a 3D image stack from a 2D binary mask
def binary_stack(infile, outdir, sigma=2):
    outfile_base = os.path.basename(infile)
    outfile_prefix, outfile_ext = os.path.splitext(outfile_base)
    
    S = binary(infile, sigma)

    zcount = len(str(S.shape[2]))
    for f in range(0, S.shape[2]):
        fname = outdir + "/" + outfile_prefix + str(f).zfill(zcount) + outfile_ext
    
        sp.misc.imsave(fname, S[:, :, f])