seam_dofs.py

#!/usr/bin/python

from __future__ import print_function, division

from math import *
from numpy import *
import numpy.linalg

try:
    from PIL import Image
except ImportError:
    import Image

## Print super-long lines instead of hard wrapping.
numpy.set_printoptions( linewidth = 200 )
## Suppress scientific notation and only print 3 decimal places.
# numpy.set_printoptions( suppress = True, precision = 3 )

def flatindex( i,j, cols ):
    return i*cols + j

def abcdFromXYandSTs( xy, flatindex, xys = None ):
    '''
    Given:
        `xy`: A 2D point (where each bilinear square encloses integer offsets)
        `flatindex`: A function taking a 2D pixel index and returning a 1D index
        (optional) `xys`: A sequence of 2D points to return in local s,t coordinates for the square
    Returns:
        The flat 1D integer indices for the a,b,c,d corners of the enclosing bilinear square
        (optional) a sequence of st coordinates for `xys`
    '''
    
    x,y = xy
    
    abcd = [
        ( int(floor( x )), int(floor( y )) ),
        ( int(ceil ( x )), int(floor( y )) ),
        ( int(floor( x )), int(ceil ( y )) ),
        ( int(ceil ( x )), int(ceil ( y )) )
        ]
    
    result = list( map( flatindex, abcd ) )
    
    if xys is None:
        return result
    else:
        return result, asarray( xys ) - array(abcd[0])[newaxis,:]

def intersecting_parameters( p0, p1 ):
    '''
    Given:
        Start and end XY endpoints `p0` and `p1` for a line segment
    Returns:
        Linear interpolation parameter values from `p0` to `p1` where the line segment intersects the boundaries of bilinear interpolation squares (e.g. integer grid lines).
    '''
    
    assert len(p0) == 2
    assert len(p1) == 2
    
    ts = []
    
    for c in (0,1):
        start = int( ceil( min( p0[c], p1[c] ) ) )
        end = int( floor( max( p0[c], p1[c] ) ) )
        
        assert start >= 0
        assert end >= 0
        
        for edge in range( start, end+1 ):
            t = ( edge - p0[c] )/( p1[c] - p0[c] )
            
            ts.append( t )
    
    return ts

def bilinear_intervals( p0, p1, q0, q1 ):
    '''
    Given:
        Start and end XY endpoints `p0` and `p1` for a line segment and `q0` and `q1` for a copy of the line segment
    Returns:
        Parameters `t` linearly along `(p0,p1)` or `(q0,q1)` where either line segment changes bilinear square.
    '''
    
    ts = [0,1] + list( intersecting_parameters( p0, p1 ) ) + list( intersecting_parameters( q0, q1 ) )
    print( "ts:", sorted( ts ) )
    assert len( ts ) == len( set( ts ) )
    ts.sort()
    return ts

def lerp( start, end, t ):
    '''
    Returns the linear interpolation of `start` and `end` according to parameter `t`.
    '''
    start = asarray(start)
    end = asarray(end)
    
    return start + t*( end - start )

def bilerp( abcd, st ):
    '''
    Given:
        `abcd`: A sequence of four values at the four corners of a bilinear square (s=0,t=0; s=1,t=0; s=0,t=1; s=1,t=1)
        `st`: A pair of values in [0,1]^2 within the square
    Returns:
        The bilinear interpolation of the values.
    '''
    
    abcd = asarray( abcd )
    a,b,c,d = abcd
    
    s,t = st
    
    return ( a + s*(b-a) )*(1-t) + ( c + s*(d-c) )*t

