IJF_Tracking.py

import time
import java.awt.Color
import ij
import ij.measure
import sys
import math
import os.path
from org.apache.commons.math.analysis import MultivariateRealFunction
from jarray import *
from org.apache.commons.math.optimization import SimpleScalarValueChecker, GoalType
from org.apache.commons.math.optimization.direct import NelderMead
from org.apache.commons.math.optimization.fitting import CurveFitter

class Gaussian2D():
#	params[0]=center(x)
#	params[1]=center(y)
#	params[2]=sigma(x)
#	params[3]=sigma(y)
#	params[4]=amp
#	params[5]=yoffset
	__nx=None;
	__ny=None;
	__data=[];
	__params=None;
    
	def __init__(self, width, height, par):
		self.__nx=width
		self.__ny=height
		self.__params=par
		self.__data=[]

	def compute(self):
		c_x=float(self.__params[0])
		c_y=float(self.__params[1])
		s_x=float(self.__params[2])
		s_y=float(self.__params[3])
		amp=float(self.__params[4])
		y_off=float(self.__params[5])
		nx=self.__nx
		ny=self.__ny
		d=self.__data.append		#points to append method to save time on object lookup inside loop
		for j in range(self.__ny):
			for i in range(self.__nx):
				d(y_off+amp*math.exp(-((i-c_x)/s_x)**2 - ((j-c_y)/s_y)**2))
		return self.__data

	def plot2D(self):
		pix=array(self.__data, 'd')
		imp=ImagePlus("Gaussian2D", FloatProcessor(self.__nx, self.__ny, pix))
		imp.show()	    

	def getParams(self):
		return self.__params
    
	def getData(self):
		return self.__data

	def setParams(self, params):
		self.__params=params
	    
class Gaussian2DMinimizer(MultivariateRealFunction):
	# for fitting, capable of dealing with constrained optimization problems
	__data=0.0
	__nx=0
	__ny=0
	__lbounds=[]
	__ubounds=[]
	__isConstrained=False
    
	def __init__(self):
		pass
    
	def setImage(self, data, width, height, lb=[], ub=[], isConstrained=False):
		self.__data=data
		self.__nx=int(width)
		self.__ny=int(height)
		self.__lbounds=lb
		self.__ubounds=ub
		self.__isConstrained=isConstrained

	def setBounds(self, lb=[],ub=[], isConst=False) :
		self.__lbounds=lb
		self.__ubounds=ub
		self.__isConstrainedFit=isConst
	
	def origToUnconstrainedParams(self, params) :
		#convert real-params to unconstrained parameters for fitting
		#assuming ub and lb are both specified
		lb=self.__lbounds
		ub=self.__ubounds
		out=[]
		for i in range(len(params)):
			if params[i]<=lb[i]:
				out.append(-math.pi/2)
			elif params[i]>=ub[i]:
				out.append(math.pi/2)
			else :
				 bla=2*(params[i]-lb[i])/float((ub[i]-lb[i]))-1
				 out.append(2*math.pi+math.asin(max(-1,min(1,bla))))
			#print "p[%f]: %f and lb,ub are (%f,%f) --> %f"  %  (i, params[i],lb[i],ub[i],out[i] )
		return out
		
  	def uncToOriginalParams(self, params) :
  		#convert unconstrained variables to real-parameter space
  		#for now, assume both lb and ub are specified
  		out=[]
  		for i in range(len(params)):
  			bla=float((math.sin(params[i])+1))/2
  			out.append(self.__lbounds[i]+(self.__ubounds[i]-self.__lbounds[i])*bla)
  		return out
    
	def value(self, params):	 #return residual value
		if self.__isConstrained :
			params=self.uncToOriginalParams(params)
			
		bla=Gaussian2D(self.__nx, self.__ny, params)
		calc=bla.compute()	    
		residual=0.0
		d=self.__data
		if len(calc)==len(self.__data) : #subtract them
			for i in range(len(calc)) :
				residual+=(calc[i]-d[i])**2
		else:
			#somehow calc and data vectors are not the same length
			raise

		return residual   


