def geodesic_shooting_diffOp(moving, target, m0, steps, mType, config): grid = moving.grid() m0.setGrid(grid) ca.ThreadMemoryManager.init(grid, mType, 1) if mType == ca.MEM_HOST: diffOp = ca.FluidKernelFFTCPU() else: diffOp = ca.FluidKernelFFTGPU() diffOp.setAlpha(config['deformation_params']['diffOpParams'][0]) diffOp.setBeta(config['deformation_params']['diffOpParams'][1]) diffOp.setGamma(config['deformation_params']['diffOpParams'][2]) diffOp.setGrid(grid) g = ca.Field3D(grid, mType) ginv = ca.Field3D(grid, mType) mt = ca.Field3D(grid, mType) It = ca.Image3D(grid, mType) It_inv = ca.Image3D(grid, mType) if (steps <= 0): time_steps = config['deformation_params']['timeSteps']; else: time_steps = steps; t = [x*1./time_steps for x in range(time_steps+1)] checkpointinds = range(1,len(t)) checkpointstates = [(ca.Field3D(grid,mType),ca.Field3D(grid,mType)) for idx in checkpointinds] scratchV1 = ca.Field3D(grid,mType) scratchV2 = ca.Field3D(grid,mType) scratchV3 = ca.Field3D(grid,mType) CAvmCommon.IntegrateGeodesic(m0,t,diffOp, mt, g, ginv,\ scratchV1,scratchV2,scratchV3,\ keepstates=checkpointstates,keepinds=checkpointinds, Ninv=config['deformation_params']['NIterForInverse'], integMethod = config['deformation_params']['integMethod']) ca.ApplyH(It,moving,ginv) ca.ApplyH(It_inv,target,g) output={'I1':It, 'I1_inv': It_inv, "phiinv":ginv} return output
def __init__(self, I0, I1, sigma, t, cpinds, cpstates, alpha, beta, gamma, nIter, stepSize, maxPert, nTimeSteps, integMethod='RK4', optMethod='FIXEDGD', nInv=10, plotEvery=None, plotSlice=None, quiverEvery=None, outputPrefix='./'): """ Initialize everything with the size and type given """ print I0 self.I0 = I0 self.I1 = I1 self.grid = I0.grid() print(self.grid) self.memtype = I0.memType() # matching param self.sigma = sigma # initial conditions self.m0 = ca.Field3D(self.grid, self.memtype) ca.SetMem(self.m0, 0.0) # state variables self.g = ca.Field3D(self.grid, self.memtype) self.ginv = ca.Field3D(self.grid, self.memtype) self.m = ca.Field3D(self.grid, self.memtype) self.I = ca.Image3D(self.grid, self.memtype) self.residualIm = ca.Image3D(self.grid, self.memtype) # adjoint variables self.madj = ca.Field3D(self.grid, self.memtype) self.Iadj = ca.Image3D(self.grid, self.memtype) self.madjtmp = ca.Field3D(self.grid, self.memtype) self.Iadjtmp = ca.Image3D(self.grid, self.memtype) # time array self.t = t # checkpointing variables self.checkpointinds = cpinds self.checkpointstates = cpstates # set up diffOp if self.memtype == ca.MEM_HOST: self.diffOp = ca.FluidKernelFFTCPU() else: self.diffOp = ca.FluidKernelFFTGPU() self.diffOp.setAlpha(alpha) self.diffOp.setBeta(beta) self.diffOp.setGamma(gamma) self.diffOp.setGrid(self.grid) # energy self.Energy = None # optimization stuff self.nIter = nIter self.stepSize = stepSize self.maxPert = maxPert self.nTimeSteps = nTimeSteps self.optMethod = optMethod self.integMethod = integMethod self.nInv = nInv # for interative update to inverse deformation # plotting variables self.plotEvery = plotEvery self.plotSlice = plotSlice self.quiverEvery = quiverEvery self.outputPrefix = outputPrefix # scratch variables self.scratchV1 = ca.Field3D(self.grid, self.memtype) self.scratchV2 = ca.Field3D(self.grid, self.memtype) self.scratchV3 = ca.Field3D(self.grid, self.memtype) self.scratchV4 = ca.Field3D(self.grid, self.memtype) self.scratchV5 = ca.Field3D(self.grid, self.memtype) self.scratchV6 = ca.Field3D(self.grid, self.memtype) self.scratchV7 = ca.Field3D(self.grid, self.memtype) self.scratchV8 = ca.Field3D(self.grid, self.memtype) self.scratchV9 = ca.Field3D(self.grid, self.memtype) self.scratchV10 = ca.Field3D(self.grid, self.memtype) self.scratchV11 = ca.Field3D(self.grid, self.memtype)
def __init__(self, grid, mType, alpha, beta, gamma, nInv, sigma, StepSize, integMethod='EULER'): """ Initialize everything with the size and type given """ self.grid = grid self.memtype = mType # initial conditions self.I0 = None # this is a reference that always points to the atlas image self.m0 = None # this is a reference that gets assigned to momenta for an individual each time # state variables self.g = ca.Field3D(self.grid, self.memtype) self.ginv = ca.Field3D(self.grid, self.memtype) self.m = ca.Field3D(self.grid, self.memtype) self.I = ca.Image3D(self.grid, self.memtype) # adjoint variables self.madj = ca.Field3D(self.grid, self.memtype) self.Iadj = ca.Image3D(self.grid, self.memtype) self.madjtmp = ca.Field3D(self.grid, self.memtype) self.Iadjtmp = ca.Image3D(self.grid, self.memtype) # image variables for closed-form template update self.sumSplatI = ca.Image3D(self.grid, self.memtype) self.sumJac = ca.Image3D(self.grid, self.memtype) # set up diffOp if self.memtype == ca.MEM_HOST: self.diffOp = ca.FluidKernelFFTCPU() else: self.diffOp = ca.FluidKernelFFTGPU() self.diffOp.setAlpha(alpha) self.diffOp.setBeta(beta) self.diffOp.setGamma(gamma) self.diffOp.setGrid(self.grid) # some extras self.nInv = nInv # for interative update to inverse deformation self.integMethod = integMethod self.sigma = sigma self.stepSize = StepSize # TODO: scratch variables to be changed to using managed memory self.scratchV1 = ca.Field3D(self.grid, self.memtype) self.scratchV2 = ca.Field3D(self.grid, self.memtype) self.scratchV3 = ca.Field3D(self.grid, self.memtype) self.scratchV4 = ca.Field3D(self.grid, self.memtype) self.scratchV5 = ca.Field3D(self.grid, self.memtype) self.scratchV6 = ca.Field3D(self.grid, self.memtype) self.scratchV7 = ca.Field3D(self.grid, self.memtype) self.scratchV8 = ca.Field3D(self.grid, self.memtype) self.scratchV9 = ca.Field3D(self.grid, self.memtype) self.scratchV10 = ca.Field3D(self.grid, self.memtype) self.scratchV11 = ca.Field3D(self.grid, self.memtype) self.scratchI1 = ca.Image3D( self.grid, self.memtype) #only used for geodesic regression with RK4
def GeodesicShooting(cf): # prepare output directory common.Mkdir_p(os.path.dirname(cf.io.outputPrefix)) # Output loaded config if cf.io.outputPrefix is not None: cfstr = Config.ConfigToYAML(GeodesicShootingConfigSpec, cf) with open(cf.io.outputPrefix + "parsedconfig.yaml", "w") as f: f.write(cfstr) mType = ca.MEM_DEVICE if cf.useCUDA else ca.MEM_HOST #common.DebugHere() I0 = common.LoadITKImage(cf.study.I0, mType) m0 = common.LoadITKField(cf.study.m0, mType) grid = I0.grid() ca.ThreadMemoryManager.init(grid, mType, 1) # set up diffOp if mType == ca.MEM_HOST: diffOp = ca.FluidKernelFFTCPU() else: diffOp = ca.FluidKernelFFTGPU() diffOp.setAlpha(cf.diffOpParams[0]) diffOp.setBeta(cf.diffOpParams[1]) diffOp.setGamma(cf.diffOpParams[2]) diffOp.setGrid(grid) g = ca.Field3D(grid, mType) ginv = ca.Field3D(grid, mType) mt = ca.Field3D(grid, mType) It = ca.Image3D(grid, mType) t = [ x * 1. / cf.integration.nTimeSteps for x in range(cf.integration.nTimeSteps + 1) ] checkpointinds = range(1, len(t)) checkpointstates = [(ca.Field3D(grid, mType), ca.Field3D(grid, mType)) for idx in checkpointinds] scratchV1 = ca.Field3D(grid, mType) scratchV2 = ca.Field3D(grid, mType) scratchV3 = ca.Field3D(grid, mType) # scale momenta to shoot cf.study.scaleMomenta = float(cf.study.scaleMomenta) if abs(cf.study.scaleMomenta) > 0.000000: ca.MulC_I(m0, float(cf.study.scaleMomenta)) CAvmCommon.IntegrateGeodesic(m0,t,diffOp, mt, g, ginv,\ scratchV1,scratchV2,scratchV3,\ keepstates=checkpointstates,keepinds=checkpointinds, Ninv=cf.integration.NIterForInverse, integMethod = cf.integration.integMethod) else: ca.Copy(It, I0) ca.Copy(mt, m0) ca.SetToIdentity(ginv) ca.SetToIdentity(g) # write output if cf.io.outputPrefix is not None: # scale back shotmomenta before writing if abs(cf.study.scaleMomenta) > 0.000000: ca.ApplyH(It, I0, ginv) ca.CoAd(mt, ginv, m0) ca.DivC_I(mt, float(cf.study.scaleMomenta)) common.SaveITKImage(It, cf.io.outputPrefix + "I1.mhd") common.SaveITKField(mt, cf.io.outputPrefix + "m1.mhd") common.SaveITKField(ginv, cf.io.outputPrefix + "phiinv.mhd") common.SaveITKField(g, cf.io.outputPrefix + "phi.mhd") GeodesicShootingPlots(g, ginv, I0, It, cf) if cf.io.saveFrames: SaveFrames(checkpointstates, checkpointinds, I0, It, m0, mt, cf)
def ElastReg(I0Orig, I1Orig, scales=[1], nIters=[1000], maxPert=[0.2], fluidParams=[0.1, 0.1, 0.001], VFC=0.2, Mask=None, plotEvery=100): mType = I0Orig.memType() origGrid = I0Orig.grid() # allocate vars I0 = ca.Image3D(origGrid, mType) I1 = ca.Image3D(origGrid, mType) u = ca.Field3D(origGrid, mType) Idef = ca.Image3D(origGrid, mType) diff = ca.Image3D(origGrid, mType) gI = ca.Field3D(origGrid, mType) gU = ca.Field3D(origGrid, mType) scratchI = ca.Image3D(origGrid, mType) scratchV = ca.Field3D(origGrid, mType) # mask if Mask != None: MaskOrig = Mask.copy() # allocate diffOp if mType == ca.MEM_HOST: diffOp = ca.FluidKernelFFTCPU() else: diffOp = ca.FluidKernelFFTGPU() # initialize some vars nScales = len(scales) scaleManager = ca.MultiscaleManager(origGrid) for s in scales: scaleManager.addScaleLevel(s) # Initalize the thread memory manager (needed for resampler) # num pools is 2 (images) + 2*3 (fields) ca.ThreadMemoryManager.init(origGrid, mType, 8) if mType == ca.MEM_HOST: resampler = ca.MultiscaleResamplerGaussCPU(origGrid) else: resampler = ca.MultiscaleResamplerGaussGPU(origGrid) def setScale(scale): global curGrid scaleManager.set(scale) curGrid = scaleManager.getCurGrid() # since this is only 2D: curGrid.spacing().z = 1.0 resampler.setScaleLevel(scaleManager) diffOp.setAlpha(fluidParams[0]) diffOp.setBeta(fluidParams[1]) diffOp.setGamma(fluidParams[2]) diffOp.setGrid(curGrid) # downsample images I0.setGrid(curGrid) I1.setGrid(curGrid) if scaleManager.isLastScale(): ca.Copy(I0, I0Orig) ca.Copy(I1, I1Orig) else: resampler.downsampleImage(I0, I0Orig) resampler.downsampleImage(I1, I1Orig) if Mask != None: if scaleManager.isLastScale(): Mask.setGrid(curGrid) ca.Copy(Mask, MaskOrig) else: resampler.downsampleImage(Mask, MaskOrig) # initialize / upsample deformation if scaleManager.isFirstScale(): u.setGrid(curGrid) ca.SetMem(u, 0.0) else: resampler.updateVField(u) # set grids gI.setGrid(curGrid) Idef.setGrid(curGrid) diff.setGrid(curGrid) gU.setGrid(curGrid) scratchI.setGrid(curGrid) scratchV.setGrid(curGrid) # end function energy = [[] for _ in xrange(3)] for scale in range(len(scales)): setScale(scale) ustep = None # update gradient ca.Gradient(gI, I0) for it in range(nIters[scale]): print 'iter %d' % it # compute deformed image ca.ApplyV(Idef, I0, u, 1.0) # update u ca.Sub(diff, I1, Idef) if Mask != None: ca.Mul_I(diff, Mask) ca.ApplyV(scratchV, gI, u, ca.BACKGROUND_STRATEGY_CLAMP) ca.Mul_I(scratchV, diff) diffOp.applyInverseOperator(gU, scratchV) vfcEn = VFC * ca.Dot(scratchV, gU) # why is this negative necessary? ca.MulC_I(gU, -1.0) # u = u*(1-VFC*ustep) + (-2.0*ustep)*gU # MulC_Add_MulC_I(u, (1-VFC*ustep), # gU, 2.0*ustep) # u = u - ustep*(VFC*u + 2.0*gU) ca.MulC_I(gU, 2.0) # subtract average if gamma is zero (result of nullspace # of L for K(L(u))) if fluidParams[2] == 0: av = ca.SumComp(u) av /= scratchI.nVox() ca.SubC(scratchV, u, av) # continue computing gradient ca.MulC(scratchV, u, VFC) ca.Add_I(gU, scratchV) ca.Magnitude(scratchI, gU) gradmax = ca.Max(scratchI) if ustep is None or ustep * gradmax > maxPert: ustep = maxPert[scale] / gradmax print 'step is %f' % ustep ca.MulC_I(gU, ustep) # apply gradient ca.Sub_I(u, gU) # compute energy energy[0].append(ca.Sum2(diff)) diffOp.applyOperator(scratchV, u) energy[1].append(0.5 * VFC * ca.Dot(u, scratchV)) energy[2].append(energy[0][-1]+\ energy[1][-1]) if plotEvery > 0 and \ ((it+1) % plotEvery == 0 or \ (scale == nScales-1 and it == nIters[scale]-1)): print 'plotting' clrlist = ['r', 'g', 'b', 'm', 'c', 'y', 'k'] plt.figure('energy') for i in range(len(energy)): plt.plot(energy[i], clrlist[i]) if i == 0: plt.hold(True) plt.hold(False) plt.draw() plt.figure('results') plt.clf() plt.subplot(3, 2, 1) display.DispImage(I0, 'I0', newFig=False) plt.subplot(3, 2, 2) display.DispImage(I1, 'I1', newFig=False) plt.subplot(3, 2, 3) display.DispImage(Idef, 'def', newFig=False) plt.subplot(3, 2, 4) display.DispImage(diff, 'diff', newFig=False) plt.colorbar() plt.subplot(3, 2, 5) display.GridPlot(u, every=4) plt.subplot(3, 2, 6) display.JacDetPlot(u) plt.colorbar() plt.draw() plt.show() # end plot # end iteration # end scale return (Idef, u, energy)