def of2dmcs(img1: np.array, img2: np.array, alpha0: float, alpha1: float,
beta0: float, beta1: float) -> (np.array, np.array):
"""Computes the L2-H1 optical flow for a 2D two-channel image sequence and
source for the second channel (optical flow 2d multi-channel with source).
Takes a two-dimensional image sequence and returns a minimiser of the
Horn-Schunck functional with source and with spatio-temporal
regularisation for both velocity and source.
Args:
img1 (np.array): A 2D image sequence of shape (t, m, n), where t is the
number of time steps and (n, n) is the number of
pixels.
img2 (np.array): A 2D image sequence of shape (t, m, n), where t is the
number of time steps and (n, n) is the number of
pixels.
alpha0 (float): The spatial regularisation parameter.
alpha1 (float): The temporal regularisation parameter.
Returns:
v (np.array): A velocity array of shape (t, m, n, 2).
k (np.array): A source array of shape (t, m, n).
"""
# Create mesh.
[t, m, n] = img1.shape
mesh = UnitCubeMesh(t-2, m-1, n-1)
# Define function space and functions.
V = FunctionSpace(mesh, 'CG', 1)
W = VectorFunctionSpace(mesh, 'CG', 1, dim=3)
v1, v2, k = TrialFunctions(W)
w1, w2, w3 = TestFunctions(W)
# Convert image to function.
f1, f2 = Function(V), Function(V)
f1.vector()[:] = dh.imgseq2funvec(img1[0:-1])
f2.vector()[:] = dh.imgseq2funvec(img2[0:-1])
# Define function to compute temporal derivative.
def time_deriv(img: np.array) -> Function:
# Evaluate function at vertices.
mc = mesh.coordinates().reshape((-1, 3))
hx, hy, hz = 1./(t-2), 1./(m-1), 1./(n-1)
x = np.array(mc[:, 0]/hx, dtype=int)
y = np.array(mc[:, 1]/hy, dtype=int)
z = np.array(mc[:, 2]/hz, dtype=int)
# Map pixel values to vertices.
d2v = dof_to_vertex_map(V)
# Compute derivative wrt. time.
imgt = img[1:] - img[0:-1]
ftv = imgt[x, y, z]
# Create function.
ft = Function(V)
ft.vector()[:] = ftv[d2v]
return ft
# Compute temporal derivatives.
f1t = time_deriv(img1)
f2t = time_deriv(img2)
# Define derivatives of data.
f1x, f1y = f1.dx(1), f1.dx(2)
f2x, f2y = f2.dx(1), f2.dx(2)
# Define weak formulation.
A = f1x*(f1x*v1 + f1y*v2)*w1*dx + f2x*(f2x*v1 + f2y*v2 - k)*w1*dx \
+ f1y*(f1x*v1 + f1y*v2)*w2*dx + f2y*(f2x*v1 + f2y*v2 - k)*w2*dx \
- (f2x*v1 + f2y*v2 - k)*w3*dx \
+ alpha0*v1.dx(1)*w1.dx(1)*dx + alpha0*v1.dx(2)*w1.dx(2)*dx \
+ alpha1*v1.dx(0)*w1.dx(0)*dx \
+ alpha0*v2.dx(1)*w2.dx(1)*dx + alpha0*v2.dx(2)*w2.dx(2)*dx \
+ alpha1*v2.dx(0)*w2.dx(0)*dx \
+ beta0*k.dx(1)*w3.dx(1)*dx + beta0*k.dx(2)*w3.dx(2)*dx \
+ beta1*k.dx(0)*w3.dx(0)*dx
b = - f1x*f1t*w1*dx - f2x*f2t*w1*dx \
- f1y*f1t*w2*dx - f2y*f2t*w2*dx \
+ f2t*w3*dx
# Compute solution.
v = Function(W)
solve(A == b, v, [], solver_parameters={"linear_solver": "cg"})
# Split solution into functions.
v1, v2, k = v.split(deepcopy=True)
# Convert back to arrays.
v1 = dh.funvec2imgseq(v1.vector().get_local(), t-1, m, n)
v2 = dh.funvec2imgseq(v2.vector().get_local(), t-1, m, n)
k = dh.funvec2imgseq(k.vector().get_local(), t-1, m, n)
return (np.stack((v1, v2), axis=3), k)
def cmscr1dnewton(img: np.array, alpha0: float, alpha1: float, alpha2: float,
alpha3: float, beta: float) \
-> (np.array, np.array, float, float, bool):
"""Same as cmscr1d but doesn't use FEniCS Newton method.
Args:
img (np.array): A 1D image sequence of shape (m, n), where m is the
number of time steps and n is the number of pixels.
alpha0 (float): The spatial regularisation parameter for v.
alpha1 (float): The temporal regularisation parameter for v.
alpha2 (float): The spatial regularisation parameter for k.
alpha3 (float): The temporal regularisation parameter for k.
beta (float): The parameter for convective regularisation.
Returns:
v: A velocity array of shape (m, n).
k: A source array of shape (m, n).
res (float): The residual.
fun (float): The function value.
converged (bool): True if Newton's method converged.
"""
# Create mesh.
[m, n] = img.shape
mesh = UnitSquareMesh(m - 1, n - 1)
# Define function space and functions.
W = VectorFunctionSpace(mesh, 'CG', 1, dim=2)
w = Function(W)
v, k = split(w)
w1, w2 = TestFunctions(W)
# Convert image to function.
V = FunctionSpace(mesh, 'CG', 1)
f = Function(V)
f.vector()[:] = dh.img2funvec(img)
# Define derivatives of data.
ft = Function(V)
ftv = np.diff(img, axis=0) * (m - 1)
ftv = np.concatenate((ftv, ftv[-1, :].reshape(1, n)), axis=0)
ft.vector()[:] = dh.img2funvec(ftv)
fx = Function(V)
fxv = np.gradient(img, 1 / (n - 1), axis=1)
fx.vector()[:] = dh.img2funvec(fxv)
# Define weak formulation.
F = - (ft + fx*v + f*v.dx(1) - k) * (fx*w1 + f*w1.dx(1))*dx \
- alpha0*v.dx(1)*w1.dx(1)*dx - alpha1*v.dx(0)*w1.dx(0)*dx \
- beta*k.dx(1)*(k.dx(0) + k.dx(1)*v)*w1*dx \
+ (ft + fx*v + f*v.dx(1) - k)*w2*dx \
- alpha2*k.dx(1)*w2.dx(1)*dx - alpha3*k.dx(0)*w2.dx(0)*dx \
- beta*(k.dx(0)*w2.dx(0) + k.dx(1)*v*w2.dx(0)
+ k.dx(0)*v*w2.dx(1) + k.dx(1)*v*v*w2.dx(1))*dx
# Define derivative.
dv, dk = TrialFunctions(W)
J = - (fx*dv + f*dv.dx(1) - dk)*(fx*w1 + f*w1.dx(1))*dx \
- alpha0*dv.dx(1)*w1.dx(1)*dx - alpha1*dv.dx(0)*w1.dx(0)*dx \
- beta*(k.dx(1)*(dk.dx(0) + dk.dx(1)*v) + k.dx(1)**2*dv
+ dk.dx(1)*(k.dx(0) + k.dx(1)*v))*w1*dx \
+ (fx*dv + f*dv.dx(1) - dk)*w2*dx \
- alpha2*dk.dx(1)*w2.dx(1)*dx - alpha3*dk.dx(0)*w2.dx(0)*dx \
- beta*(dv*k.dx(1) + dk.dx(0) + v*dk.dx(1))*w2.dx(0)*dx \
- beta*(dv*k.dx(0) + 2*v*dv*k.dx(1) + v*dk.dx(0)
+ v*v*dk.dx(1))*w2.dx(1)*dx
# Alternatively, use automatic differentiation.
# J = derivative(F, w)
# Define algorithm parameters.
tol = 1e-10
maxiter = 100
# Define increment.
dw = Function(W)
# Run Newton iteration.
niter = 0
eps = 1
res = 1
while res > tol and niter < maxiter:
niter += 1
# Solve for increment.
solve(J == -F, dw)
# Update solution.
w.assign(w + dw)
# Compute norm of increment.
eps = dw.vector().norm('l2')
# Compute norm of residual.
a = assemble(F)
res = a.norm('l2')
# Print statistics.
print("Iteration {0} eps={1} res={2}".format(niter, eps, res))
# Split solution.
v, k = w.split(deepcopy=True)
# Evaluate and print residual and functional value.
res = abs(ft + fx * v + f * v.dx(1) - k)
fun = 0.5 * (res**2 + alpha0 * v.dx(1)**2 + alpha1 * v.dx(0)**2 +
alpha2 * k.dx(1)**2 + alpha3 * k.dx(0)**2 + beta *
(k.dx(0) + k.dx(1) * v)**2)
print('Res={0}, Func={1}'.format(assemble(res * dx), assemble(fun * dx)))
# Convert back to array.
vel = dh.funvec2img(v.vector().get_local(), m, n)
k = dh.funvec2img(k.vector().get_local(), m, n)
return vel, k, assemble(res * dx), assemble(fun * dx), res > tol