class GaussianFit:
	__width=0
	__height=0
	__realData=0.0
	__calcData=0.0
	__steps=[]
	__firstGuess=[]
	__bestGuess=[]
	__maxIter=1000
	__lbounds=[]
	__ubounds=[]
	__ip=[]
	__objSigma=3		#assume 3 pixel object width for first guess
	__fitResult=[]
	__convergenceVal=1e-4
    
	def reshapeArray(self,arr):
    #reshape 2DJava arrays from ImageJ for fitting in my code
		px=[]
		col =len(arr)
		row=len(arr[0])			#reshape list of pixels to match formatting for 2D Gaussian Fitting routines
		for i in range(row):
			for j in range(col):
				px.append(arr[j][i])
		return px

    	
	def __init__(self, data=0.0, width=0, height=0):
		#constructor needs to check whether it got an ImagePlus, ImageProcessor or a data array
		self.__ip=[]
		pixels=[]
		if isinstance(data, ij.ImagePlus) : #check to see type of object we were sent
			data=data.getProcessor()  #if ImagePlus, just keep the processor
		
		if isinstance(data, ij.process.FloatProcessor):
			self.__ip=data
			pixels=self.__ip.getFloatArray()
		elif isinstance(data, ij.process.ByteProcessor) or isinstance(data, ij.process.ShortProcessor):
			self.__ip=data
			pixels=self.__ip.getIntArray()
		elif isinstance(data, list):  #just got a list of values
			self.__realData=data
			self.__width=int(width)
			self.__height=int(height)
		else:
			sys.exit('Gaussian Fitter could not initialize because data was not recognized')
		
		#assign width, height and pixel values based on processor, if one was specified
		if self.__ip!=[]:
			self.__width=self.__ip.getWidth()
			self.__height=self.__ip.getHeight()
			self.__realData=self.reshapeArray(pixels)

	def getBounds(self):
		return (self.__lbounds, self.__ubounds)

	def getData(self):
		return self.__realData

	def getFitParams(self):
		pass

	def getWidth(self):
		return self.__width
	    
	def getHeight(self):
		return self.__height

	def getSteps(self):
		return self.__steps

	def getFirstGuess(self):
		return self.__firstGuess

	def getBestGuess(self):
		return self.__bestGuess

	def getResidual(self, p=[]):
		if p==[]:
			p=self.__bestGuess
		calcMin=Gaussian2DMinimizer()
		calcMin.setImage(self.__realData, self.__width, self.__height)
		return calcMin.value(p)

	
  	def getConvergenceValue(self):
  		return self.__convergenceVal
			
	def estimateGaussianParametersCOM(self):
	#try to estimate first guess for fitting a Gaussian based on center of mass
		if self.__ip!=[] :
			bla=self.__ip.getStatistics()
			
			return [bla.xCenterOfMass, bla.yCenterOfMass, self.__objSigma, self.__objSigma, bla.max, bla.min]		
		else:
			sys.error('Could not estimate centroid, ImageProcessor not defined!')

	def estimateGaussianParametersBinary(self):
	#try to estimate first guess for fitting a Gaussian based on largest binary particle
		if self.__ip!=[] :
			ip_binary=self.__ip.duplicate()
			ip_binary.setAutoThreshold(ij.process.AutoThresholder.Method.IsoData, True) 
			#ip_binary.setAutoThreshold(ImageProcessor.ISODATA2, True) 
			ip_binary.threshold(int(ip_binary.getMinThreshold()))
			imp=ImagePlus('temp',ip_binary)
			#not finished yet
			print 'Threshold: ', int(ip_binary.getMinThreshold())
			
			#return [bla.xCenterOfMass, bla.yCenterOfMass, self.__objSigma, self.__objSigma, bla.max, bla.min]		
		else:
			sys.error('Could not estimate centroid, ImageProcessor not defined!')
		
	def setFirstGuess(self, p=[]) :	
		if p!=[] :
			self.__firstGuess=p
		else:
			pass #eventually have code to assemble first guess
		
	def setBestGuess(self, p=[]) :	
		if p!=[] :
			self.__bestGuess=p
		else:
			pass #eventually have code to assemble first guess
		    
	def setSteps(self, mutiplier=.05, offset=0.05):
		self.__steps=[]
		p=self.__firstGuess
		for i in range(len(p)) :	    #set up array of initial steps for simplex
			if p[i] !=0 :
				self.__steps.append(p[i]*mutiplier)
			else :
				self.__steps.append(offset)
	
	def setBounds(self, lb, ub):
  		self.__lbounds=lb
  		self.__ubounds=ub

  	def setConvergenceValue(self,val):
  		self.__convergenceVal=val

	def doFit(self, maxIter=10000):
		self.__maxIter=abs(maxIter)
		#set up Gaussain2DMinimizer function, include options for constrained fitting
		#constrained optimization is done as in fminsearchbnd.m MATLAB routine (google it)
		calcMin=Gaussian2DMinimizer()
		
		calcMin.setBounds(self.__lbounds,self.__ubounds)
		
		
		#check to see if there are constraints
		if len(self.__lbounds)==len(self.__firstGuess) and len(self.__ubounds)==len(self.__firstGuess):
			calcMin.setImage(self.__realData, self.__width, self.__height, self.__lbounds,self.__ubounds ,True)
			#printGuessBounds(self.__firstGuess, self.__lbounds, self.__ubounds)
			
			self.__firstGuess=calcMin.origToUnconstrainedParams(self.__firstGuess)
			#print "Constrained Fit"
			isConstrained=True
		else:
			calcMin.setImage(self.__realData, self.__width, self.__height, [],[] ,False)
			#print "UN Constrained Fit"
			isConstrained=False
		
		self.setSteps()
		n=array(self.__firstGuess, 'd')
		#set up NelderMead optimizer settings
		nMeadOptimizer=NelderMead()
		convergeChecker=SimpleScalarValueChecker(self.__convergenceVal,-1)	 #play with this later to speed up fitting?
		nMeadOptimizer.setStartConfiguration(array(self.__steps, 'd'))  #define step size in simplex
		nMeadOptimizer.setConvergenceChecker(convergeChecker)
		nMeadOptimizer.setMaxIterations(self.__maxIter);

		#t=time.time()
		#begin fitting at this point
		try:  #catch error if fit fails
			out=nMeadOptimizer.optimize(calcMin, GoalType.MINIMIZE, n)
		except:
			self.__fitResult=None
		else:
			#print "Number of iterations: %f" % nMeadOptimizer.getIterations()
			#print "Time: %.3f" % (time.time()-t)
		#print out.getPoint().tolist()
			if isConstrained:
				self.__fitResult= calcMin.uncToOriginalParams(out.getPoint().tolist())
			else:
				self.__fitResult= out.getPoint().tolist()

		return self.__fitResult  #return None if fit failed
		

	def setDataFromImg(self, img): #sets fitting data from ImagePlus
		ip=img
		if isinstance(img, ij.ImagePlus) : #check to see type of object we were sent
			ip=img.getProcessor()
		else :
			pass	 #error of some sort

		if isinstance(ip, ij.process.FloatProcessor):
			bla=ip.getFloatArray()
		elif isinstance(ip, ij.process.ByteProcessor) or isinstance(ip, ij.process.ShortProcessor):
			bla=ip.getIntArray()
		else:
			print 'weird format'
			raise

		px=[]
		col=len(bla)
		row=len(bla[0])			#reshape list of pixels to match formatting for 2D Gaussian Fitting routines
		for i in range(row):
			for j in range(col):
				px.append(bla[j][i])
		self.__realData=px
		    

