Homec4science
Diffusion Motion correction paper (cluster_pipeline) Browse

R6289/code/of_l1_l2_fm_admm.py
dcd42431cf9ccluster_pipeline

Last Change
Raw File

of_l1_l2_fm_admm.py
View Options

from scipy.ndimage import gaussian_filter
import math
from skimage import transform
import numpy as np
def of_l1_l2_fm_admm(I1, I2, sz0, matches, param, v, disp):
#print('I1',I1)
#print('I2',I2)
#print('matches',matches)
#print('sz0',sz0)
sigmaPreproc = 0.9
#TODO compare results to m script
I1 = gaussian_filter(I1, sigmaPreproc, mode='mirror')
I2 = gaussian_filter(I2, sigmaPreproc, mode='mirror')
deriv_filter = np.array([1, -8, 0, 8, -1])/12.
#coarse-to-fine parameters
minSizeC2f = 10
c2fLevels = int(math.ceil(math.log(minSizeC2f / max(I1.shape))) / math.log(1/param['c2fSpacing']))
factor = math.sqrt(2)
smooth_sigma = math.sqrt(param['c2fSpacing'])/factor
#TODO test if these perform the same as in m script
I1C2f = transform.pyramid_gaussian(I1, c2fLevels, param['c2fSpacing'], smooth_sigma)
I2C2f = transform.pyramid_gaussian(I2, c2fLevels, param['c2fSpacing'], smooth_sigma)
lambdaC2f = np.empty(c2fLevels)
sigmaSSegC2f = np.empty(c2fLevels)
gammaC2f = np.empty(c2fLevels)
matchesC2f = np.empty((c2fLevels,matches.shape[0],matches.shape[1]))
matchesl = matches
for i in range(c2fLevels):
lambdaC2f[i] = param['lmbd']
sigmaSSegC2f[i] = param['sigmaS']/(param['c2fSpacing']**i)
gammaC2f[i] = param['gamma']*((c2fLevels-i+1)/c2fLevels)**1.8
matchesl[0,:] = np.maximum(np.around(matches[0,:]/param['c2fSpacing']**i),1)
matchesl[1,:] = np.maximum(np.around(matches[1,:]/param['c2fSpacing']**i),1)
matchesl[2,:] = matches[2,:]/param['c2fSpacing']**i
matchesl[3,:] = matches[3,:]/param['c2fSpacing']**i
matchesl[4,:] = matches[4,:]
matchesC2f[i,:] = matchesl
#Initationlization
wl = np.zeros((I1.shape[0], I1.shape[1], 2))
occ = np.ones(I1.shape)
#Coarse to fine
#next(I1C2f) #skip first element of pyramids
#next(I2C2f)
for l, I1, I2 in zip(range(c2fLevels), I1C2f, I2C2f):
if v:
print('\nScale {}\n'.format(l))
#Scaled data
#I1 = I1C2f[l]
#I2 = I2C2f[l]
#print('l',l)
#print('image1 size', I1.shape)
#print('image2 size', I2.shape)
sigmaS = sigmaSSegC2f[l]
matchesl = matchesC2f[l]
lmbd = lambdaC2f[l]
param['gamma'] = gammaC2f[l]
#Resacle flow
ratio = I1.shape[0] / wl[:,:,0].shape[0]
ul = transform.resize(wl[:,:,0],I1.shape, order=3) * ratio
ratio = I1.shape[1] / wl[:,:,1].shape[1]
vl = transform.resize(wl[:,:,1],I1.shape, order=3) * ratio
wl = np.dstack((ul,vl))
sz0 = np.floor(np.array(sz0)*ratio)
#Create binary and motion vectors fields
c = np.zeros(I1.shape)
m = np.zeros((I1.shape[0], I1.shape[1], 3))
matchesl = matchesl.astype(np.int_)
#print('matchesl',matchesl)
#print('matchesl shape',matchesl.shape)
#print('c shape', c.shape)
for i in range(0,matches.shape[1]):
c[matchesl[1,i],matchesl[0,i]] = 1
if matchesl[4,i] > m[matchesl[1,i], matchesl[0,i], 2]:
m[matchesl[1,i],matchesl[0,i],0] = matchesl[2,i]
m[matchesl[1,i],matchesl[0,i],1] = matchesl[1,i]
m[matchesl[1,i],matchesl[0,i],2] = matchesl[0,i]
# no weights
m[:,:,2] = np.ones(I1.shape)
occ = transform.resize(occ, I1.shape, order=0)
occ = occ.astype(np.bool_)
mu = param['mu']
nu = param['nu']
#TODO lin operator
for iWarp = range(parma['nbWarps']):
if v:
print('Warp {}\n'.format(iWarp))
wPrev = wl
dul = np.zeros(wl.shape[0], wl.shape[1])
dvl = np.zeros(wl.shape[0], wl.shape[1])
uu1 = np.zeros(wl.shape[0], wl.shape[1])
uu2 = np.zeros(wl.shape[0], wl.shape[1])
dwl = np.zeros(wl.shape)
#ADMM variables
alpha = np.zeros((I1.shape[0],I1.shape[1],2))
beta = np.zeros((I1.shape[0],I1.shape[1],2))
z = np.zeros((I1.shape[0],I1.shape[1],2))
u = np.zeros((I1.shape[0],I1.shape[1],2))
#No occlusion handling
occ = np.ones(I1.shape)
#Pre-computations
[It, Ix, Iy] = #TODO figure out how to do this
#Main iterations loop
for it = range(param['maxIters']):
thresh = Igrad/(mu+nu)
thresh2 = param['gamma']*m[:,:,2]/nu #TODO check if operations are element-wise
#Data update
r1 = z - wl - alpha/mu
r2 = u - wl + beta/nu
t = (mu*r1+nu*r2)/(mu+nu)
t1 = t[:,:,1]
t2 = t[:,:,2]
rho = It + t1*Ix + t2*Iy #TODO check if multiplication is element-wise
idx1 = rho < -thresh
idx2 = rho > thresh
idx3 = np.abs(rho) <= thresh
dul[idx1] = t1[idx1] + Ix[idx1]/[mu+nu]
dvl[idx1] = t2[idx1] + Iy[idx1]/[mu+nu]
dul[idx2] = t1[idx2] - Ix[idx2]/[mu+nu]
dvl[idx2] = t2[idx2] - Iy[idx2]/[mu+nu]
dul[idx3] = t1[idx3] - rho[idx3]*Ix[idx3]/Igrad[idx3]#TODO check element-wise multiplication and devision
dvl[idx3] = t2[idx3] - rho[idx3]*Iy[idx3]/Igrad[idx3]#TODO check element-wise multiplication and devision
dul[idocc] = t1[idocc]
dvl[idocc] = t2[idocc]
dwl = np.dstack(dul,dvl)
#Regularization update
z[:,:,0] = real(ifft2(fft2(mu*(dwl(:,:,1)+wl(:,:,1))+alpha(:,:,1))./eigsDtD))
z[:,:,1] = real(ifft2(fft2(mu*(dwl[:,:,1]+wl[:,:,1])+alpha[:,:,1])./eigsDtD))
#Matching update
u0 = wl+dwl - beta/nu
u00 = u0[:,:,0]
u01 = u0[:,:,1]
m0 = m[:,:,0]
m1 = m[:,:,1]
idx = [c == 0]
rho = u0-m[:,:,:2]
# idx1 = logical((rho(:,:,1) < - thresh2) .* (c == 1))
# idx2 = logical((rho(:,:,1) > thresh2) .* (c == 1))
# idx3 = logical((abs(rho(:,:,1)) <= thresh2) .* (c == 1))
#
# uu1(idx1) = u01(idx1) + thresh2(idx1)
# uu1(idx2) = u01(idx2) - thresh2(idx2)
# uu1(idx3) = m1(idx3)
# uu1(idx) = u01(idx)
# #uu1(idocc) = z01(idocc)
#
# idx1 = logical((rho(:,:,2) < - thresh2) .* (c == 1))
# idx2 = logical((rho(:,:,2) > thresh2) .* (c == 1))
# idx3 = logical((abs(rho(:,:,2)) <= thresh2) .* (c == 1))
#
# uu2(idx1) = u02(idx1) + thresh2(idx1)
# uu2(idx2) = u02(idx2) - thresh2(idx2)
# uu2(idx3) = m2(idx3)
# uu2(idx) = u02(idx)
# #uu2(idocc) = z02(idocc)
#
# u = cat(3,uu1,uu2)
#
# #Lagrange parameters update
return wl
c4science · Help