def constraints_for_seam_pair( p0, p1, q0, q1, rows, cols ):
    '''
    Given:
        Start and end XY endpoints `p0` and `p1` for a line segment and `q0` and `q1` for a copy of the line segment
        The dimensions of the pixel grid `rows` by `cols`
    Returns:
        A sequence of linear constraints so that the bilinear reconstruction along `(p0,p1)` exactly matches the one along `(q0,q1)`.
    '''
    
    ts = bilinear_intervals( p0, p1, q0, q1 )
    assert len( ts ) >= 2
    
    def fi( ij ):
        return flatindex( ij[0], ij[1], cols )
    
    constraints = zeros(( 3*(len(ts)-1), rows*cols ))
    
    for i in range( len( ts ) - 1 ):
        t0 = ts[i]
        t1 = ts[i+1]
        
        ## Add constraint for t0 to t1
        ### Get the bilinear square for p. Use the midpoint of the interval, because the endpoints are borderline.
        (ap,bp,cp,dp), ((p_s0,p_t0),(p_s1,p_t1)) = abcdFromXYandSTs( lerp( p0, p1, .5*(t0+t1) ), fi, [ lerp( p0, p1, t0 ), lerp( p0, p1, t1 ) ] )
        (aq,bq,cq,dq), ((q_s0,q_t0),(q_s1,q_t1)) = abcdFromXYandSTs( lerp( q0, q1, .5*(t0+t1) ), fi, [ lerp( q0, q1, t0 ), lerp( q0, q1, t1 ) ] )
        
        ### Constant term.
        #### a + s0*(b-a) + t0*(c-a) + s0*t0*(d-c+a-b)
        constraints[ 3*i + 0, ap ] += 1. - p_s0 - p_t0 + p_s0*p_t0
        constraints[ 3*i + 0, aq ] -= 1. - q_s0 - q_t0 + q_s0*q_t0
        constraints[ 3*i + 0, bp ] += p_s0 - p_s0*p_t0
        constraints[ 3*i + 0, bq ] -= q_s0 - q_s0*q_t0
        constraints[ 3*i + 0, cp ] += p_t0 - p_s0*p_t0
        constraints[ 3*i + 0, cq ] -= q_t0 - q_s0*q_t0
        constraints[ 3*i + 0, dp ] += p_s0*p_t0
        constraints[ 3*i + 0, dq ] -= q_s0*q_t0
        
        ### Linear term.
        #### (t1-t0)*(c-a) + (s1-s0)*(b-a) + ( t0*(s1-s0) + s0*(t1-t0) )*( d-c+a-b )
        constraints[ 3*i + 1, ap ] += -(p_t1-p_t0) - (p_s1-p_s0) + ( p_t0*(p_s1-p_s0) + p_s0*(p_t1-p_t0) )
        constraints[ 3*i + 1, aq ] -= -(q_t1-q_t0) - (q_s1-q_s0) + ( q_t0*(q_s1-q_s0) + q_s0*(q_t1-q_t0) )
        constraints[ 3*i + 1, bp ] += (p_s1-p_s0) - ( p_t0*(p_s1-p_s0) + p_s0*(p_t1-p_t0) )
        constraints[ 3*i + 1, bq ] -= (q_s1-q_s0) - ( q_t0*(q_s1-q_s0) + q_s0*(q_t1-q_t0) )
        constraints[ 3*i + 1, cp ] += (p_t1-p_t0) - ( p_t0*(p_s1-p_s0) + p_s0*(p_t1-p_t0) )
        constraints[ 3*i + 1, cq ] -= (q_t1-q_t0) - ( q_t0*(q_s1-q_s0) + q_s0*(q_t1-q_t0) )
        constraints[ 3*i + 1, dp ] += ( p_t0*(p_s1-p_s0) + p_s0*(p_t1-p_t0) )
        constraints[ 3*i + 1, dq ] -= ( q_t0*(q_s1-q_s0) + q_s0*(q_t1-q_t0) )
        
        ### Quadratic term.
        #### (s1-s0)*(t1-t0)*(d-c+a-b)
        constraints[ 3*i + 2, ap ] += (p_s1-p_s0)*(p_t1-p_t0)
        constraints[ 3*i + 2, aq ] -= (q_s1-q_s0)*(q_t1-q_t0)
        constraints[ 3*i + 2, bp ] += -(p_s1-p_s0)*(p_t1-p_t0)
        constraints[ 3*i + 2, bq ] -= -(q_s1-q_s0)*(q_t1-q_t0)
        constraints[ 3*i + 2, cp ] += -(p_s1-p_s0)*(p_t1-p_t0)
        constraints[ 3*i + 2, cq ] -= -(q_s1-q_s0)*(q_t1-q_t0)
        constraints[ 3*i + 2, dp ] += (p_s1-p_s0)*(p_t1-p_t0)
        constraints[ 3*i + 2, dq ] -= (q_s1-q_s0)*(q_t1-q_t0)
    
    return constraints

def constraints_linear( rows, cols ):
    '''
    Given:
        The dimensions of the pixel grid `rows` by `cols`
    Returns:
        A sequence of linear constraints so that the bilinear reconstruction is always exactly linear.
    '''
    constraints = zeros( ( (rows-1)*(cols-1), rows*cols ) )
    
    count = 0
    for r in range(rows-1):
        for c in range(cols-1):
            ## (d-c+a-b) = 0
            a,b,c,d = abcdFromXYandSTs( ( c + 0.5, r + 0.5 ), lambda ij: flatindex( ij[0], ij[1], cols ) )
            constraints[ count, a ] += 1
            constraints[ count, b ] += -1
            constraints[ count, c ] += -1
            constraints[ count, d ] += 1
            count += 1
    
    return constraints

