Homec4science
Diffusion Motion correction paper (workstation) Browse

R6289/code/of_l1_l2_fm_admm.py
ff9f520a76c9workstation

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
import matlab.engine
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
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
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
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
eng = matlab.engine.start_matlab()
eigsDtD = eng.eigsDtD(int(I1.shape[0]), int(I1.shape[1]), float(lmbd), float(mu))
for iWarp in range(param['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
[matlabIt, matlabIx, matlabIy] = eng.py_partial_deriv(matlab.double(I1.tolist()),matlab.double(I2.tolist()),matlab.double(wl.tolist()),matlab.double(deriv_filter.tolist()), nargout=3)#TODO figure out how to do this
It = np.asarray(matlabIt)
Ix = np.asarray(matlabIx)
Iy = np.asarray(matlabIy)
Igrad = Ix**2 + Iy**2 + 1e-3
idocc = (occ==0);
#Main iterations loop
for it in range(param['maxIters']):
thresh = Igrad/(mu+nu)
thresh2 = param['gamma']*m[:,:,2]/nu
#Data update
r1 = z - wl - alpha/mu
r2 = u - wl + beta/nu
t = (mu*r1+nu*r2)/(mu+nu)
t1 = t[:,:,0]
t2 = t[:,:,1]
rho = It + t1*Ix + t2*Iy
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]
dvl[idx3] = t2[idx3] - rho[idx3]*Iy[idx3]/Igrad[idx3]
dul[idocc] = t1[idocc]
dvl[idocc] = t2[idocc]
dwl = np.dstack((dul,dvl))
#Regularization update
#TODO translate
#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
u01 = u0[:,:,0]
u02 = u0[:,:,1]
m1 = m[:,:,0]
m2 = m[:,:,1]
idx = [c == 0]
rho = u0-m[:,:,:2]
idx1 = (rho[:,:,0] < - thresh2) * (c == 1)
idx2 = (rho[:,:,0] > thresh2) * (c == 1)
idx3 = (abs(rho[:,:,0]) <= 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 = (rho[:,:,1] < - thresh2) * (c == 1)
idx2 = (rho[:,:,1] > thresh2) * (c == 1)
idx3 = (abs(rho[:,:,1]) <= 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 = np.dstack((uu1,uu2))
#Lagrange parameters update
alpha = alpha + mu*(wl+dwl - z)
beta= beta + nu*(u-wl-dwl)
#Post-processing
if (it % param['iWM']) == 0:
w0 = wl+dwl
#[w0,occTmp] = post_process(w0,I1,I2,sigmaS,param['sigmaC'])
matlab.double(w0.tolist())
matlab.double(I1.tolist())
matlab.double(I2.tolist())
matlab.double([sigmaS])
print(param['sigmaC'])
matlab.double([param['sigmaC']])
[matlabw0,matlaboccTmp] = eng.post_process(matlab.double(w0.tolist()),matlab.double(I1.tolist()),matlab.double(I2.tolist()),matlab.double([sigmaS]),matlab.double([param['sigmaC']]),nargout=2)
w0 = np.asarray(matlabw0)
occTmp = np.asarray(matlaboccTmp)
dwl = w0-wl
#End of iterations checking
w = wl+dwl
change = np.linalg.norm(w.flatten()-wPrev.flatten())/np.linalg.norm(wPrev.flatten())
if change<param['changeTol']:
break
wPrev = w
wl = wl + dwl;
#[wl,occ] = post_process(wl,I1,I2,sigmaS,param.sigmaC);
[matlabw0,matlabocc] = eng.post_process(matlab.double(w0.tolist()),matlab.double(I1.tolist()),matlab.double(I2.tolist()),matlab.double([sigmaS]),matlab.double([param['sigmaC']]),nargout=2)
w0 = np.asarray(matlabw0)
occ = np.asarray(matlabocc)
occ = occ > param['occThresh']
#TODO chage displaying to python
#if disp:
# figure(11);imagesc(crop_fit_size_center(flowToColor(wl),[sz0(1),sz0(2),3]));colormap gray; axis image; drawnow;
return wl
c4science · Help