onilton/gist:9b197b0003385eefb1e74602b2a0dc0d

## gistfile1.py
#-----------------------------------------------------------------------------------------
#
# Spectral Matting
#
import time
import logging

import numpy as np
from scipy import ndimage, sparse
import scipy.sparse.linalg as sparse_linalg

TEST_SMALL_IMAGE = False # takes the top left corner of the image
CONVERT_TO_HSV = False

def printt(arr):
    for x in range(len(arr)):
        for y in range(len(arr[x])):
            print arr[x][y],
        print


def change_img_colorsys(img, cspace="hsv"):
    """Change img colorspace from rgb to hsv, hls or yiq
    """
    if cspace not in ("hsv", "hls", "yiq"): raise ValueError
    import colorsys

    if cspace == "hsv":   fn = colorsys.rgb_to_hsv
    elif cspace == "hls": fn = colorsys.rgb_to_hls
    elif cspace == "yiq": fn = colorsys.rgb_to_yiq

    new_img = np.empty(img.shape)
    for h in range(len(img)):
        for w in range(len(img[h])):
            new_img[h,w] = fn(*img[h,w])
    return new_img

def add_img_colorsys(img, cspaces=("hsv",)):
    """Add img colorspace(s) to matrix -- turning the image from a w*h*3 image to
    a w*h*6|9|12 image
    """
    if any(cspace not in ("hsv", "hls", "yiq") for cspace in cspaces): raise ValueError
    import colorsys

    new_imgs = []
    for cspace in cspaces:
        if cspace == "hsv":   fn = colorsys.rgb_to_hsv
        elif cspace == "hls": fn = colorsys.rgb_to_hls
        elif cspace == "yiq": fn = colorsys.rgb_to_yiq

        new_imgs.append(np.empty(img.shape))
        for h in range(len(img)):
            for w in range(len(img[h])):
                new_imgs[-1][h,w] = fn(*img[h,w])

    stacked_img = np.copy(img)
    for new_img in new_imgs:
        stacked_img = np.dstack((img, new_img))
    return stacked_img

bname = "woman_small"
save_partial_results = True

#-----------------------------------------------------------------------------------------
# Globals
#
# win_size=1 actually makes a 3x3 window -- more like a radius
#
win_size = 1
win_diam = (win_size*2)+1
neb_size = win_diam**2
epsilon = 0.00001

# 212x165x3 image
# EDU>> I(1:4,1:4, 1)
#    0.1922    0.1961    0.1843    0.1765
#    0.1608    0.1608    0.1451    0.1529
#    0.1882    0.1765    0.1765    0.2039
#    0.2000    0.1961    0.2039    0.2118

#-----------------------------------------------------------------------------------------
# Calculate local window variances
#
# Python-derived image is different to the one stored in my_images.mat
# Therefore just read in woman.png array from my_images.mat to ensure the same as matlab
# See test.py for comparison
# To run matlab:
# plant_medium = double(imread('plant_medium.png'))/255.;
# save('plant_medium.mat', 'plant_medium');
#

img = ndimage.imread("%s.png" % bname).astype(np.float) / 255.0
#img = np.array([float(val) for line in open("woman.mat.textfile") for val in line.strip().split(",")])
#img = np.reshape(img, (212,3,165))
#img = np.transpose(img, axes=(0,2,1))

if TEST_SMALL_IMAGE:
    img = img[:60,:120,:3]

if CONVERT_TO_HSV:
    img = add_img_colorsys(img, cspaces=("hls",))

h,w,c = img.shape
print "loaded image", img.shape, "TEST_SMALL_IMAGE", TEST_SMALL_IMAGE
img_size = w*h

sparse_rows, sparse_cols, sparse_vals = [], [], []