def build_system( p0s, p1s, q0s, q1s, with_linear = None ):
    '''
    Builds a constraint system
    
    
    '''
    
    '''
    Given:
        Sequences of start and end XY endpoints `p0s` and `p1s` for a sequence of line segments and `q0s` and `q1s` for where the copy of each line segment should be. To be specific, the line seqment from `p0s[i]` to `p1s[i]` should have a matching bilinearly reconstructed function with the line segment from `q0s[i]` to `q1s[i]`.
        The dimensions of the pixel grid `rows` by `cols`
        (optional) `with_linear`: A boolean specifying whether to add the constraint that each bilinear square should be a linear function.
    Returns:
        A system of homogeneous linear equations which ensures that the bilinear reconstruction along `(p0,p1)` exactly matches the one along `(q0,q1)` for each `p0,p1` and `q0,q1` in the sequences. If `with_linear` is true, some of the constraints also ensure that the bilinear reconstruction is linear in each square.
    '''
    
    if with_linear is None:
        with_linear = False
    
    p0s = asarray( p0s )
    p1s = asarray( p1s )
    q0s = asarray( q0s )
    q1s = asarray( q1s )
    
    ## All points should be 2D
    assert p0s.shape[1] == 2
    assert p0s.shape == q0s.shape
    assert p0s.shape == p1s.shape
    assert q0s.shape == q1s.shape
    
    cols  = int(ceil( concatenate( ( p0s, p1s, q0s, q1s ), axis = 0 )[:,0].max() ))+1
    rows = int(ceil( concatenate( ( p0s, p1s, q0s, q1s ), axis = 0 )[:,1].max() ))+1
    print( "cols:", cols )
    print( "rows:", rows )
    
    all_cs = []
    for p0, p1, q0, q1 in zip( p0s, p1s, q0s, q1s ):
        cs = constraints_for_seam_pair( p0, p1, q0, q1, rows, cols )
        assert cs.shape[0] % 3 == 0
        assert cs.shape[1] == rows*cols
        
        all_cs.append( cs )
    
    if with_linear:
        cs = constraints_linear( rows, cols )
        print( "Adding linear constraints (bottom %s)." % len( cs ) )
        all_cs.append( cs )
    
    assert len( all_cs ) > 0
    M = concatenate( all_cs, axis = 0 )
    assert len( M.shape ) == 2
    assert with_linear or M.shape[0] % 3 == 0
    assert M.shape[1] == rows*cols
    
    return M, rows, cols

def nnz( v, eps = 1e-9 ):
    '''
    Returns the number of non-zeros in a sequence of numbers.
    '''
    return sum( abs( asarray(v) ) > 1e-9 )

def analyze_M( M, rows, cols, randomly_sample = None ):
    
    print( "M.shape:", M.shape )
    
    # print( "M:" )
    # print( M )
    
    eps = 1e-10
    
    # print( "rank:", numpy.linalg.matrix_rank(M) )
    singvals = numpy.linalg.svd( M, compute_uv = False )
    print( "svd singular values:", singvals )
    nonzero_singvals = singvals > eps
    print( "svd singular values > %s:" % eps, nonzero_singvals )
    print( "number of singular values > %s:" % eps, sum( nonzero_singvals ) )
    
    print( "Does a constant (non-zero) function satisfy all constraints? If so, this will be all zeros:" )
    print( nnz( dot( M, ones( M.shape[1] ) ) ) )
    
    print( "Does a linear (non-constant) function satisfy all constraints? If so, this will be all zeros:" )
    v = zeros( M.shape[1] )
    for r in range(rows):
        for c in range(cols):
            i = flatindex( r,c, cols )
            v[i] = r+c
    print( nnz( dot( M, v ) ) )
    
    ## Let's modify v minimally to get a perfect solution.
    vp = v
    for it in range(10):
        vp = linalg.solve( 10**7*dot( M.T, M ) + identity(M.shape[1]), vp.copy() )
    print( "How many non-zeros after solving for a closest satisfying solution?" )
    print( nnz( dot( M, vp ) ) )
    
    if randomly_sample is not None:
        constant_diff = randomly_sample( ones( rows*cols ) )
        before_diff = randomly_sample( v )
        after_diff = randomly_sample( vp )
        print( "Average absolute difference along seams for a constant function (should be zero, tests bilerp()):", constant_diff )
        print( "Average absolute difference along seams for a function before optimization (should be non-zero):", before_diff )
        print( "Average absolute difference along seams for a function after optimization (should be zero):", after_diff )

