xftomo_Gleber.py

#!/usr/bin/env python

#imports
from skimage.transform import radon, iradon
import scipy as sp
import Image
from pylab import *
import pdb
from readMDA import readMDA
import h5py
import os
import phantom
import pad
import vine_utils as vu
from images2gif import writeGif
import mayavi.mlab as mlab
import pickle
import scipy.ndimage as spn
import scipy.fftpack as spf

#constants
base_direc = '/home/david/data/Trahey'
images_direc = '/home/david/data/Trahey/images'
align_direc = '/home/david/data/Trahey/align'
sinograms_direc = '/home/david/data/Trahey/sinograms'
tomograms_direc = '/home/david/data/Trahey/tomograms'
raw_data_direc = '/home/david/data/Trahey/h5'
animations_direc = '/home/david/data/Trahey/animations'

direcs = [base_direc, images_direc, align_direc, sinograms_direc, tomograms_direc, 
			raw_data_direc, animations_direc]

fstring = '2xfm_{:04d}.h5'

# Default pad is the amount added to all projections so that when they are aligned usign the cross correlation
# there is no 'overhang' wrapped around the array. This amount is added to both sides of the array. So this number
# should larger than the largest alignment shift
default_pad = 50

def rotate(x, y=1):
	"""
	out = rotate(x, y)

	Cyclically rotate a list.

	Inputs
	------
	x: the list to rotate.

	y: the number of elements to rotate by.

	Outputs
	-------
	out: the rotated list.

	Example
	-------
	l = [1,2,3,4,5]
	rotate(l, 1)
	>>> [2,3,4,5,1]

	"""
	if len(x) == 0:
	  return x
	y = y % len(x) # Normalize y, using modulo - even works for negative y

	return x[y:] + x[:y]

def extract_angles(fn_string, fn_range):
	"""
	Convenience function for extracting the angles from a bunch of mda files

	Inputs
	------
	fn_string: the path to the mda files with {:04d} (or similar) for specifying the number.
	fn_range: a two element tuple specifying the dtart and end numbers

	Outputs
	-------

	A pretty list of mda number and angle which can be cut and pasted into a google doc.

	"""
	l=[]
	for i in range(fn_range[0], fn_range[1]):
		mda = readMDA(fn_string.format(i))
		l.append((i, mda[0]['2xfm:userStringCalc10.DD'][2]))

	print 'MDA\t\tAngle (degrees)'
	pstr = '\t{:04d}\t\t{:3.2f}'
	for element in l:
		printpstr.format(element[0], element[1])

