# Anisotropic Diffusion Filter
# Written by Robert A. Brown
# Calgary, Alberta, Canada
# 2002

#	This code may be used freely for non-commercial purposes, but please
# leave the paragraph above intact and attached and mention my name somewhere.  
# This code is intended as an example only and I don't offer any guarantees -- 
# it might blow up your computer, kill your patients or cause any number of other 
# misfortunes.  If you do find a use, a bug, or want to make a comment,
# please let me know at software@robbtech.com.


from Numeric import *

def shift(a,ny,nx):
	b = concatenate((a[:,ny:,:],a[:,0:ny,:]),1)
	c = concatenate((b[:,:,nx:],b[:,:,0:nx]),2)
	return c	


def flow(grad,k):
	magimage = sum(grad*grad,0)
	magimage = magimage / (k*k)
	small = 0.00001
	flowimage = exp(-magimage)
	w = (abs(flowimage) > small)
	flowimage = flowimage * w
	return flowimage.astype(Float32)


def anisoFilter(images,k,iterations=1):
	""" Images is a list of 2D Numpy arrays with greyscale pixel data.
			Multiple images are treated as independent channels, and the information
			is used to perform better filtering.
			k is the diffusion parameter.  Higher k filters more aggressively, but will
			not respect weaker edges.
			iterations is the number of times to run the filter.  3 is a fairly standard
			value."""
	images = array(images.astype(Float32))
	if (iterations <= 0): iterations = 1
	for j in range(0,iterations):
		dN = images - shift(images,1,0)
		dS = images - shift(images,-1,0)
		dE = images - shift(images,0,1)
		dW = images - shift(images,0,-1)

		cN = flow(dN,k)
		cS = flow(dS,k)
		cE = flow(dE,k)
		cW = flow(dW,k)
		
		for i in range(0,shape(images)[0]):
			dN[i] = dN[i]*cN
			dS[i] = dS[i]*cS
			dE[i] = dE[i]*cE
			dW[i] = dW[i]*cW
		
		images = images - 0.2*(dN+dS+dW+dE)	
	
	return images