from interpolation import pv
from numpy import array, matrix, zeros, dot
from numpy.linalg.linalg import lstsq, inv, svd

# url: http://www.leptonica.com/affine.html

def affine_transform(image, p1, p2, p3, width, height):
    """Warps the parallelogram defined by the three points (x1,y1), (x2,y2) and
    (x3,y3) onto a new image of size MxN."""

    x1, y1 = p1
    x2, y2 = p2
    x3, y3 = p3

    M = array([[x1, y1, 1, 0, 0, 0],
               [0, 0, 0, x1, y1, 1],
               [x2, y2, 1, 0, 0, 0],
               [0, 0, 0, x2, y2, 1],
               [x3, y3, 1, 0, 0, 0],
               [0, 0, 0, x3, y3, 1]])

    q = array([[0],
               [0],
               [width],
               [0],
               [width],
               [height]])

    # Calculate transform matrix
    a, b, c, d, e, f = lstsq(M, q)[0].T.tolist()[0]
    A = matrix([[a, b, c],
                [d, e, f],
                [0, 0, 1]]).I

    # Construct the transformed image
    result = zeros((width, height))

    for x in xrange(width):
        for y in xrange(height):
            orig_pos = A * array([[x, y, 1]]).T
            result[x][y] = pv(image, orig_pos[0], orig_pos[1], 'linear')

    return result

def perspective_transform(image, p1, p2, p3, p4, width, height):
    """Warps the parallelogram defined by four corner points onto a new image
    of size MxN."""

    x1, y1 = p1
    x2, y2 = p2
    x3, y3 = p3
    x4, y4 = p4

    M = array([[x1, y1, 1, 0, 0, 0,            0,            0,       0],
               [0, 0, 0, x1, y1, 1,            0,            0,       0],
               [x2, y2, 1, 0, 0, 0,  -width * x2,  -width * y2,  -width],
               [0, 0, 0, x2, y2, 1,            0,            0,       0],
               [x3, y3, 1, 0, 0, 0,            0,            0,       0],
               [0, 0, 0, x3, y3, 1, -height * x3, -height * y3, -height],
               [x4, y4, 1, 0, 0, 0,  -width * x4,  -width * y4,  -width],
               [0, 0, 0, x4, y4, 1, -height * x4, -height * y4, -height]])

    # Calculate transform matrix
    A = inv(svd(M)[2][-1].reshape(3, 3))

    # Construct the transformed image
    result = zeros((width, height, 4))

    for x in xrange(width):
        for y in xrange(height):
            p = dot(A, array([[x], [y], [1]]))
            result[x][y] = pv(image, p[1] / p[2], p[0] / p[2], 'linear')

    return result

if __name__ == '__main__':
    from pylab import imread, imshow, plot, subplot, show, axis

    # Affine transformation
    subplot(221)
    image = imread('cameraman.png')
    half_width = image.shape[0] / 2
    half_height = image.shape[1] / 2
    x = [0, half_width, image.shape[0], half_width, 0]
    y = [half_height, 0, half_height, image.shape[1], half_height]
    plot(x, y, 'r-')
    axis([0, image.shape[0], 0, image.shape[1]])
    imshow(image, cmap='gray')
    subplot(222)
    result = affine_transform(image, (0., half_height), (half_width, 0.),
                              (image.shape[0], half_height), 128, 128)
    imshow(result, cmap='gray')

    # Perspective transformation
    subplot(223)
    image = imread('flyeronground.png')
    x, y = zip((147, 588), (100, 370), (300, 205), (392, 392), (147, 588))
    plot(x, y, 'r-')
    imshow(image)
    subplot(224)
    result = perspective_transform(image, (147, 588), (392, 392), (100, 370), \
                                   (300, 205), 100, 150)
    imshow(result)
    show()