for col in range(win_size,w-win_size,1):
    for row in range(win_size,h-win_size,1):
        winI = img[row-win_size:row+win_size+1,
                   col-win_size:col+win_size+1,
                   :]
        winI = np.transpose(winI, (1,0,2)).reshape(neb_size,c)

        win_mu = winI.mean(axis=0)

        # win_var1 and win_var2 are almost identical
        # numpy variance uses N, numpy covariance uses N-1 (sampling variance?)
        # matlab variance uses N-1, matlab covariance uses N-1
        win_var1 = np.cov(winI.T) # 3x3

        # uses N not N-1 so it's slightly difference to cov
        win_var2 = np.matrix(winI).T*np.matrix(winI)/neb_size - np.matrix(win_mu).T*np.matrix(win_mu)

        # win_var3 uses array notation instead of matrix notation
        # matlab: win_var3 = inv(winI'*winI/neb_size-win_mu*win_mu' +epsilon/neb_size*eye(c));
        # I need to turn win_mu into 2D column and row vectors from 1D vectors
        win_var3 = np.dot(winI.T, winI)/neb_size - np.dot(win_mu[:,np.newaxis], win_mu[np.newaxis,:])

        # win_var2/win_var3 are very close but not identical to matlab
        win_var = win_var3

        assert np.array_equal(win_var2, win_var3)
        assert np.abs(win_var2[0,0] - np.var(winI[:,0])) < 1e-4

        win_var_plus = np.linalg.inv(win_var + (epsilon/neb_size)*np.eye(c))
        # Final win_var_plus is almost equal to matlab win_var_plus
        # minor differences due to numerical instability?

        n_winI = winI - win_mu

        # matrix version converted to array
        # tvals = np.asarray((1 + np.matrix(n_winI) * np.matrix(win_var_plus) * np.matrix(n_winI).T) / neb_size)
        tvals = (1 + np.dot(np.dot(n_winI, win_var_plus), n_winI.T)) / neb_size

        assert winI.shape == (neb_size,c)
        assert win_var_plus.shape == (c,c)
        assert tvals.shape == (neb_size,neb_size)
        #print "tvals", row, col
        #print tvals

        # rc:
        # [[-1 -1 -1  0  0  0  1  1  1]
        #  [-1  0  1 -1  0  1 -1  0  1]]
        # move out of loop
        rc = np.array([np.repeat(range(-win_size,win_size+1),win_diam),
                       np.tile(range(-win_size,win_size+1),win_diam)])

        for t_r in range(len(tvals)):
            new_row1 = row + rc[1][t_r]
            new_col1 = col + rc[0][t_r]
            ind1 = new_col1*h + new_row1
            #print "new_row1", row, rc[0][t_r], t_r
            #print "ind1", ind1, new_row1, new_col1

            for t_c in range(len(tvals[t_r])):
                new_row2 = row + rc[1][t_c]
                new_col2 = col + rc[0][t_c]
                ind2 = new_col2*h + new_row2
                sparse_rows.append(ind1)
                sparse_cols.append(ind2)
                sparse_vals.append(tvals[t_r][t_c])

# sparse array is very similar but not exactly the same as matlab
# things start to diverge at win_var3, and win_var_plus, due to numerical instability
print "sparse"
print sparse_vals[:10], sparse_vals[-10:]
print sparse_rows[:10], sparse_rows[-10:]
print sparse_cols[:10], sparse_cols[-10:]
print w*h


def sum_identical_row_cols(sparse_rows, sparse_cols, sparse_vals):
    """DEPRECATED
    Sum values that have the same row and column (there is some repetition)
    to exactly match matlab code -- it's possible to remove it later if it was
    not on purpose in matlab.

    http://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.coo_matrix.html
    By default when converting [from COO] to CSR or CSC format, duplicate (i,j) entries will be summed together.
    This facilitates efficient construction of finite element matrices and the like. (see example)
    """
    assert len(sparse_rows) == len(sparse_cols) == len(sparse_vals)
    zipped_rc = zip(sparse_rows, sparse_cols, sparse_vals, range(len(sparse_rows)))
    zipped_rc.sort()
    new_zipped_rc = []
    i = 0
    while i < len(zipped_rc):
        row, col, val, ind = zipped_rc[i]

        j = i + 1
        while j < len(zipped_rc):
            jrow, jcol, jval, jind = zipped_rc[j]
            if jrow == row and jcol == col:
                val = val + jval
                j += 1
            else:
                break
        new_zipped_rc.append((row, col, val, ind))
        i = j

    # go back to being sorted by the original index
    new_zipped_rc.sort(key=lambda x:x[3])
    sparse_rows = [i[0] for i in new_zipped_rc]
    sparse_cols = [i[1] for i in new_zipped_rc]
    sparse_vals = [i[2] for i in new_zipped_rc]

    return sparse_rows, sparse_cols, sparse_vals

