OF

from math import sqrt

import numpy

def I(pixel, opt=0):
    R, G, B = pixel
    return (0.2126 * R) + (0.7152 * G) + (0.0722 * B)


def image_to_intensity(image, opt=0):
    pix = image.load()
    return [[I(pix[y, x]) for y in range(image.size[0])] for x in range(image.size[1])]


def lucas_kanade(im1, im2, win=2):
    width = len(im1[0])
    height = len(im1)

    assert len(im1[0]) == len(im2[0])
    assert len(im1) == len(im2)

    opfl = [[(0.0, 0.0)] * width for x in range(height)]

    for y in range(win + 1, height - win - 1):
        for x in range(win + 1, width - win - 1):
            neighbours = [(Y, X) for X in range(x - win, x + win + 1) for Y in range(y - win, y + win + 1)]

            A = numpy.matrix([[
                (im1[p[0]][p[1] + 1] - im1[p[0]][p[1] - 1]) / 2,
                (im1[p[0] + 1][p[1]] - im1[p[0] - 1][p[1]]) / 2
                ] for p in neighbours
            ])

            b = numpy.matrix([[
                im1[p[0]][p[1]] - im2[p[0]][p[1]]
                ] for p in neighbours
            ])

            try:
                v = numpy.linalg.solve(numpy.transpose(A) * A, numpy.transpose(A) * b)
            except(Exception):
                v = [[0], [0]]
            opfl[y][x] = (v[0][0], v[1][0])

    return opfl


import gc
import Image
from OpticalFlows import image_to_intensity
from OpticalFlows import lucas_kanade
from sys import argv as arg


for i in range(int(arg[1]), int(arg[2])):
    n1 = 'frame%04d.jpg' % i
    n2 = 'frame%04d.jpg' % (i + 1)

    print 'Computing Optical Flow between {0} and {1}'.format(n1, n2)

    im1 = Image.open(n1)
    im2 = Image.open(n2)

    i1 = image_to_intensity(im1)
    i2 = image_to_intensity(im2)

    win = 2

    opfl = lucas_kanade(i1, i2, win)

    # check GC, suprisingly values returned after each iteration are exactly the same, or less than 1% different
    print len(gc.get_objects())
    gc.collect()
    print len(gc.get_objects())

    # prepare iblack-red image, presenting the optical flow and blend it with original one
    normal = 255.0 / ((2*win+1)*2**(0.5))

    width, height = im1.size

    out = Image.new("RGB", (width,height))
    pix = out.load()
    for y in range(height):
      for x in range(width):
        pix[x,y] = (min(255,int(normal*sqrt(abs(opfl[y][x][0])*abs(opfl[y][x][0]) + abs(opfl[y][x][1])*abs(opfl[y][x][1])))),0,0)

    out.save('a'+n2)
    i3 = Image.blend(im2, out, 0.5)
    i3 = i3.point(lambda i: i * 2)
    i3.save('b'+n2)