class tomo_data(object):
	"""
	data3d = tomo_data([file, path],mda2theta=None, ref_channel=None, ref_projection=None)

	Provides access to tomo data.

	Inputs
	------

	file, path: a file or directory containing the projections.

	mda2theta: a file relating mda filenumber to projection angle (theta, degrees).
			This file is generated by creating a google docs spreadsheet with 
			first column mda and second column corresponding angle. Download as 
			text file (.tsv)

	ref_channel: the reference detector channel that will be used for 
				 registering images for alignment. Typically the highest
				 signal-to-noise detector channel should be chosen.

	ref_projection: the reference projection to begin the projection
					projection registration alignment.
	

	Member Variables
	----------------
	self.theta: a list of the tomographic angles

	self.data: a dictionary key, value pair of angle, 2D projections (arrays).

	self.alignment: a dictionary key, value pair of angle, ordered pair describing the horizontal and vertical shift to be applied to each projection
					to align them.


	"""
	def __init__(self, *args, **kwargs):
		super(tomo_data, self).__init__()
		self.arg = args
		self.theta = None
		self.data = None
		self.alignment = None
		for kwarg in kwargs:
			setattr(self, kwarg, kwargs[kwarg])

		for direc in direcs:
			if not os.path.isdir(direc):
				os.mkdir(direc)

		try:
			if self.mda2theta_fn: 
				self.mda2theta = np.genfromtxt(self.mda2theta_fn)[1:]
			if os.path.isfile(args[0]):
				self.read_single_file(args[0])
			if os.path.isdir(args[0]):
				self.read_multi_files(args[0])
		except:
			pass
		

	def phantom(self):
		self.theta = range(0,180)
		self.data = dict(zip(self.theta, radon(phantom.phantom())))
		

	def read_multi_files(self, dir_path, theta_filenumber=None, channel=None, hdfpath=None, hdf_channel_names=None):
		"""

		read_multi_files(dir_path, theta_filenumber=None, channel=0)

		Read in tomo data where each file contains a single 2D projection.

		Inputs
		------

		dir_path: directory containing the files.
				  Files must end in _xxxx.[mda, h5] where xxxx is an integer.

		theta_filenumber: a list of point pairs describing the angle for a given file.
						  e.g [(0, 256), (3, 257)] meaning 0 degrees corresponds to file 256, 3 degrees in 257.
						  Default: assumes every file in directory belongs to same dataset and projections are
						  spaced evenly between 0 & 180 degrees.

		channel: the detector channel to read in.
				 Reconstructions on XFM measurements with many det channels are performed separately.

		hdfpath: the HDF5 directory path pointing to the data.

		Outputs
		-------
		self.data, self.theta: populated with raw data
		"""
		
		if theta_filenumber is None: theta_filenumber = self.mda2theta
		if channel is None: channel = self.ref_channel
		if hdfpath is None: hdfpath = self.hdfpath
		if hdf_channel_names is None: hdf_channel_names = self.hdf_channel_names

		try:
			n_files = len(theta_filenumber)
		except: 
			n_files = len([name for name in os.listdir(dir_path) if os.path.isfile(name)])
			theta_filenumber = sp.arange(0,180, 180/n_files)

		no_nans = True
		self.data, self.theta = {}, []

		if not os.path.exists(images_direc):
			os.mkdir(images_direc)
		l=[]
		for (fnumber, theta) in theta_filenumber:
			fn = '/'.join([raw_data_direc, fstring]).format(int(fnumber))
			if math.isnan(theta): 
				l.append(int(fnumber))
			else:
				try:
					data = self.get_data(fn, channel, hdfpath=hdfpath, hdf_channel_names=hdf_channel_names)
					# Check for nans
					if math.isnan(sum(data)):
						if no_nans: print 'WARNING: Found some NaNs. Will set them to 0 and silence this message.'
						no_nans = False
						data=np.nan_to_num(data)

					self.data['{}'.format(theta)] = data
					sp.misc.toimage(data).save(images_direc+'/raw_channel_{0:d}_angle_{1:3.2f}.png'.format(int(channel), theta))
				except (IOError, KeyError):
					print 'Could not open file: {}'.format(fn)

		if len(l)>0:
			print 'These mda files were skipped: ', l


	def get_data(self, path, channel, hdfpath='/MAPS/XRF_fits', hdf_channel_names='MAPS/channel_names'):
		"""
		get_data(path, channel, hdfpath=None)

		Open a file and return the 2D or 3D dataset corresponding to channel.

		Inputs
		------
		path: path to filename

		channel: Detector channel number.

		hdfpath: The HDF5 directory structure path.

		Outputs
		-------
		returns a 2D or 3D raw data set.

		Example: data = get_data('file_xxxx.h5', 0, hdfpath = '/MAPS/XRF_fits')
				 data = get_data('file_xxxx.mda', 0)


		"""
		self.channel = channel

		try:	
			if path.split('.')[-1].lower()=='mda':
				print 'MDA files not implemented yet.'

			elif path.split('.')[-1].lower() in ['h5', 'hdf5']:
				f = h5py.File(path)
				self.num_channels = len(f[hdfpath])
				try:
					self.channel_names = f[hdf_channel_names].value
				except:
					self.channel_names = None
				return f[hdfpath][channel]
				
		except (NameError, IndexError): 
			print 'The mda or hdf tag or channel specified could not be found.'
			print_file_metadata(path)
			return None

	def print_file_metadata(self, path):
		"""
		print_file_metadata(path)

		Print metadata from hdf5 or mda files.

		Inputs
		------

		path: path to file to be interrogated.

		Outputs
		-------

		Dataset is described onscreen.

		"""

		if path.split('.')[-1].lower() in ['hdf5', 'h5']:
			print 'HDF5 File: '+path
			print '-'*80
			L=[]
			try:
				f=h5py.File(path)
			except:
				'{} is not a valid hdf5 file'.format(path)
			f.visit(L.append)
		for element in L:
			print 'H5 Tag {}:\n\tDatatype: {}\n\tShape: {}\n\tValue: {}'.format(element, 
				f[element].dtype, f[element].shape, f[element].value)
			print '-'*80


	def read_single_file(self, path):
		"""Single file contains all projections"""
		pass

	def shift(self, arr1, nx, ny):
	    """
	    Shifts an array by nx and ny respectively.
	    """
	    
	    xpix, ypix = arr1.shape
	    
	    if ((nx % 1. == 0.) and (ny % 1. ==0)):
	        return sp.roll(sp.roll(arr1, int(nx), axis=0),
	                       int(ny), axis=1 )
	    else:
	        xfreqs, yfreqs = spf.fftfreq(xpix), spf.fftfreq(ypix)
	        phaseFactor = sp.zeros((xpix,ypix),dtype=complex)
	        for i in xrange(xpix):
	            for j in xrange(ypix):
	                phaseFactor[i,j] = sp.exp(complex(0., -2.*sp.pi)*(ny*yfreqs[j]+nx*xfreqs[i]))
	        tmp = spf.ifft2(spf.fft2(arr1)*phaseFactor)
	        return sp.real(tmp.copy())


