Iteration 3 of NNMF project for blending pixel art

icon_nnmf.py

from scipy.misc import imread
import numpy as np

# These are flower images from the PROJCAM garden bundle.
path = "C:\Users\Mark\Documents\ProcJam\PROCJAM2016-Tess2D\PROCJAM2016-Tess2D\Garden\\"

names = [ "flower{:02d}.png".format( i ) for i in xrange( 0, 22 ) ]
original = [ imread( path + name ) for name in names ]

def flatten( image ):
    return np.reshape( image, [-1] )

palette = set()
for image in original:
    for row in image:
        for color in row:
            palette.add( tuple( color ) )

transparent = False
for (r,g,b,a) in list( palette ):
    if a == 0: # transparent
        transparent = True
        palette.remove( (r,g,b,a) )

palette = list( palette )

if transparent:
    # Make sure this is color 0
    palette = [ (255, 255, 255, 0) ] + palette
    
print "Palette", palette

def convertToPalette( image ):
    px = []
    for pixel in np.reshape( image, [-1, 4] ):
        (r, g, b, a ) = pixel
        if a == 0:
            pixel = ( 255, 255, 255, 0 )
        vec = [ int( tuple(pixel) == palette[i] ) for i in xrange( len( palette ) ) ]
        px.append( vec )
    newRep = np.array( px )
    return np.reshape( newRep, [-1] )

from math import sqrt

def convertToRgb_max( image ):
    y = len( palette )
    pixels = np.reshape( image,[-1, y] )
    maxPixel = np.argmax( pixels, 1 )                           

    rgbVec = np.array( [ palette[i] for i in maxPixel ] )    
    dim = int( sqrt( len( rgbVec ) ) )
    rgbMat = np.reshape( rgbVec, [dim, dim, 4] )
    return np.array( rgbMat, dtype="uint8" )

def convertToRgb_alphablend( image ):
    y = len( palette )
    pixels = np.reshape( image, [-1, y] )
    pa = np.array( palette )

    # The RGB component of the transparent color should
    # not be included in the mix.
    rgbVec = []
    transparent = 0     # palette.index( (255, 255, 255, 0) )
    for p in pixels:
        row_sum = sum( p[1:] )
        alpha = p[0]

        if row_sum == 0:
            color = palette[0]
        else:
            color = ( p[1:] / row_sum ).dot( pa[1:] )
            color[3] = (1.0-alpha) * 255
        #print p, color
        
        rgbVec.append( color )
                           
    rgbVec = np.array( rgbVec )
    dim = int( sqrt( len( rgbVec ) ) )
    rgbMat = np.reshape( rgbVec, [dim, dim, 4] )
    return np.array( rgbMat, dtype="uint8" )

from scipy.spatial import distance

def convertToRgb_neighbor( image ):
    y = len( palette )
    pixels = np.reshape( image, [-1, y] )
    pa = np.array( palette[1:] )
    
    rgbVec = []
    for p in pixels:
        # If more than 50 percent background weight, make transparent
        if p[0] >= 0.5:
            color = palette[0]
        else:
            # Otherwise average the colors and pick the closest neighbor
            # from the palette
            row_sum = sum( p[1:] )
            blendcolor = ( p[1:] / row_sum ).dot( pa )
            color = pa[distance.cdist( [blendcolor], pa ).argmin() ]            
            #print p, blendcolor, color
            
        #print p, color
        
        rgbVec.append( color )
                           
    rgbVec = np.array( rgbVec )
    dim = int( sqrt( len( rgbVec ) ) )
    rgbMat = np.reshape( rgbVec, [dim, dim, 4] )
    return np.array( rgbMat, dtype="uint8" )

def convertToRgb_blend( image ):
    y = len( palette )
    pixels = np.reshape( image, [-1, y] )
    pa = np.array( palette )
    
    # Normalize to 1.0
    row_sums = pixels.sum( axis = 1, keepdims = True )
    normalized = pixels / row_sums
                           
    rgbVec = normalized.dot( pa )
    dim = int( sqrt( len( rgbVec ) ) )
    rgbMat = np.reshape( rgbVec, [dim, dim, 4] )
    return np.array( rgbMat, dtype="uint8" )

def randDist( d ):
    return np.random.choice( range( len( d ) ),
                             p = d )
    
def convertToRgb_rand( image ):
    y = len( palette )
    pixels = np.reshape( image,[-1, y] )

    # Normalize to 1.0
    row_sums = pixels.sum( axis = 1, keepdims = True )
    normalized = pixels / row_sums
    
    rgbVec = np.array( [ palette[randDist(i)] for i in normalized ] )    
    dim = int( sqrt( len( rgbVec ) ) )
    rgbMat = np.reshape( rgbVec, [dim, dim, 4] )
    return np.array( rgbMat, dtype="uint8" )