#	params[0]=center(x)
#	params[1]=center(y)
#	params[2]=sigma(x)
#	params[3]=sigma(y)
#	params[4]=amp
#	params[5]=yoffset

class StackFitter:
	#stack fitting methods
	__imp=[]
	__nSlices=[]
	__width=[]
	__height=[]
	__prevFitParams=[]
	__initialGuess=[]
	__fitResults={}
	__sliceFitter=[]
	__keyFrames={}
	__objSigma=3		#assume 3 pixel object width for first guess
	headings=["x_c","y_c","x_width","y_width","Amplitude","offset"]

	def __init__(self, imp=[]):
		self.__imp=imp
		if isinstance(imp, ij.ImagePlus):
			self.__width=imp.getWidth()
			self.__height=imp.getHeight()
			self.__nSlices=imp.getImageStackSize()
			self.__sliceFitter=GaussianFit(imp.getStack().getProcessor(1))
		else:
			sys.exit('Cant fit stack because ImagePlus was not specified')
			
	def getInitialGuess(self):
		return self.__initialGuess

	def getFitResults(self):
		return self.__fitResults

	def printFitResults(self):
		for fr in sorted(self.__fitResults):
			print 'Frame: ', fr, 'Fit: ', self.__fitResults[fr]
	
	def showFitResultsTable(self):
		bla=ij.measure.ResultsTable() #new results table
		#self.__fitResults
		#self.printFitResults()
		
		for fr in sorted(self.__fitResults):
			bla.incrementCounter()
			bla.addValue('frame',fr)
			for i in range(6):
				bla.addValue(''.join(['LB ',self.headings[i]]), self.__fitResults[fr][1][i])
				bla.addValue(self.headings[i], self.__fitResults[fr][0][i])
				bla.addValue(''.join(['UB ',self.headings[i]]), self.__fitResults[fr][2][i])
				
			bla.addValue("Residual", self.__fitResults[fr][3])
		bla.show("Fit Results")

	def showFitResultsPlotWindow(self, fitColumn=0):
		#make lists of dictionary
		a=[]
		frames=sorted(self.__fitResults)
		for i in frames:
  			a.append(self.__fitResults[i][0][fitColumn])
		bla=ij.gui.Plot(''.join(['Fit Result for: ',self.headings[fitColumn]]), "Frame #", self.headings[fitColumn], array(frames,'f'),array(a,'f')) #new results table
		bla.show()

	def showFitResultsRGBImg(self):
		#create a new RGB stack with the fit oval marked in one color and original data in another
		impResult=ij.ImagePlus("Results",self.__imp.getStack())
		IJ.run(impResult, "RGB Color","")
		frames=sorted(self.__fitResults)
		for fr in frames:
			fitout=self.__fitResults[fr][0]
			ip=	impResult.getStack().getProcessor(fr)
			ip.setColor(java.awt.Color.MAGENTA)
			ip.drawOval(int(round(fitout[0]-fitout[2],0)),int(round(fitout[1]-fitout[3],0)),int(round(fitout[2]*2,0)),int(round(fitout[3]*2,0)))
		impResult.show()
		return impResult

	def getKeyFrames(self):
		return self.__keyFrames

	def setKeyFramesMultiROI(self):
		bla=RoiManager(True)
		rm=bla.getInstance()
		if rm!=None and rm.getCount()>0:
			roiArr=rm.getRoisAsArray()
			for i in range(len(roiArr)):
				name=rm.getName("%s"%i) 
				fr=rm.getSliceNumber(name)
				if fr>0:
					self.__keyFrames[fr]=[roiArr[i].getPolygon().xpoints[0], roiArr[i].getPolygon().ypoints[0]]

	def fitAllSlices(self):
		self.__fitResults={}
		lb=[]
		ub=[]
		params=[]
		moveSpeed=2 # max pixels allowed to move from frame to frame
		maxIntensityFluct=.75 #75% percentage intensity fluctuations allowed
		maxSigmaFluct=0.75 #75% percentage width fluctuations allowed
		self.__initialGuess=[]
		for i in range(1,self.__nSlices+1) :
			if i in self.__keyFrames:  #check to see if keyFrame is defined for this slice, if so, constrained fit
				print "Key frame defined in slice %i -- x: %i, y: %i" % (i, self.__keyFrames[i][0], self.__keyFrames[i][1])
				bla=self.__imp.getImageStack().getProcessor(i).getStatistics()
				#print bla
				self.__initialGuess=[int(self.__keyFrames[i][0]), int(self.__keyFrames[i][1]), self.__objSigma, self.__objSigma, int(bla.max), int(bla.min)]	

			params,residual=self.fitSlice(i,lb,ub)  #do the fit
			#if fit is reasonable, make these parameters first guess for next frame, also adjust lb,ub
			print "Frame %i of %i (%f done)" % (i, self.__nSlices+1,float(i)/(self.__nSlices+1))
			if params[2]<10 and params[2]>1 and params[3]<10 and params[3]>1 and params[4]>3*params[5]:
				self.__initialGuess=params #set intial guess for next frame
				#center position constrains
				lb=[0,0,0,0,0,0]
				ub=[0,0,0,0,0,0]
				lb[0]=max(params[0]-moveSpeed,0)
				lb[1]=max(params[1]-moveSpeed,0)
				ub[0]=min(params[0]+moveSpeed,self.__width)
				ub[1]=min(params[1]+moveSpeed,self.__height)

				#sigma constrains
				lb[2]=params[2]*maxSigmaFluct
				lb[3]=params[2]*maxSigmaFluct
				ub[2]=params[2]*(maxSigmaFluct+1)
				ub[3]=params[2]*(maxSigmaFluct+1)

				#amplitude
				lb[4]=params[4]*maxIntensityFluct
				ub[4]=params[4]*(maxIntensityFluct+1)

				#offset
				lb[5]=0
				ub[5]=params[4]
				self.__fitResults[i]=[params, lb,ub,residual]
			else: #fit was unreasobable, maybe object blinked, for now do unconstrained fit
				lb=[]
				ub=[]
				#self.__initialGuess=[]
				print "Unconstrained fit at frame: %i" % i
			#end if	
			

	def fitSlice(self, sl=1, lb=[], ub=[], numIter=10000):
	#fits 2D gaussian to slice in stack
		gFit=GaussianFit(self.__imp.getImageStack().getProcessor(sl))
		if self.__initialGuess==[]:  #check to see if initial guess was defined
			#print gFit.estimateGaussianParametersBinary()
			self.__initialGuess=gFit.estimateGaussianParametersCOM()
		
		if len(lb)==len(self.__initialGuess) and len(ub)==len(self.__initialGuess): #bounded fitting 
			#printGuessBounds(self.__initialGuess,lb,ub)
			gFit.setBounds(lb,ub)
		
		gFit.setFirstGuess(self.__initialGuess) 
		#lb=[1,1,1,1,0.25*guess[4],0.5*guess[5]]
		#ub=[w-1,h-1,w,h,2*guess[4],0.5*guess[4]]
		#test.setBounds(lb,ub)
		bestVal=gFit.doFit(numIter)
		return (bestVal,gFit.getResidual(bestVal))

	def setInitialGuess(self, params):
		self.__initialGuess=params

	def writeFittingResults(self, filename=""):
		if filename==None or filename=="":  #try to place filename in same folder as imp
			bla=self.__imp.getOriginalFileInfo()  #determine info on original file
			try:	#see if the file has been saved already
				dirpath= bla.directory #determine where it was saved
				fn,dummy=os.path.splitext(bla.fileName)
				fn+="_fit.txt"
				filename=os.path.join(dirpath, fn)
				print "Data saved to: "+filename
			except:
				IJ.showMessage("Save stack to file or define data export file name!")
				raise RuntimeException("Save stack to file")

		#now write to file		
		try:
			f= open(filename, "w" )
			#write header row
			
			for fr in sorted(self.__fitResults):
				st="%s" % fr
				for i in range(6):
					st+="\t%s"%self.__fitResults[fr][0][i] #saved just the fit values, not the bounds
				f.write(st+'\n')
		finally:
			f.close()

	def writeFitResultsRGB(self, impResult, filename=""):
		if filename==None or filename=="":  #try to place filename in same folder as imp
			bla=self.__imp.getOriginalFileInfo()  #determine info on original file
			try:	#see if the file has been saved already
				dirpath= bla.directory #determine where it was saved
				fn,dummy=os.path.splitext(bla.fileName)
				fn+="_Results.tif"
				filename=os.path.join(dirpath, fn)
				
				print "Results saved to: "+filename
			except:
				IJ.showMessage("Save stack to file or define data export file name!")
				raise RuntimeException("Save stack to file")

		#now write to file		
		IJ.save(impResult, filename)
    