class tomo_recon(tomo_data):
	"""
	recon3d = tomo_recon()

	Functions defined on tomo_data objects.

	Inputs
	------
	None

	Outputs
	-------
	None

	"""
	def __init__(self, *args, **kwargs):
		super(tomo_recon, self).__init__(*args, **kwargs)
		self.args = args
		self.kwargs = kwargs

	def align_data(self, align_to = None):
		"""
		data.align_data(align_to = None)

		If data.alignment is None, uses a phase cross correlation to align 2D projections.
		Else, uses data.alignment to align the data.

		The phase correlation is between the initial projection (specified by align_to) and the next increasing value projection.
		It then wraps at the end of the stack and continues from the beginning.

		Typically you want to use the storngest signal to calculate the alignment 'shift' values and then 
		use those for all other detector channels.

		Inputs
		------
		data: an instance of tomo_data containing the projections

		align_to: choose which projection to begin the alignment.

		Outputs
		-------
		data: an in-place shifted version of the raw data which should now be aligned

		self.alignment: If self.alignment was None, an dictionary of theta: (x,y) values describing the shifts calculated.
						Else, no change.

		"""
		
		assert self.data is not None and self.theta is not None, 'tomo_data instance does not contain any data.'

		if self.channel_names is not None:
			path = align_direc+'/channel_{}'.format(self.channel_names[self.channel])
		else:
			align_direc+'/channel_{:02d}'.format(self.channel)
		if not os.path.exists(path):
			os.mkdir(path)

		self.theta = self.data.keys()
		self.theta.sort(key=float)
		np.save(align_direc+'/theta', np.array(self.theta))

		# Zero pad all arrays to the same shape
		# Step 1 - Determine array shapes
		shapes = []
		[shapes.append(self.data[key].shape) for key in self.data if self.data[key].shape not in shapes]
		print 'Found {:d} different array shapes'.format(len(shapes))+(' {} '*len(shapes)).format(*shapes)
		maxx, maxy = 0,0
		for x, y in shapes:
			if x>maxx: maxx=x
			if y>maxy: maxy=y
		for key in self.data:
			shape = self.data[key].shape
			self.data[key] = pad.with_constant(self.data[key], pad_width=((default_pad, maxx+default_pad-shape[0]  ),
																		  (default_pad, maxy+default_pad-shape[1])))
		
		shape = self.data[self.data.keys()[0]].shape
		print 'New uniform array size: {:d}x{:d}'.format(shape[0], shape[1])

		if not self.alignment:
			n_projections = len(self.data)
			assert align_to in self.data.keys(), "The reference projection 'align_to' must be one of the projection angles: {}".format(self.data.keys())
			
			sp.save(align_direc+'/theta', sp.array(self.theta))

			self.alignment = {}
			proj_num = 0
			for i in range(n_projections):
				if self.theta[i]==align_to:
					proj_num = i
			
			correl_range = range(n_projections)
			correl_range = correl_range[proj_num:]+correl_range[:proj_num]
			print 'Aligning images...'
			for i in correl_range:
				try:
					x, y = sp.unravel_index(sp.argmax(sp.real(vu.phase_corr(self.data[self.theta[i]], self.data[self.theta[i+1]]))), \
							shape)
					if x>shape[0]/2: x -= shape[0]
					if y>shape[1]/2: y -= shape[1]
					self.alignment[self.theta[i+1]] = (y,x)
					self.data[self.theta[i+1]] = vu.shift(self.data[self.theta[i+1]], *self.alignment[self.theta[i+1]])
					sp.misc.toimage(self.data[self.theta[i+1]]).save(path+'/align_theta_{}.png'.format(self.theta[i+1]))
					print 'Aligning image {}. Shift by x: {}, y: {}'.format(i, x, y)

				except IndexError:
					x, y = sp.unravel_index(sp.argmax(sp.real(vu.phase_corr(self.data[self.theta[i]], self.data[self.theta[0]]))), \
							self.data[self.theta[0]].shape)
					self.alignment[self.theta[0]] = (y,x)
		else:
			for i in range(len(self.theta)):
				y, x = self.alignment[self.theta[i]]
				self.data[self.theta[i]] = vu.shift(self.data[self.theta[i]], *self.alignment[self.theta[i]])
				sp.misc.toimage(self.data[self.theta[i]]).save(path+'/align_theta_{}.png'.format(self.theta[i]))
				print 'Aligning image {}. Shift by x: {}, y: {}'.format(i, x, y)
		print 'Writing GIF...'
		l=[]
		for theta in self.theta:
			element = self.data[theta]
			element -= element.min()
			element /= element.max()
			l.append(element)
		writeGif(path+'/aligned.gif', l, duration=0.2)
		

	def plot_ccorrelation(self):
		"""
		plot_ccorrelation()
		
		Plot the cross correlation between subsequest frames to gauge how well the
		tomographic alignment has been done.

		Inputs
		------

		self.data: the data on which to perform the correlation.

		Outputs
		-------

		plot: the cross correlation as a function of theta

		"""
		assert self.data is not None and self.theta is not None, 'tomo_data instance does not contain any data.'

		ccorr = sp.zeros(len(self.theta))
		for i in xrange(len(self.theta)):
			ccorr[i] = vu.cross_corr(self.data[self.theta[i]], self.data[self.theta[i+1]])

		plot(self.theta, ccorr)
		
	def create_sinogram(self, axis=0, chan_name=None):
		"""
		create_sinogram(axis=1)

		Takes a stack of (x,y) at theta projections and creates a stack of
		(x, theta) at y (or, (y, theta) at x) sinograms.

		Inputs
		------

		axis: specifies whether to express the sinogram as a function of x (0) or y (1).

		self.data: a dictionary of theta: (x,y) arrays representing the aligned
			projections.

		chan_name: an optional name that will be used in the filename of the sinogram pngs.

		Outputs
		-------

		self.sino3d: a 3D array in order xytheta.

		self.sino_axis: record the user specified free axis of the sinogram

		PNG file for each sinogram stored in ./sinograms_direc


		"""
		print 'Creating sinogram...'
		assert self.data is not None and self.theta is not None, 'tomo_data instance does not contain any data.'
		
		if self.channel_names is not None:
			path = sinograms_direc+'/channel_{}'.format(self.channel_names[self.channel])
		else: 
			path = sinograms_direc+'/channel_{:02d}'.format(self.channel)

		if not os.path.exists(path):
			os.mkdir(path)

		self.sino_axis = axis

		x_len, y_len = self.data[self.theta[0]].shape
		self.sino3d = sp.zeros((x_len, y_len, len(self.theta)))
		for i, theta in enumerate(self.theta):
			self.sino3d[:,:,i] = self.data[theta]

		dim = {
			'0': x_len,
			'1': y_len
		}['{}'.format(axis)]

		np.save(path+'/sino', self.sino3d )
		for i in range(dim):
			if axis==0:
				sp.misc.toimage(self.sino3d[i,:,:]).save(path+'/sino_{}.png'.format(i))
			if axis==1:
				sp.misc.toimage(self.sino3d[:,i,:]).save(path+'/sino_{}.png'.format(i))

		print 'Writing GIF...'
		l=[]
		for i in range(dim):
			if axis==0:
				element = self.sino3d[i,:,:].copy()
			elif axis==1:
				element = self.sino3d[:,i,:]
			element -= element.min()
			element /= element.max()
			l.append(element)
		writeGif(path+'/sino.gif', l, duration=0.2)

		print 'Creating sinogram... Done.'

	def reconstruct_sinogram(self, sino_offset=None, tomo_slice=None, channel_name=None,
		tomo_filter='shepp-logan', tomo_interpolation='linear'):

		"""
		reconstruct_sinogram()

		Loops over self.sino3d sequentially applying the Inverse Radon transform

		Inputs
		------
		self.sino3d: the 3D cube of xytheta data.

		self.theta: a list of theta values corresponding to axis 2 of self.sino3d

		self.channel_names: optional name to use when saving the reconstruction

		sino_offset: offset the sinogram in the array by this much to center the rotation axis in the array

		slice: If specified a single slice will be reconstructed. If None, all will be reconstructed.

		Outputs
		-------
		self.tomo: a 3D array containing the reconstruction. This is written to the tomograms_direc
			and saved under tomo[_chan_name].dat.

		pngs of the 2D tomogram slices stored in tomograms_direc

		"""
		assert self.sino3d is not None and self.theta is not None, 'tomo_data instance does not contain the sinogram.'

		path = tomograms_direc+'/channel_{}'
		try:
			self.channel_names
		except AttributeError:
			self.channel_names=None

		if self.channel_names is not None:
			path=path.format(self.channel_names[self.channel])
			channel = self.channel_names[self.channel]
		elif channel_name is not None:
			path=path.format(channel_name)
			channel = channel_name
		else:
			path = tomograms_direc+'/channel_{:2d}'.format(self.channel)
			channel = self.channel

		if not os.path.exists(path):
			os.mkdir(path)

		x_len, y_len = self.sino3d.shape[:2]

		if tomo_slice is None:
			dim_range = range(x_len)
		else:
			dim_range = range(tomo_slice, tomo_slice+1)

		theta = sp.array([float(theta) for theta in self.theta ])
		
		# Shift sinogram in array for tomo alignment
		if sino_offset==None:
			sino_offset = self.sino3d.shape[1]/2-spn.center_of_mass(self.sino3d)[1]
		print 'Sinogram offset applied: {:2.2e}'.format(sino_offset)
		self.sino3d = sp.roll(self.sino3d, int(np.round(sino_offset)), axis=1)
		
		first_run = True
		for i in dim_range:
			print 'Reconstructing slice: {:3d} out of {:3d} for channel {}'.format(i, x_len, channel)
			
			recon = iradon(self.sino3d[i,:,:], theta, filter=tomo_filter, interpolation=tomo_interpolation)
			if first_run:
				self.tomo = sp.zeros((x_len, recon.shape[0], recon.shape[1]))
				first_run = False
			self.tomo[i,:,:] = recon

			sp.misc.toimage(recon).save(path+'/tomo_{}.png'.format( i ))
				
		f_tomo_shape = open(path+'/tomo_shape.txt', 'w')
		f_tomo_shape.write('{:}'.format(self.tomo.shape))
		f_tomo_shape.close()
		np.save(path+'/tomo', self.tomo )
		# Export tomo in a format suitable for import into drishti
		self.tomo -= self.tomo.min()
		self.tomo *= 255/self.tomo.max()
		self.tomo = np.uint8(self.tomo)
		self.tomo.tofile(path+'/tomo.raw')


	def animate_tomogram(self):
		if self.channel_names is not None:
			path = tomograms_direc+'/channel_{}'.format(self.channel_names[self.channel])
		else:
			path = tomograms_direc+'/channel_{:2d}'.format(self.channel)

		s=mlab.contour3d(self.tomo)

		for i in range(720):
			s.scene.camera.azimuth(i)
			s.scene.render()
			s.scene.save_png(path+'/anim_{:04d}.png'.format(i))


		os.system('ffmpeg -qscale 5 -r 20 -i anim_%04d.png {}'.format(path+'/movie.mp4'))