def test_sum_identical():
    sparse_rows = [0,0,0,1,1,1,2,2,2]
    sparse_cols = [0,0,0,1,2,1,2,1,2]
    sparse_vals = [1,2,3,4,5,6,7,8,9]
    assert sum_identical_row_cols(sparse_rows, sparse_cols, sparse_vals) == \
           ([0, 1, 1, 2, 2], [0, 1, 2, 2, 1], [6, 10, 5, 16, 8])

test_sum_identical()

#sparse_rows, sparse_cols, sparse_vals = sum_identical_row_cols(sparse_rows, sparse_cols, sparse_vals)

#-----------------------------------------------------------------------------------------
# Make a sparse matrix
# http://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.coo_matrix.html
# By default when converting [from COO] to CSR or CSC format, duplicate (i,j) entries will be summed together.
# This facilitates efficient construction of finite element matrices and the like. (see example)
#
A = sparse.coo_matrix((sparse_vals, (sparse_rows,sparse_cols)), shape=(w*h,w*h))
A = sparse.csc_matrix(A)

# works the same as matlab sumA except matlab is not sparse
sumA = A.sum(1)

import scipy.sparse
A = scipy.sparse.spdiags(sumA.ravel(1),0,img_size,img_size)-A;

#-----------------------------------------------------------------------------------------
# Create the L (Laplacian) matrix (eigenvectors)
# matlab calcMattingComponent -- produces the L matrix
#
import scipy.sparse.linalg

eigs_num = 50
nclust = 10
maxItr = 20
exp_a = 0.9
I = img
L = A

st = time.time()
d00, u00 = scipy.sparse.linalg.eigs(L, k=eigs_num, sigma=0)
print "eig time", time.time()-st

if d00[0] > d00[-1]:
    raise Exception, "d00 is not sorted correctly -- matlab code reverses under these \
                      circumstances but it may never happen anyway"

print d00
# Although the eigenvalues do return in approximately the right order, there can be
# misordering in there. In a test with 50 eigenvalues there were 2 out of order in the middle
# assert np.array_equal(np.sort(d00), d00), "%s %s" % (np.sort(d00), d00)

# FIXFIX d00 should be 51 not 50 (like in matlab) -- check this
d = np.diag(d00)
u = u00

print "u00", u00.shape
print u00[0:10,0:10]
print u00[-10:,-10:]
print "d00", d00.shape
print d00


s_eigs_num = 20
eigs_weights = 1.0/(d00[1:s_eigs_num+1]**0.5)
print "eigs_weights", eigs_weights

eigss=u00[:,1:s_eigs_num+1]*eigs_weights;
eigss=eigss.astype(float)
print "eigss", eigss.shape, eigss

#-----------------------------------------------------------------------------------------
# Apply KMEANS
#
# scipy.cluster.vq.kmeans does not work since it uses distance, not squared distance
# kmeans2 appears to work. vq also does not work, possibly for the same reason
# use labels instead of vq "codebook"
# codebook, distortion = kmeans(eigss, nclust)
# code, dist = vq(eigss, codebook)
#
from scipy.cluster.vq import (kmeans2, vq)
#from z_kmeans import kmeans