def printGuessBounds(guess, lb, ub):
	for i in range(len(guess)):
		print "Parameter [%i]: %f < %f < %f" % (i, lb[i],guess[i], ub[i])

if __name__ == "__main__":

	w=10
	h=20
	p=[5, 8, 2.5,3.5,.7,0]
	p2=[5,5,5,5,10,0]
	maxIter=10000
    
	#bla=Gaussian2D(w,h,p)
	#bla.compute()
	#bla.plot2D()
	imp=IJ.getImage()
	testingStack=True
	if testingStack:
		test=StackFitter(imp)
		test.setKeyFramesMultiROI()
		print test.getKeyFrames()
		test.fitAllSlices()
		test.writeFittingResults()
		#out=test.getFitResults()
		impResult=test.showFitResultsRGBImg()
		test.writeFitResultsRGB(impResult)
		#test.showFitResultsTable()
		test.showFitResultsPlotWindow(1)
	else:
		test=GaussianFit(imp.getProcessor())
		guess=test.estimateGaussianParametersCOM()
		test.setFirstGuess(guess) 
		lb=[1,1,1,1,0.25*guess[4],0.5*guess[5]]
		ub=[w-1,h-1,w,h,2*guess[4],0.5*guess[4]]

		guess[1]=10
		printGuessBounds(guess, lb, ub)
		print "Residual from first guess: %f" % test.getResidual(guess)
	#lb=[1,1,1,1,1,0]
	#ub=[6,6,6,6,1,6]
		test.setBounds(lb,ub)
		bestVal=test.doFit(10000)
		print "Residual from fit result: %f" % test.getResidual(bestVal)
#	print test.uncToOriginalParams([7.268, 7.0129, 1.57, 1.57, 1.57,1.56])
#	print test.doFit()