# Matrix V has shape [samples, feature]
# Each sample is an image, each feature is the presence of a palette color
# in a pixel
V = np.matrix( [ convertToPalette( i ) for i in original ] )
(numSamples, numFeatures) = V.shape
print "Samples: ", numSamples
print "Features: ", numFeatures

from sklearn.decomposition import NMF
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

def trial( numComponents = 22, max_iter = 10000 ):
    model = NMF( n_components= numComponents, max_iter=max_iter )
    W = model.fit_transform( V )
    H = model.components_
    # Matrix W has shape [ samples, components ]
    # Matrix H has shape [ components, features ]

    print "Error", model.reconstruction_err_

    X2 = model.inverse_transform( W )

    plt.figure( 1 )
    numImages = len( names )
    for n in xrange( numImages ):
        # original
        o = original[n]
        plt.subplot( numImages, 2, n * 2 + 1 )
        plt.imshow( o, interpolation = "none" )
        plt.axis( "off" )
    
        # reconstruction
        p = convertToRgb( X2[n] )
        plt.subplot( numImages, 2, n * 2 + 2 )
        plt.imshow( p, interpolation = "none" )
        plt.axis( "off" )    

    plt.show()

#    return ( model, W, H )

def mixAll( numComponents = 22, max_iter = 10000 ):
    model = NMF( n_components= numComponents, max_iter=max_iter )
    W = model.fit_transform( V )
    H = model.components_
    # Matrix W has shape [ samples, components ]
    # Matrix H has shape [ components, features ]

    fig = plt.figure( 1 )
    numImages = len( names )
    for a in xrange( numImages ):
        for b in xrange( numImages ):
            plt.subplot( numImages, numImages, numImages * a + b + 1 )
            if a == b:
                plt.imshow( original[a], interpolation = "none" )
            elif a < b:
                wa = W[a]
                wb = W[b]
                wi = ( wa + wb ) * 0.5
                x = wi.dot( H )
                plt.imshow( convertToRgb_alphablend( x ), interpolation = "none" )
            plt.axis( "off" )

    fig.text( 0.1, 0.1, "50% blend with " + str( numComponents ) + " components\nalpha-blend conversion" )                     
    plt.show()

def plot_interpolate( a, b, t, W, H, func, ax ): 
    ax.set_title( str( int( t * 100 ) ) + "%" )
    wa = W[a]
    wb = W[b]
    wi = wa * (1.0 - t) + wb * t
    x = wi.dot( H )
    ax.imshow( func( x ), interpolation = "none" )
    ax.axis( "off" )

def row_interpolate( a, b, W, H, func, axs ):
    axs[0].imshow( original[a], interpolation = "none" )
    axs[0].axis( "off" )

    axs[-1].imshow( original[b], interpolation = "none" )
    axs[-1].axis( "off" )

    numIntermediate = len( axs ) - 2
    for i in xrange( numIntermediate ):
        ax = axs[i + 1]
        plot_interpolate( a, b,
                         ( i + 1 ) * 1.0 / ( numIntermediate + 1 ),
                         W, H, func, ax )
    
def interpolate( a, b, numComponents = 22, max_iter = 10000, func = convertToRgb_alphablend ):
    model = NMF( n_components= numComponents, max_iter=max_iter )
    W = model.fit_transform( V )
    H = model.components_

    f, axs = plt.subplots( ncols = 11 )
    row_interpolate( a, b, W, H, func, axs )
    plt.show()

def compareInterpolate( a, b, numComponents = 22, max_iter = 10000 ):
    funcs = [
        ( "max", convertToRgb_max ),
        ( "blend", convertToRgb_blend ),
        ( "alphablend", convertToRgb_alphablend ),
        ( "rand", convertToRgb_rand ),
        ( "neighbor", convertToRgb_neighbor )
        ]

    model = NMF( n_components= numComponents, max_iter=max_iter )
    W = model.fit_transform( V )
    H = model.components_

    f, axs = plt.subplots( nrows = len( funcs ),
                           ncols = 11 )

    f.suptitle( "Comparing RGB conversions, " + str( numComponents ) + " components" )

    for i in xrange( len( funcs ) ):
        row_interpolate( a, b, W, H, funcs[i][1], axs[i] )
        axs[i][0].set_title( funcs[i][0], loc="left" )
        
    plt.show()