def standard_tomo_reconstruction():

	"""
	standard_tomo_reconstruction()

	This routine defines a standard tomographic reconstruction
	from raw data to final reconstruction. 


	"""
	path = '/home/david/data/Trahey/h5'
	mda2theta_fn = '/home/david/data/Trahey/Trahey-Sample1-MDA_Angles.tsv'
	ref_channel = 10
	sino_offset = None
	tomo_filter = ['ramp', 'shepp-logan', 'cosine', 'hamming', 'han'][1]
	tomo_interpolation = ['linear', 'nearest'][1]

	# Reconstruct ref_channel first
	data = tomo_recon(path, 
					 mda2theta_fn=mda2theta_fn,
					 ref_channel=ref_channel,
					 hdfpath='/MAPS/XRF_fits', 
					 hdf_channel_names='MAPS/channel_names'
					)
	data.align_data(align_to = '-178.5')
	data.create_sinogram()
	data.reconstruct_sinogram(sino_offset=sino_offset, tomo_filter=tomo_filter, tomo_interpolation=tomo_interpolation)
	
	# Reconstruct remaining XRF channels
	l=range(data.num_channels)
	l.pop(ref_channel)
	for i in l:
		data.read_multi_files(path, channel=i, hdfpath='/MAPS/XRF_fits', hdf_channel_names='MAPS/channel_names')
		data.align_data()
		data.create_sinogram()
		data.reconstruct_sinogram(sino_offset=sino_offset, tomo_filter=tomo_filter, tomo_interpolation=tomo_interpolation)
	"""
	# Recosntruct DPC channels
	l=range(5,11)
	for i in l:
		data.read_multi_files(path, channel=i, hdfpath='/MAPS/scalers', hdf_channel_names='/MAPS/scaler_names')
		data.align_data()
		data.create_sinogram()
		data.reconstruct_sinogram()
	"""


def reconstruct_sinogram():

	# This function is useful for reconstructing a single tomo slice from an existing sinogram
	# to play with the sino offset to reduce tomo artifacts

	data = tomo_recon(
		sino3d = np.load('/home/david/data/Trahey/sinograms/channel_dia1_dpc_cfg'),
		theta = list(np.load('/home/david/data/Trahey/align/theta.npy'))
		)

	data.reconstruct_sinogram(sino_offset = 10, tomo_slice=101, channel_name='Mn')

def crop_hdf5(filename, hdf_string, crop, axis):
	# Do symmetric crop of all hdf channels
	f=h5py.File(filename)
	data = f[hdf_string]
	if 

if __name__ == '__main__': standard_tomo_reconstruction()