def save_eigenvectors( M, rows, cols, prefix, uniformly_sample ):
    '''
    Saves the zero eigenvectors of `M.T @ M` as numbered rows-by-cols greyscale images
    with the file beginning with `prefix`.
    '''
    
    eps = 1e-10
    
    ## Let's plot eigenfunctions of the square matrix.
    eigenvals, eigenvecs = linalg.eig( dot( M.T, M ) )
    count = 0
    for eigenval, eigenvec in zip( eigenvals, eigenvecs ):
        if eigenval > eps: continue
        
        image = eigenvec.reshape( rows, cols )
        print( "Eigenval: ", eigenval )
        uniformly_sample( eigenvec ) 
        outpath = "%s zero eigenvector %s.png" % ( prefix, count )
        Image.fromarray( (image*255.).round().clip(0,255).astype(uint8) ).save( outpath )
        print( "Saved:", outpath )
        
        count += 1

def uniformly_sample( rows, cols, values, p0s, p1s, q0s, q1s, num_samples ):
    '''
    Given:
        `rows` and `cols`: The dimensions of a grid of values.
        `values`: A flat grid of values.
        Sequences of start and end XY endpoints `p0s` and `p1s` for a sequence of line segments and `q0s` and `q1s` for where the copy of each line segment should be. To be specific, the line seqment from `p0s[i]` to `p1s[i]` should have a matching bilinearly reconstructed function with the line segment from `q0s[i]` to `q1s[i]`.
    Returns:
        The average absolute difference of the bilinearly interpolated values at randomly sampled points along seam edges.
    '''
    
    p0s = asarray( p0s )
    p1s = asarray( p1s )
    q0s = asarray( q0s )
    q1s = asarray( q1s )
    
    ## All points should be 2D
    assert p0s.shape[1] == 2
    assert p0s.shape == q0s.shape
    assert p0s.shape == p1s.shape
    assert q0s.shape == q1s.shape
    
    def fi( ij ):
        return flatindex( ij[0], ij[1], cols )
    
    total = 0.
    for edge_index in range(len( p0s )):
        for t in linspace(0,1,100):
            p = lerp( p0s[edge_index], p1s[edge_index], t )
            q = lerp( q0s[edge_index], q1s[edge_index], t )
            p_abcd, [p_st] = abcdFromXYandSTs( p, fi, [p] )
            q_abcd, [q_st] = abcdFromXYandSTs( q, fi, [q] )
            p_val = bilerp( values[ p_abcd ], p_st )
            q_val = bilerp( values[ q_abcd ], q_st )
            print(p_val,q_val);
            total += abs( p_val - q_val )
    
    return total / num_samples

def randomly_sample( rows, cols, values, p0s, p1s, q0s, q1s, num_samples ):
    '''
    Given:
        `rows` and `cols`: The dimensions of a grid of values.
        `values`: A flat grid of values.
        Sequences of start and end XY endpoints `p0s` and `p1s` for a sequence of line segments and `q0s` and `q1s` for where the copy of each line segment should be. To be specific, the line seqment from `p0s[i]` to `p1s[i]` should have a matching bilinearly reconstructed function with the line segment from `q0s[i]` to `q1s[i]`.
    Returns:
        The average absolute difference of the bilinearly interpolated values at randomly sampled points along seam edges.
    '''
    
    p0s = asarray( p0s )
    p1s = asarray( p1s )
    q0s = asarray( q0s )
    q1s = asarray( q1s )
    
    ## All points should be 2D
    assert p0s.shape[1] == 2
    assert p0s.shape == q0s.shape
    assert p0s.shape == p1s.shape
    assert q0s.shape == q1s.shape
    
    def fi( ij ):
        return flatindex( ij[0], ij[1], cols )
    
    total = 0.
    for ri in range(num_samples):
        edge_index = random.randint( 0, len( p0s ) )
        t = random.uniform()
        p = lerp( p0s[edge_index], p1s[edge_index], t )
        q = lerp( q0s[edge_index], q1s[edge_index], t )
        p_abcd, [p_st] = abcdFromXYandSTs( p, fi, [p] )
        q_abcd, [q_st] = abcdFromXYandSTs( q, fi, [q] )
        p_val = bilerp( values[ p_abcd ], p_st )
        q_val = bilerp( values[ q_abcd ], q_st )
        total += abs( p_val - q_val )
    
    return total / num_samples