#
# there is quite a lot of variability in total dist so rerun kmeans2 several times to
# get a good answer
#
dist = np.inf
replicates = 2
for i in range(replicates):
    st = time.time()
    pre_centroids, pre_labels = kmeans2(eigss, nclust)
    #pre_centroids, pre_labels, iters = kmeans(eigss, nclust)
    print "kmeans", int(time.time()-st)
    #print "centroid1", pre_centroids.shape, pre_labels.shape

    #code, dist = vq(eigss, centroids)
    #print "kmeans stuff, iter", i

    pre_dist = 0
    for l in range(nclust):
        pre_dist += np.sum(np.sum((eigss[pre_labels==l] - pre_centroids[l])**2, axis=1))

    # print "centroid", pre_centroids.shape, "labels", pre_labels.shape, "iters", iters, \
    #       "dist", pre_dist

    print "centroid", pre_centroids.shape, "labels", pre_labels.shape, "dist", pre_dist

    if pre_dist < dist:
        dist, centroids, labels = pre_dist, pre_centroids, pre_labels


# Remove randomness from kmeans for testing
TESTING_USE_MATLAB_KMEANS = False
if TESTING_USE_MATLAB_KMEANS:
    labels = np.array([int(i)-1 for i in open("kmeans.out")])

print "labels", labels.shape, " ".join([str(i) for i in labels])

# convert labels to indicator array
x = np.zeros((len(labels),nclust), np.float64) - 1
for l in range(nclust):
    x[:,l] = (labels==l)

assert sum(x.ravel()<0)==0, "every point must have a cluster"

print "x", x.shape, x[:15,:15]

#-----------------------------------------------------------------------------------------
# Set up for iterateProjComps
#
thr_e=1e-10
w1=0.3
w0=0.3
e1 = (w1**exp_a)*np.maximum(np.abs(x-1),thr_e)**(exp_a-2)
e0 = (w0**exp_a)*np.maximum(np.abs(x),thr_e)**(exp_a-2)

scld=1
eig_vectors = u[:,0:eigs_num]
eig_values = d[0:eigs_num]

omaxItr = maxItr
maxItr = int(np.ceil(omaxItr/4))

#-----------------------------------------------------------------------------------------
# inputs to iterateProjComps
# e0, e1
# L, eig_vectors, eig_values
#

