
"""

Automatic sketch maker.

I am sorry, normally Python is readable. This is program, not so much.

"""

import sys, Image, numpy, os, random
from numpy import fft


def load(filename):
    i = Image.open(filename)
    a = numpy.asarray(i).astype('float64')
    return a

def pad(a):
    pad_y = 1
    while pad_y < a.shape[0]*1.5: pad_y *= 2

    pad_x = 1
    while pad_x < a.shape[1]*1.5: pad_x *= 2
    
    a2 = numpy.zeros((pad_y,pad_x,3),'float64')
    in_y = (pad_y-a.shape[0])//2
    in_x = (pad_x-a.shape[1])//2
    a2[in_y:a.shape[0]+in_y,in_x:a.shape[1]+in_x] = a
    return a2, a.shape[0], a.shape[1]

def save(a, filename):
    a = numpy.clip(a,0,255)
    i = Image.fromarray(a.astype('uint8'))    
    i.save(filename)


if len(sys.argv) != 2:
    print 'Usage:'
    print
    print '    python %s image_file.jpg/png/etc' % sys.argv[0]
    print
    print 'Will write result to "output.jpg".'
    sys.exit(1)

a = load(sys.argv[1])
a, orig_height, orig_width = pad(a)

y,x = numpy.indices((a.shape[0],a.shape[1]))

def make_op(power):
    op = (y*y+x*x+1.0) ** -power
    numpy.maximum(op[:,1:], op[:,:-1][:,::-1], op[:,1:])
    numpy.maximum(op[1:,:], op[:-1,:][::-1,:], op[1:,:])
    op /= numpy.sum(op)
    return fft.fft2(op)

def blur(img, fop): 
    return fft.ifft2( fft.fft2(img) * fop ).real

op1 = make_op(1.25)
op2 = make_op(1.5)

gray = numpy.sum(a,2)

gray1 = blur(gray, op1)
gray2 = blur(gray, op2)

ao = numpy.empty(a.shape, 'float64')
ao[:] = 120

def derivative_y(a):
    result = numpy.zeros(a.shape)
    result[1:-1] = a[2:]*0.5
    result[1:-1] -= a[:-2] * 0.5
    return result

def derivative_x(a):
    return derivative_y(a.transpose()).transpose()

dy = derivative_y(gray2)
dx = derivative_x(gray2)
dydy = derivative_y(dy)
dxdx = derivative_x(dx)
dydx = derivative_y(dx)

det = dydy*dxdx - dydx*dydx 
trace = dydy+dxdx

i1 = 0.5*(trace + numpy.sqrt(trace*trace - 4*det))
i2 = 0.5*(trace - numpy.sqrt(trace*trace - 4*det))

def energy(a, op):
    return blur(a*a, op) ** 0.5

minish = energy(i1, op2)*2.0
ao[:,:,0:3] -= (  numpy.minimum(i2 + minish,0.0) )[:,:,None] * 20.0
minish = energy(i2, op2)*2.0
ao[:,:,0:3] -= (  numpy.maximum(i1 - minish,0.0) )[:,:,None] * 20.0

ao += ( (gray2 > gray1)*50.0 )[:,:,None] * [[[1.0,0.5,0.0]]]


pad = min(orig_height,orig_width)*0.2
pad_height = orig_height+pad
pad_width = orig_width+pad
in_y = (ao.shape[0]-pad_height)//2
in_x = (ao.shape[1]-pad_width)//2
save(ao[in_y:pad_height+in_y,in_x:pad_width+in_x],'output.jpg')


