Source code for neutompy.recon.nnfbp.astra_plugin

#-----------------------------------------------------------------------
#Copyright 2015 Daniel M. Pelt
#
#Contact: D.M.Pelt@cwi.nl
#Website: http://www.dmpelt.com
#
#
#This file is part of the PyNN-FBP, a Python implementation of the
#NN-FBP tomographic reconstruction method.
#
#PyNN-FBP is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.
#
#PyNN-FBP is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with PyNN-FBP. If not, see <http://www.gnu.org/licenses/>.
#
#-----------------------------------------------------------------------
import astra
import numpy as np
from six.moves import range
from scipy.signal import fftconvolve
import scipy.ndimage.filters as snf
import numpy.linalg as na
import glob
import tifffile
import scipy.io as sio

import os, errno

import random

from .TrainingData import MATTrainingData
from neutompy.image.image import read_image

[docs]def mkdir_p(path): try: os.makedirs(path) except OSError as exc: # Python >2.5 if exc.errno == errno.EEXIST and os.path.isdir(path): pass else: raise
[docs]def customFBP(W, f, s): sf = np.zeros_like(s) padded = np.zeros(s.shape[1]*2) l = int(s.shape[1]/2.) r = l+s.shape[1] bl = f.shape[0]/len(s) if len(f.shape)==1: f = np.tile(f,(sf.shape[0],1)) for i in range(sf.shape[0]): padded[l:r] = s[i] padded[:l] = padded[l] padded[r:] = padded[r-1] sf[i] = fftconvolve(padded,f[i],'same')[l:r] return (W.T*sf).reshape(W.vshape)
class plugin_prepare(astra.plugin.ReconstructionAlgorithm2D): """Prepares a training set for the NN-FBP method [1]. Options: 'hqrecfiles': HQ reconstruction files (in TIFF format), e.g. '/path/to/rec*.tiff'. 'traindir': folder where output training files should be stored 'z_id': z-index of current slice 'npick' (optional): number of random pixels to pick (per slice) 'nlinear' (optional): number of linear steps in exponential binning 'extra_ids' (optional): extra ASTRA data ids to use during reconstruction 'angle_dependent' (optional): compute an angle-dependent filter if set to "True" [1] Pelt, D. M., & Batenburg, K. J. (2013). Fast tomographic reconstruction from limited data using artificial neural networks. Image Processing, IEEE Transactions on, 22(12), 5238-5251. """ astra_name="NN-FBP-prepare" def initialize(self, cfg, hqrecfiles, z_id, traindir, nlinear=2, npick=100, extra_ids=None, angle_dependent=False): if extra_ids==None: extra_ids = [] self.extra_s = [astra.data2d.get_shared(i) for i in extra_ids] self.W = astra.OpTomo(cfg['ProjectorId']) self.npick = npick self.td = traindir mkdir_p(traindir) # self.rec = tifffile.imread(sorted(glob.glob(hqrecfiles))[z_id]) self.rec = read_image(sorted(glob.glob(hqrecfiles))[z_id]) self.bck = (self.W.T*np.ones_like(self.s)<self.s.shape[0]-0.5).reshape(self.v.shape) self.outfn = traindir + os.sep + "train_{:05d}.mat".format(z_id) self.z_id = z_id fs = self.s.shape[1] if fs%2==0: fs += 1 mf = int(fs/2) w=1 c=mf bas = np.zeros(fs,dtype=np.float32) self.basis = [] count=0 while c<fs: bas[:]=0 l = c r = c+w if r>fs: r=fs bas[l:r]=1 if l!=0: l = fs-c-w r = l+w if l<0: l=0 bas[l:r]=1 self.basis.append(bas.copy()) c += w count += 1 if count>nlinear: w=2*w if angle_dependent=="True": basis_ang = [] bas_ang = np.zeros((self.s.shape[0],fs),dtype=np.float32) for bas in self.basis: for i in range(self.s.shape[0]): bas_ang[i] = bas basis_ang.append(bas_ang.copy()) bas_ang[i][:] = 0 self.basis = basis_ang self.nf = len(self.basis) def run(self, iterations): out = np.zeros((self.npick,(1+len(self.extra_s))*self.nf+1)) pl = np.random.random((self.npick,2)) pl[:,0]*=self.v.shape[0] pl[:,1]*=self.v.shape[1] pl = pl.astype(np.int) for i in range(self.npick): while self.bck[pl[i,0],pl[i,1]]: pl[i,0] = np.random.random(1)*self.v.shape[0] pl[i,1] = np.random.random(1)*self.v.shape[1] out[:,-1] = self.rec[pl[:,0],pl[:,1]] for i, bas in enumerate(self.basis): img = customFBP(self.W, bas, self.s) out[:,i] = img[pl[:,0],pl[:,1]] for l, s in enumerate(self.extra_s): img = customFBP(self.W, bas, s) out[:,i+(l+1)*self.nf] = img[pl[:,0],pl[:,1]] sio.savemat(self.outfn,{'mat':out},do_compression=True) #astra.log.info('Slice {} done...'.format(self.z_id)) import numexpr import time import scipy.sparse as ss import scipy.linalg as la try: import scipy.linalg.fblas as fblas hasfblas=True except: hasfblas=False
[docs]def sigmoid(x): '''Sigmoid function''' return numexpr.evaluate("1./(1.+exp(-x))")
[docs]class Network(object): ''' The neural network object that performs all training and reconstruction. :param nHiddenNodes: The number of hidden nodes in the network. :type nHiddenNodes: :class:`int` :param projector: The projector to use. :type projector: A ``Projector`` object (see, for example: :mod:`nnfbp.SimpleCPUProjector`) :param trainData: The training data set. :type trainData: A ``DataSet`` object (see: :mod:`nnfbp.DataSet`) :param valData: The validation data set. :type valData: A ``DataSet`` object (see: :mod:`nnfbp.DataSet`) :param reductor: Optional reductor to use. :type reductor: A ``Reductor`` object (see: :mod:`nnfbp.Reductors`, default:``LogSymReductor``) :param nTrain: Number of pixels to pick out of training set. :type nTrain: :class:`int` :param nVal: Number of pixels to pick out of validation set. :type nVal: :class:`int` :param tmpDir: Optional temporary directory to use. :type tmpDir: :class:`string` :param createEmptyClass: Used internally when loading from disk, to create an empty object. Do not use directly. :type createEmptyClass: :class:`boolean` ''' def __init__(self, nHiddenNodes, trainData, valData, setinit=None): self.tTD = trainData self.vTD = valData self.nHid = nHiddenNodes self.nIn = self.tTD.getDataBlock(0).shape[1]-1 self.jacDiff = np.zeros((self.nHid) * (self.nIn+1) + self.nHid + 1); self.jac2 = np.zeros(((self.nHid) * (self.nIn+1) + self.nHid + 1, (self.nHid) * (self.nIn+1) + self.nHid + 1)) self.setinit = setinit def __inittrain(self): '''Initialize training parameters, create actual training and validation sets by picking random pixels from the datasets''' self.l1 = 2 * np.random.rand(self.nIn+1, self.nHid) - 1 if self.setinit is not None: self.l1.fill(0) nd = self.nIn/self.setinit[0] for i,j in enumerate(self.setinit[1]): self.l1[j*nd:(j+1)*nd,i] = 2 * np.random.rand(nd) - 1 self.l1[-1,i] = 2 * np.random.rand(1) - 1 beta = 0.7 * self.nHid ** (1. / (self.nIn)) l1norm = np.linalg.norm(self.l1) self.l1 *= beta / l1norm self.l2 = 2 * np.random.rand(self.nHid + 1) - 1 self.l2 /= np.linalg.norm(self.l2) self.minl1 = self.l1.copy() self.minl2 = self.l2.copy() self.minmax = self.tTD.getMinMax() self.tTD.normalizeData(self.minmax[0], self.minmax[1], self.minmax[2], self.minmax[3]) self.vTD.normalizeData(self.minmax[0], self.minmax[1], self.minmax[2], self.minmax[3]) self.ident = np.eye((self.nHid) * (self.nIn+1) + self.nHid + 1) def __processDataBlock(self,data): ''' Returns output values (``vals``), 'correct' output values (``valOut``) and hidden node output values (``hiddenOut``) from a block of data.''' valOut = data[:, -1].copy() data[:, -1] = -np.ones(data.shape[0]) hiddenOut = np.empty((data.shape[0],self.l1.shape[1]+1)) hiddenOut[:,0:self.l1.shape[1]] = sigmoid(np.dot(data, self.l1)) hiddenOut[:,-1] = -1 rawVals = np.dot(hiddenOut, self.l2) vals = sigmoid(rawVals) return vals,valOut,hiddenOut def __getTSE(self, dat): '''Returns the total squared error of a data block''' tse = 0. for i in range(dat.nBlocks): data = dat.getDataBlock(i) vals,valOut,hiddenOut = self.__processDataBlock(data) tse += numexpr.evaluate('sum((vals - valOut)**2)') return tse def __setJac2(self): '''Calculates :math:`J^T J` and :math:`J^T e` for the training data. Used for Levenberg-Marquardt method.''' self.jac2.fill(0) self.jacDiff.fill(0) for i in range(self.tTD.nBlocks): data = self.tTD.getDataBlock(i) vals,valOut,hiddenOut = self.__processDataBlock(data) diffs = numexpr.evaluate('valOut - vals') jac = np.empty((data.shape[0], (self.nHid) * (self.nIn+1) + self.nHid + 1)) d0 = numexpr.evaluate('-vals * (1 - vals)') ot = (np.outer(d0, self.l2)) dj = numexpr.evaluate('hiddenOut * (1 - hiddenOut) * ot') I = np.tile(np.arange(data.shape[0]), (self.nHid + 1, 1)).flatten('F') J = np.arange(data.shape[0] * (self.nHid + 1)) Q = ss.csc_matrix((dj.flatten(), np.vstack((J, I))), (data.shape[0] * (self.nHid + 1), data.shape[0])) jac[:, 0:self.nHid + 1] = ss.spdiags(d0, 0, data.shape[0], data.shape[0]).dot(hiddenOut) Q2 = np.reshape(Q.dot(data), (data.shape[0],(self.nIn+1) * (self.nHid + 1))) jac[:, self.nHid + 1:jac.shape[1]] = Q2[:, 0:Q2.shape[1] - (self.nIn+1)] if hasfblas: self.jac2 += fblas.dgemm(1.0,a=jac.T,b=jac.T,trans_b=True) self.jacDiff += fblas.dgemv(1.0,a=jac.T,x=diffs) else: self.jac2 += np.dot(jac.T,jac) self.jacDiff += np.dot(jac.T,diffs)
[docs] def train(self): '''Train the network using the Levenberg-Marquardt method.''' self.__inittrain() mu = 100000.; muUpdate = 10; prevValError = np.Inf bestCounter = 0 tse = self.__getTSE(self.tTD) curTime = time.time() self.allls = [] for i in range(1000000): self.__setJac2() try: dw = -la.cho_solve(la.cho_factor(self.jac2 + mu * self.ident), self.jacDiff) except la.LinAlgError: break done = -1 while done <= 0: self.l2 += dw[0:self.nHid + 1] for k in range(self.nHid): start = self.nHid + 1 + k * (self.nIn+1); if self.setinit is not None: nd = self.nIn/self.setinit[0] j = self.setinit[1][k] self.l1[j*nd:(j+1)*nd,k] += dw[start+j*nd:start + (j+1)*nd] self.l1[-1,k] += dw[start+self.nIn] else: self.l1[:, k] += dw[start:start + self.nIn+1] newtse = self.__getTSE(self.tTD) if newtse < tse: if done == -1: mu /= muUpdate if mu <= 1e-100: mu = 1e-99 done = 1; else: done = 0; mu *= muUpdate if mu >= 1e20: done = 2 break; self.l2 -= dw[0:self.nHid + 1] for k in range(self.nHid): start = self.nHid + 1 + k * (self.nIn+1); if self.setinit is not None: nd = self.nIn/self.setinit[0] j = self.setinit[1][k] self.l1[j*nd:(j+1)*nd,k] -= dw[start+j*nd:start + (j+1)*nd] self.l1[-1,k] -= dw[start+self.nIn] else: self.l1[:, k] -= dw[start:start + self.nIn+1] try: dw = -la.cho_solve(la.cho_factor(self.jac2 + mu * self.ident), self.jacDiff) except la.LinAlgError: done=2 gradSize = np.linalg.norm(self.jacDiff) if done == 2: break curValErr = self.__getTSE(self.vTD) if curValErr > prevValError: bestCounter += 1 else: prevValError = curValErr self.minl1 = self.l1.copy() self.minl2 = self.l2.copy() if (newtse / tse < 0.999): bestCounter = 0 else: bestCounter +=1 if bestCounter == 50: break if(gradSize < 1e-8): break tse = newtse astra.log.info('Validation set error: {}'.format(prevValError)) self.allls.append([self.minl1,self.minl2]) self.l1 = self.minl1 self.l2 = self.minl2 self.valErr = prevValError
[docs] def saveAllToDisk(self,fn): for i,k in enumerate(self.allls): sio.savemat(fn+"{}.mat".format(i),{'l1':k[0], 'l2':k[1], 'minmax':self.minmax},do_compression=True)
[docs] def saveToDisk(self,fn): '''Save a trained network to disk, so that it can be used later without retraining. :param fn: Filename to save it to. :type fn: :class:`string` ''' sio.savemat(fn,{'l1':self.l1,'l2':self.l2,'minmax':self.minmax},do_compression=True)
[docs]def plugin_train(traindir, nhid, filter_file, val_rat=0.5, setinit=None, saveAll=False): """Traing filters and weights using the NN-FBP method [1]. Options: 'traindir': folder where training files are stored 'nhid': number of hidden nodes to use 'filter_file': file to store trained filters in 'val_rat' (optional): fraction of training examples to use as validation 'saveAll' (optional): save filters at each iteration instead of final filter only [1] Pelt, D. M., & Batenburg, K. J. (2013). Fast tomographic reconstruction from limited data using artificial neural networks. Image Processing, IEEE Transactions on, 22(12), 5238-5251. """ fls = glob.glob(traindir + os.sep + '*.mat') random.shuffle(fls) nval = int(val_rat*len(fls)) val = MATTrainingData(fls[:nval]) trn = MATTrainingData(fls[nval:]) n = Network(nhid, trn, val,setinit=setinit) n.train() if saveAll: n.saveAllToDisk(filter_file) else: n.saveToDisk(filter_file)
class plugin_rec(astra.plugin.ReconstructionAlgorithm2D): """Reconstructs using the NN-FBP method [1]. Options: 'filter_file': file with trained filters to use 'nlinear' (optional): number of linear steps in exponential binning 'extra_ids' (optional): extra ASTRA data ids to use during reconstruction 'angle_dependent' (optional): compute an angle-dependent filter if set to "True" [1] Pelt, D. M., & Batenburg, K. J. (2013). Fast tomographic reconstruction from limited data using artificial neural networks. Image Processing, IEEE Transactions on, 22(12), 5238-5251. """ astra_name="NN-FBP" def initialize(self, cfg, filter_file, nlinear=2, extra_ids=None, angle_dependent=False): if extra_ids==None: extra_ids = [] self.extra_s = [astra.data2d.get_shared(i) for i in extra_ids] self.W = astra.OpTomo(cfg['ProjectorId']) fs = self.s.shape[1] if fs%2==0: fs += 1 mf = int(fs/2) w=1 c=mf bas = np.zeros(fs,dtype=np.float32) self.basis = [] count=0 while c<fs: bas[:]=0 l = c r = c+w if r>fs: r=fs bas[l:r]=1 if l!=0: l = fs-c-w r = l+w if l<0: l=0 bas[l:r]=1 self.basis.append(bas.copy()) c += w count += 1 if count>nlinear: w=2*w if angle_dependent=="True": basis_ang = [] bas_ang = np.zeros((self.s.shape[0],fs),dtype=np.float32) for bas in self.basis: for i in range(self.s.shape[0]): bas_ang[i] = bas basis_ang.append(bas_ang.copy()) bas_ang[i][:] = 0 self.basis = basis_ang self.nf = len(self.basis) fl = sio.loadmat(filter_file) self.l1 = fl['l1'] self.l2 = fl['l2'].transpose() minmax = fl['minmax'][0] minL = minmax[0] maxL = minmax[1] self.minIn = minmax[2] self.maxIn = minmax[3] mindivmax = minL/(maxL-minL) mindivmax[np.isnan(mindivmax)]=0 mindivmax[np.isinf(mindivmax)]=0 divmaxmin = 1./(maxL-minL) divmaxmin[np.isnan(divmaxmin)]=0 divmaxmin[np.isinf(divmaxmin)]=0 nHid = self.l1.shape[1] nsl = len(extra_ids)+1 dims = [nHid,nsl,] dims.extend(self.basis[0].shape) self.filters = np.empty(dims) self.offsets = np.empty(nHid) for i in range(nHid): wv = (2*self.l1[0:self.l1.shape[0]-1,i]*divmaxmin).transpose() self.filters[i] = np.zeros(dims[1:]) for t, bas in enumerate(self.basis): for l in range(nsl): self.filters[i,l] += wv[t+l*len(self.basis)]*bas self.offsets[i] = 2*np.dot(self.l1[0:self.l1.shape[0]-1,i],mindivmax.transpose()) + np.sum(self.l1[:,i]) def run(self, iterations): self.v[:]=0 for i in range(self.l2.shape[0]-1): mult = float(self.l2[i]) offs = float(self.offsets[i]) back = customFBP(self.W, self.filters[i,0],self.s) for l, s in enumerate(self.extra_s): back += customFBP(self.W, self.filters[i,l+1],s) self.v[:] += numexpr.evaluate('mult/(1.+exp(-(back-offs)))') self.v[:] = sigmoid(self.v-self.l2[-1]) self.v[:] = (self.v-0.25)*2*(self.maxIn-self.minIn) + self.minIn