def iterateProjComps(e0,e1,x, TEST_IPC=False):
    """This function works exactly the same as matlab"""
    #global x, y

    '''if TEST_IPC:
        # e0: 1600x10, e1: 1600x10
        # L: 1600x1600, eig_values: 50, eig_vectors: 1600x50
        # u: 1600x50, d: 50,
        # FIXFIX d should be 51 according to matlab, not 50
        #
        print e0.shape, e1.shape
        print L.shape, eig_values.shape, eig_vectors.shape
        print u.shape, d.shape'''

    # eig_values needs to be a diagonal matrix, not a 1D array, for this fn. to work
    assert len(eig_values.shape) == 2 and eig_values.shape[0]==eig_values.shape[1]

    # default params
    exp_a=0.9; thr_e=1e-10; w1=0.3; w0=0.3

    for itr in range(maxItr):
        tA = np.zeros(((nclust-1)*eigs_num, (nclust-1)*eigs_num))
        tb = np.zeros((nclust-1)*eigs_num)
        eig_vectors = u[:,:eigs_num]

        for k in range(nclust-1):
            weighted_eigs = np.tile(e1[:,k] + e0[:,k], (eigs_num, 1)).T * eig_vectors

            tA[k*eigs_num:(k+1)*eigs_num, k*eigs_num:(k+1)*eigs_num] = np.dot(eig_vectors.conj().T, weighted_eigs) + scld*eig_values
            tb[k*eigs_num:(k+1)*eigs_num] = np.dot(eig_vectors.conj().T, e1[:,k])

        k = nclust-1 # last element: nclust in matlab, nclust-1 here
        weighted_eigs = np.tile(e1[:,k] + e0[:,k], (eigs_num, 1)).T*eig_vectors

        ttA = np.dot(eig_vectors.conj().T, weighted_eigs) + scld*eig_values

        # reshape(-1) turns a 1D column (made from the L matrix), e.g. [[1][1][1]] into a 1D array
        # replace np.array(scld*(np.sum(np.dot(eig_vectors.T,  L.todense()), axis=1))).reshape(-1)
        # which dies due to L.todense() being too big. There are some accumulated differences
        # though once you go above 1000x1000

        #print scipy.sparse.csc_matrix(eig_vectors.T)
        #print type(L), L.shape
        #print type(scipy.sparse.csc_matrix(eig_vectors.T)), scipy.sparse.csc_matrix(eig_vectors.T).shape
        #print np.dot(scipy.sparse.csc_matrix(eig_vectors.T),  L)
        #print "Yes"
        #print np.array(scld*(np.dot(scipy.sparse.csc_matrix(eig_vectors.T),  L)))
        #print np.array(scld*(np.dot(scipy.sparse.csc_matrix(eig_vectors.T),  L).sum(axis=1)))
        #print np.array(scld*(np.dot(scipy.sparse.csc_matrix(eig_vectors.T),  L).sum(axis=1))).reshape(-1)

        #ttb = np.dot(eig_vectors.conj().T, e0[:,k]) + \
        #      np.array(scld*(np.dot(scipy.sparse.csc_matrix(eig_vectors.T),  L).sum(axis=1))).reshape(-1)

        ttb = (eig_vectors.conj().T).dot(e0[:,k]) + \
              np.array(scld*( (scipy.sparse.csc_matrix(eig_vectors.T).dot(L)).sum(axis=1) )).reshape(-1)

        # tA is 450x450, tb is (450,)
        tA = tA + np.tile(ttA,(nclust-1,nclust-1))
        tb = tb + np.tile(ttb,nclust-1)

        # a\b in matlab translates to np.linalg.solve(a, b) (for square a)
        y = np.reshape(np.linalg.solve(tA,tb), (nclust-1, eigs_num)).T

        x = np.dot(u[:,0:eigs_num], y)

        # each row is supposed to sum to one, but each element in x is >1e4,
        # so it may be a bug in the matlab code. The column added onto the end is
        # approximately -1e5 in my test case
        #x(:,nclust)=1-sum(x(:,1:nclust-1),2);
        x = np.column_stack([x, 1 - np.sum(x[:,0:nclust-1],axis=1)])

        if itr < maxItr-1:
            e1 = (w1**exp_a)*np.maximum(np.abs(x-1),thr_e)**(exp_a-2)
            e0 = (w0**exp_a)*np.maximum(np.abs(x),thr_e)**(exp_a-2)

    # return e0, e1 or they go out of scope
    return e0, e1, x


# Test the iterateProjComps function using test data
TEST_ITERATE = False
if TEST_ITERATE:
    F = 1600
    eigs_num = 50
    nclust = 10
    scld = 1
    maxItr = 20
    e0 = 0 * np.reshape(range(1,F*nclust+1), (nclust,F)).T / 1e6
    e1 = 0 * np.reshape(range(1,F*nclust+1), (nclust,F)).T / 1e6
    L  = scipy.sparse.csc_matrix(np.mod(np.reshape(range(1,F*F+1), (F,F)).T, 0.927)-.5)
    eig_values = np.diag(np.array(range(1,eigs_num+1)) / 1e3)
    eig_vectors = None #np.reshape(range(1,F*eigs_num+1), (eigs_num,F)).T / 1e6
    u = (np.mod(np.reshape(range(1,F*eigs_num+1), (eigs_num,F)).T, 0.859)-.5) / 1e2
    d = None #np.array(range(1,eigs_num+2)) / 1e4