def test_two_edges( with_linear = None, prefix = None ):
    print( "=== test_two_edges( %s ) ===" % with_linear )
    
    ## Two edges in a 3x3 pixel grid
    p0s = [[ .3, .5 ]]
    p1s = [[ .5, 1.5 ]]
    q0s = [[ 1.4, .6 ]]
    q1s = [[ 1.5, 1.5 ]]
    
    M, rows, cols = build_system( p0s, p1s, q0s, q1s, with_linear = with_linear )
    analyze_M( M, rows, cols, lambda v: randomly_sample( rows, cols, v, p0s, p1s, q0s, q1s, 100 ) )
    if prefix is not None: save_eigenvectors( M, rows, cols, prefix, lambda v: uniformly_sample(rows, cols, v, p0s, p1s, q0s, q1s, 100 ) )

def test_four_edges1( with_linear = None, prefix = None ):
    print( "=== test_four_edges1( %s ) ===" % with_linear )
    
    ## Four edges in a 3x3 pixel grid
    p0s = [[ 0.9, 0.1 ], [ 0.2, 1.1 ]]
    p1s = [[ 0.2, 1.1 ], [ 1.9, 1.2 ]]
    
    q0s = [[ 1.9, 1.2 ], [ 0.5, 1.5 ]]
    q1s = [[ 0.5, 1.5 ], [ 0.9, 0.1 ]]
    
    M, rows, cols = build_system( p0s, p1s, q0s, q1s, with_linear = with_linear )
    analyze_M( M, rows, cols, lambda v: randomly_sample( rows, cols, v, p0s, p1s, q0s, q1s, 100 ) )
    if prefix is not None: save_eigenvectors( M, rows, cols, prefix , lambda v: uniformly_sample(rows, cols, v, p0s, p1s, q0s, q1s, 100 ) )

def test_four_edges2( with_linear = None, prefix = None ):
    print( "=== test_four_edges2( %s ) ===" % with_linear )
    
    ## Four edges in a 2+offX-by-2+offY pixel grid.
    ### If one of offX or offY stays less than 2,
    ### then the dimensions in one direction will be
    ### 4 texels. 4 texels means that every
    ### texel in the image will still be part
    ### of a bilinear square with a seam edge
    ### running through it (no degrees of freedom
    ### are uninvolved and hence have 0 singular
    ### values "by default").
    offX = 2
    offY = 5
    p0s = [[ 0.9, 0.1        ], [ 0.2, offY + 0.1 ]]
    p1s = [[ 0.2, offY + 0.1 ], [ offX + 0.9, offY + 0.2 ]]
    
    q0s = [[ offX + 0.9, offY + 0.2 ], [ 0.5, offY + 0.5 ]]
    q1s = [[ 0.5, offY + 0.5 ], [ 0.9, 0.1 ]]
    
    ## verify that we didn't mess anything up.
    VERIFY = False
    if VERIFY:
        print( "verifying..." )
        p0s0 = [[ 0.9, 0.1 ], [ 0.2, 1.1 ]]
        p1s0 = [[ 0.2, 1.1 ], [ 1.9, 1.2 ]]
        q0s0 = [[ 1.9, 1.2 ], [ 0.5, 1.5 ]]
        q1s0 = [[ 0.5, 1.5 ], [ 0.9, 0.1 ]]
        print( abs(asarray(p0s0) - p0s).sum() )
        print( abs(asarray(p1s0) - p1s).sum() )
        print( abs(asarray(q0s0) - q0s).sum() )
        print( abs(asarray(q1s0) - q1s).sum() )
    
    M, rows, cols = build_system( p0s, p1s, q0s, q1s, with_linear = with_linear )
    analyze_M( M, rows, cols, lambda v: randomly_sample( rows, cols, v, p0s, p1s, q0s, q1s, 100 ) )
    if prefix is not None: save_eigenvectors( M, rows, cols, prefix , lambda v: uniformly_sample(rows, cols, v, p0s, p1s, q0s, q1s, 100 ))

if __name__ == '__main__':
    test_two_edges( True )
    test_two_edges( False )
    
    test_four_edges1( True )
    test_four_edges1( False )
    
    ## This version stretches into more DOFs.
    test_four_edges2( True, "large linear" )
    test_four_edges2( False, "large bilinear" )