Spatio temporal ISM

stism.py

from math import sqrt, log, pi
from random import random
import numpy as np

def build_st_ism(frames, r, s):
	nframes = len(frames)
	M, N = frames[0].shape
	patches = map(ism, get_patches(frames, r, s))
	spatial_ism = stitch_space(patches, nframes, M, N, r, s)
	slices = map(ism, get_slices(frames))
	temporal_ism = stitch_time(slices, nframes, M, N)
	return spatial_ism + temporal_ism

def get_patches(frames, r=10, s=10):
	nframes = len(frames)
	M, N = frames[0].shape
	npatches = (M*N)/(r*s)
	patches = []

	for i in xrange(nframes):
		for j in xrange(npatches):
			start_x = (j // (N // s)) * r
			start_y = (j % (N // s)) * s
			patches.append(list(frames[i][start_x:(start_x+r), start_y:(start_y+s)].flat))
	return patches

def get_slices(frames):
	slices = []
	M, N = frames[0].shape
	nframes = len(frames)

	for i in xrange(M):
		for j in xrange(N):
			slices.append([frames[k][i, j] for k in xrange(nframes)])

	return slices
	
def ism(patch):
	patch = np.array(patch)
	mu = np.mean(patch)
	sigma = np.std(patch) or 1 # prevent division by 0
	term1 = -log(1 / sqrt(2*pi) / sigma)
	gaussian = term1 + 0.5 * ((patch - mu) ** 2 / (sigma**2)) #Vectorised operation, executes C 
	return (mu, gaussian)

def stitch_space(patches, nframes, M, N, r=10, s=10):
	data = [patches[i][0] for i in xrange(len(patches))]
	mu = np.mean(data)
	sigma = np.std(data)

	npatches = (M*N) / (r*s)

	result = np.zeros((nframes, M, N))

	for i, (weight, patch) in enumerate(patches):
		frame_no = i / npatches
		patch_no = i % npatches
		info = gaussian_info(weight, mu, sigma)
		start_x = (patch_no // (N // s)) * r
		start_y = (patch_no % (N // s)) * s

		for j, val in enumerate(patch):
			a = j // s
			b = j % s
			result[frame_no, start_x + a, start_y + b] = info + val

	return result


def stitch_time(slices, nframes, M, N):
	data = [slices[i][0] for i in xrange(len(slices))]
	mu = np.mean(data)
	sigma = np.std(data)
	result = np.zeros((nframes, M, N))
	for i, (weight, slice) in enumerate(slices):
		info = gaussian_info(weight, mu, sigma)
		x = i // N
		y = i % N
		for j, val in enumerate(slice):
			result[j, x, y] = info + val

	return result

def gaussian_info(val, mu, sigma):
	#prefer vectorised numpy version over this
	return .5 * (val - mu) ** 2 / (sigma**2)

def test():
	video = [np.array([[4, 4, 4, 4],
					   [4, 4, 1, 4],
					   [4, 4, 4, 4],
					   [4, 4, 4, 4]]), 
		np.array([[4, 4, 4, 34],
					   [4, 4, 10, 23],
					   [4, 4, 4, 4],
					   [4, 4, 4, 4]])]
	print build_st_ism(video, 2, 2)

if __name__ == '__main__':
	test()