for nzitr in range(3):
  # compute matting component with sparsity prior.
  e0, e1, x = iterateProjComps(e0, e1, x)

  # remove matting components that are close to zero.
  nzii = np.where(np.max(np.abs(x), axis=0) > 0.1)[0]
  nclust = len(nzii)
  x = x[:,nzii]
  print "nzii", nzii
  print "x", x.shape, x
  e1 = (w1**exp_a)*np.maximum(np.abs(x-1),thr_e)**(exp_a-2)
  e0 = (w0**exp_a)*np.maximum(np.abs(x),thr_e)**(exp_a-2)

# Final iterateProjComps
# matlab and numpy do diverge for e0, e1, x, but not enough to indicate a bug
# I am still not sure what causes this
e0, e1, x = iterateProjComps(e0, e1, x)

alpha_comps = np.copy(x)
print alpha_comps
print "e0", e0[:10,:10]
print "e1", e1[:10,:10]
print "x", x[:10,:10]

for l in range(nclust):
  if save_partial_results:
      compname = '%s_components%.2d.png' % (bname, l)
      rotated = ((alpha_comps[:,l].T).reshape(w,h).T).astype(np.float128)

      scipy.misc.imsave(compname, rotated) #np.reshape(alpha_comps.astype(np.float128)[:,l], (h,w)).T)

      print "l", l, np.sum(alpha_comps.astype(np.float128)[:,l])

#-----------------------------------------------------------------------------------------
# unsupervised GroupMatting
#
# alpha=unsp_GroupMattingComponents(L, alpha_comps, I, bname, save_partial_results);

SizeThresh = 0.35
nclust = alpha_comps.shape[1]
# FIXFIX check that n,m match w,h for non-square matrices
[n,m,c] = I.shape
assert n==h and m==w, "n %s m %s w %s h %s" % (n,m,w,h)
N = n*m

perrM = np.zeros((nclust, nclust))
clust_size = np.zeros(nclust)
for i in range(nclust):
    clust_size[i] = np.sum(np.sum(alpha_comps[:,i]))
    for j in range(i+1):
      # array * sparse_matrix == dot(array, array)
      # dot(array, sparse_matrix) == array of sparse matrices(!)
      terrp = np.dot(alpha_comps[:,i].T * L, alpha_comps[:,j])
      perrM[i,j]=terrp
      perrM[j,i]=terrp

# clust_size matches matlab +/-3%, but definitely in the right order
# perrM also matches +/-3% or so
free_clust_len = nclust

# find grouping for which the sum of costs is small, and
# whose size is within SizeThresh to (1-SizeThresh) of the image.
scr = np.ones(2**(free_clust_len-1))
sizes = np.ones(2**(free_clust_len-1))


for subg in range(2**(free_clust_len-1)):
  #bits = ind2bin(subg, free_clust_len)
  bits = np.array([i for i in reversed(np.binary_repr(subg+1, width=free_clust_len))], np.float128)
  scr[subg] = np.dot(np.dot(bits.T,perrM),bits)
  sum_alpha = np.sum(np.dot(clust_size,bits))
  sizes[subg] = sum_alpha
  if ((sum_alpha < SizeThresh*N) or (sum_alpha > (1-SizeThresh)*N)):
      scr[subg] = 1e7

print "scr", scr.shape, scr
print "sizes", sizes.shape, sizes
# ok up to here -- scr and sizes match almost perfectly

def decideFB(alpha):
    """ foreground vs background """
    falpha = np.copy(alpha)

    ones_mat = np.ones(alpha.shape)
    # flatten() for row-wise ; flatten(1) is column-wise
    # in this case it doesn't matter since the border is symmetrical
    border_size = np.sum(ones_mat.flatten(1)) - np.sum(np.sum(ones_mat[1:-1,1:-1]))
    border_alpha = sum(alpha.flatten(1)) - np.sum(np.sum(alpha[1:-1,1:-1]))

    print "border_size", border_size
    print "border_alpha", border_alpha

    flipped = (border_alpha/border_size) > 0.5

    if flipped: falpha = 1 - falpha
    return falpha

nbest = 10
mi = np.argsort(scr)
mv = scr[mi]

# mi does not match matlab -- there are small differences
# I do not think it's a bug, but not sure. In my test dataset:
# mi[5] and mi[6] are identical to 4 decimal places, so they could be switched
# However, mi[7-9] should be in the correct order, unless they are just unstable
# e.g., mi (python): mi [ 13  14 497 213 495  69 197  78 209 206]
#       mi (matlab):      14  15 498 214 496 198  70 215 207 210

print "mi", mi[:10]
print "mv", mv[:10]

best_x = np.zeros((N,nbest))
for l in range(nbest):
    subg = mi[l]
    bits = np.array([i for i in reversed(np.binary_repr(subg+1, width=free_clust_len))], np.float128)
    print "alpha_comps", alpha_comps
    # FIXFIX make sure this works with non-square images (n!=m)
    # because numpy is row-first and matlab is column-first, use m,n and transpose
    alpha = np.reshape(np.dot(alpha_comps, bits), (m,n)).T
    print "alpha", alpha[:10,:10], alpha[-10:,-10:]

    alpha = decideFB(alpha)

    # again, column-first flattening
    best_x[:,l] = alpha.flatten(1)
    print best_x[:,l][:10], best_x[:,l][-10:]

if (save_partial_results):
    # in matlab, plot everything as one image
    #imwrite(flat3DArray(reshape(best_x,n,m,nbest),2),[bname,'alpha_all_cands.tif']);
    for l in range(len(best_x[0])):
        print "saving l", l
        compname = "%s_alpha_all_cands_%d.png" % (bname,l)
        scipy.misc.imsave(compname, np.reshape(best_x[:,l], (w,h)).T)
print "done"

# now work on the best one...
best_idx = 0
subg=mi[best_idx]
bits = np.array([i for i in reversed(np.binary_repr(subg+1, width=free_clust_len))], np.float128)
alpha = np.reshape(np.dot(alpha_comps,bits),(h,w)).T
alpha = decideFB(alpha)

if (save_partial_results):
    scipy.misc.imsave("%s_best_alpha.png" % bname, alpha.T.astype(np.float128))

apply_final_enhancement = True

def finalEnhancment(alpha,eigs_num,eig_vectors,eig_values):
    # do some iterations of the Newton's method, to get a sparser alpha matte
    su = eig_vectors
    sd = eig_values
    scld = 1
    exp_a=0.9
    thr_e=1e-10
    w1=0.3
    w0=0.3
    wa=1
    orig_exp_a = exp_a

    x = alpha.flatten(1)
    for i in range(30):
      e1 = (w1**exp_a)*np.maximum(np.abs(x-1),thr_e)**(exp_a-2)
      e0 = (w0**exp_a)*np.maximum(np.abs(x),thr_e)**(exp_a-2)
      print e1.shape, e0.shape, (e1+e0).shape, eigs_num
      ssu = np.tile(e1+e0, (eigs_num,1)).T * su

      # y=(su'*ssu+scld*sd)\(su'*e1);
      print type((np.dot(su.T,ssu) + scld*sd)[0,0])
      print type(np.dot(su.T,e1)[0])
      y = np.linalg.solve((np.dot(su.T,ssu) + scld*sd).astype(np.float),
                          np.dot(su.T,e1).astype(np.float))
      x = np.dot(su,y)
      print x

      if i==20:
        exp_a = exp_a-0.2;

    n,m = alpha.shape

    # FIXFIX make sure this works with non-square matrix
    alpha = np.reshape(x,(h,w)).T

    exp_a = orig_exp_a

    return alpha

if apply_final_enhancement:
    alpha = finalEnhancment(alpha, eigs_num, eig_vectors, eig_values)

if (save_partial_results):
    rotated = ((alpha.T).reshape(w,h).T).astype(np.float128)
    # If I do not truncate at [0,1] then imsave doesn't work properly (white is light grey)
    rotated[rotated<0] = 0
    rotated[rotated>1] = 1
    scipy.misc.imsave("%s_alpha.png" % bname, rotated)

np.savetxt("alpha_vals.py.txt", alpha)