--- /dev/null Thu Jan 01 00:00:00 1970 +0000
+++ b/src/OpticalFlow.jl Tue Apr 07 14:19:48 2020 -0500
@@ -0,0 +1,280 @@
+################################
+# Code relevant to optical flow
+################################
+
+__precompile__()
+
+module OpticalFlow
+
+using AlgTools.Util
+using ImageTools.Gradient
+import ImageTools.Translate
+using ImageTools.ImFilter
+
+##########
+# Exports
+##########
+
+export flow!,
+ pdflow!,
+ flow_grad!,
+ flow_interp!,
+ estimate_Λ²,
+ estimate_linear_Λ²,
+ pointwise_gradiprod_2d!,
+ pointwise_gradiprod_2dᵀ!,
+ horn_schunck_reg_prox!,
+ horn_schunck_reg_prox_op!,
+ mldivide_step_plus_sym2x2!,
+ linearised_optical_flow_error,
+ Image, AbstractImage, ImageSize,
+ Gradient, Displacement,
+ DisplacementFull, DisplacementConstant,
+ HornSchunckData,
+ filter_hs
+
+###############################################
+# Types (several imported from ImageTools.Translate)
+###############################################
+
+Image = Translate.Image
+AbstractImage = AbstractArray{Float64,2}
+Displacement = Translate.Displacement
+DisplacementFull = Translate.DisplacementFull
+DisplacementConstant = Translate.DisplacementConstant
+Gradient = Array{Float64,3}
+ImageSize = Tuple{Int64,Int64}
+
+#################################
+# Displacement field based flow
+#################################
+
+function flow_interp!(x::AbstractImage, u::Displacement, tmp::AbstractImage;
+ threads = false)
+ tmp .= x
+ Translate.translate_image!(x, tmp, u; threads=threads)
+end
+
+function flow_interp!(x::AbstractImage, u::Displacement;
+ threads = false)
+ tmp = copy(x)
+ Translate.translate_image!(x, tmp, u; threads=threads)
+end
+
+flow! = flow_interp!
+
+function pdflow!(x, Δx, y, Δy, u, dual_flow; threads=:none)
+ if dual_flow
+ #flow!((x, @view(y[1, :, :]), @view(y[2, :, :])), diffu,
+ # (Δx, @view(Δy[1, :, :]), @view(Δy[2, :, :])))
+ @backgroundif (threads==:outer) begin
+ flow!(x, u, Δx; threads=(threads==:inner))
+ end begin
+ flow!(@view(y[1, :, :]), u, @view(Δy[1, :, :]); threads=(threads==:inner))
+ flow!(@view(y[2, :, :]), u, @view(Δy[2, :, :]); threads=(threads==:inner))
+ end
+ else
+ flow!(x, u, Δx)
+ end
+end
+
+function pdflow!(x, Δx, y, Δy, z, Δz, u, dual_flow; threads=:none)
+ if dual_flow
+ @backgroundif (threads==:outer) begin
+ flow!(x, u, Δx; threads=(threads==:inner))
+ flow!(z, u, Δz; threads=(threads==:inner))
+ end begin
+ flow!(@view(y[1, :, :]), u, @view(Δy[1, :, :]); threads=(threads==:inner))
+ flow!(@view(y[2, :, :]), u, @view(Δy[2, :, :]); threads=(threads==:inner))
+ end
+ else
+ flow!(x, u, Δx; threads=(threads==:inner))
+ flow!(z, u, Δz; threads=(threads==:inner))
+ end
+end
+
+##########################
+# Linearised optical flow
+##########################
+
+# ⟨⟨u, ∇b⟩⟩
+function pointwise_gradiprod_2d!(y::Image, vtmp::Gradient,
+ u::DisplacementFull, b::Image;
+ add = false)
+ ∇₂c!(vtmp, b)
+
+ u′=reshape(u, (size(u, 1), prod(size(u)[2:end])))
+ vtmp′=reshape(vtmp, (size(vtmp, 1), prod(size(vtmp)[2:end])))
+ y′=reshape(y, prod(size(y)))
+
+ if add
+ @simd for i = 1:length(y′)
+ @inbounds y′[i] += dot(@view(u′[:, i]), @view(vtmp′[:, i]))
+ end
+ else
+ @simd for i = 1:length(y′)
+ @inbounds y′[i] = dot(@view(u′[:, i]), @view(vtmp′[:, i]))
+ end
+ end
+end
+
+function pointwise_gradiprod_2d!(y::Image, vtmp::Gradient,
+ u::DisplacementConstant, b::Image;
+ add = false)
+ ∇₂c!(vtmp, b)
+
+ vtmp′=reshape(vtmp, (size(vtmp, 1), prod(size(vtmp)[2:end])))
+ y′=reshape(y, prod(size(y)))
+
+ if add
+ @simd for i = 1:length(y′)
+ @inbounds y′[i] += dot(u, @view(vtmp′[:, i]))
+ end
+ else
+ @simd for i = 1:length(y′)
+ @inbounds y′[i] = dot(u, @view(vtmp′[:, i]))
+ end
+ end
+end
+
+# ∇b ⋅ y
+function pointwise_gradiprod_2dᵀ!(u::DisplacementFull, y::Image, b::Image)
+ ∇₂c!(u, b)
+
+ u′=reshape(u, (size(u, 1), prod(size(u)[2:end])))
+ y′=reshape(y, prod(size(y)))
+
+ @simd for i=1:length(y′)
+ @inbounds @. u′[:, i] *= y′[i]
+ end
+end
+
+function pointwise_gradiprod_2dᵀ!(u::DisplacementConstant, y::Image, b::Image)
+ @assert(size(y)==size(b) && size(u)==(2,))
+ u .= 0
+ ∇₂cfold!(b, nothing) do g, st, (i, j)
+ @inbounds u .+= g.*y[i, j]
+ return st
+ end
+ # Reweight to be with respect to 𝟙^*𝟙 inner product.
+ u ./= prod(size(b))
+end
+
+mutable struct ConstantDisplacementHornSchunckData
+ M₀::Array{Float64,2}
+ z::Array{Float64,1}
+ Mv::Array{Float64,2}
+ av::Array{Float64,1}
+ cv::Float64
+
+ function ConstantDisplacementHornSchunckData()
+ return new(zeros(2, 2), zeros(2), zeros(2,2), zeros(2), 0)
+ end
+end
+
+# For DisplacementConstant, for the simple prox step
+#
+# (1) argmin_u 1/(2τ)|u-ũ|^2 + (θ/2)|b⁺-b+<<u-ŭ,∇b>>|^2 + (λ/2)|u-ŭ|^2,
+#
+# construct matrix M₀ and vector z such that we can solve u from
+#
+# (2) (I/τ+M₀)u = M₀ŭ + ũ/τ - z
+#
+# Note that the problem
+#
+# argmin_u 1/(2τ)|u-ũ|^2 + (θ/2)|b⁺-b+<<u-ŭ,∇b>>|^2 + (λ/2)|u-ŭ|^2
+# + (θ/2)|b⁺⁺-b⁺+<<uʹ-u,∇b⁺>>|^2 + (λ/2)|u-uʹ|^2
+#
+# has with respect to u the system
+#
+# (I/τ+M₀+M₀ʹ)u = M₀ŭ + M₀ʹuʹ + ũ/τ - z + zʹ,
+#
+# where the primed variables correspond to (2) for (1) for uʹ in place of u:
+#
+# argmin_uʹ 1/(2τ)|uʹ-ũʹ|^2 + (θ/2)|b⁺⁺-b⁺+<<uʹ-u,∇b⁺>>|^2 + (λ/2)|uʹ-u|^2
+#
+function horn_schunck_reg_prox_op!(hs::ConstantDisplacementHornSchunckData,
+ bnext::Image, b::Image, θ, λ, T)
+ @assert(size(b)==size(bnext))
+ w = prod(size(b))
+ z = hs.z
+ cv = 0
+ # Factors of symmetric matrix [a c; c d]
+ a, c, d = 0.0, 0.0, 0.0
+ # This used to use ∇₂cfold but it is faster to allocate temporary
+ # storage for the full gradient due to probably better memory and SIMD
+ # instruction usage.
+ g = zeros(2, size(b)...)
+ ∇₂c!(g, b)
+ @inbounds for i=1:size(b, 1)
+ for j=1:size(b, 2)
+ δ = bnext[i,j]-b[i,j]
+ @. z += g[:,i,j]*δ
+ cv += δ*δ
+ a += g[1,i,j]*g[1,i,j]
+ c += g[1,i,j]*g[2,i,j]
+ d += g[2,i,j]*g[2,i,j]
+ end
+ end
+ w₀ = λ
+ w₂ = θ/w
+ aʹ = w₀ + w₂*a
+ cʹ = w₂*c
+ dʹ = w₀ + w₂*d
+ hs.M₀ .= [aʹ cʹ; cʹ dʹ]
+ hs.Mv .= [w*λ+θ*a θ*c; θ*c w*λ+θ*d]
+ hs.cv = cv*θ
+ hs.av .= hs.z.*θ
+ hs.z .*= w₂/T
+end
+
+# Solve the 2D system (I/τ+M₀)u = z
+@inline function mldivide_step_plus_sym2x2!(u, M₀, z, τ)
+ a = 1/τ+M₀[1, 1]
+ c = M₀[1, 2]
+ d = 1/τ+M₀[2, 2]
+ u .= ([d -c; -c a]*z)./(a*d-c*c)
+end
+
+function horn_schunck_reg_prox!(u::DisplacementConstant, bnext::Image, b::Image,
+ θ, λ, T, τ)
+ hs=ConstantDisplacementHornSchunckData()
+ horn_schunck_reg_prox_op!(hs, bnext, b, θ, λ, T)
+ mldivide_step_plus_sym2x2!(u, hs.M₀, (u./τ)-hs.z, τ)
+end
+
+function flow_grad!(x::Image, vtmp::Gradient, u::Displacement; δ=nothing)
+ if !isnothing(δ)
+ u = δ.*u
+ end
+ pointwise_gradiprod_2d!(x, vtmp, u, x; add=true)
+end
+
+# Error b-b_prev+⟨⟨u, ∇b⟩⟩ for Horn–Schunck type penalisation
+function linearised_optical_flow_error(u::Displacement, b::Image, b_prev::Image)
+ imdim = size(b)
+ vtmp = zeros(2, imdim...)
+ tmp = b-b_prev
+ pointwise_gradiprod_2d!(tmp, vtmp, u, b_prev; add=true)
+ return tmp
+end
+
+##############################################
+# Helper to smooth data for Horn–Schunck term
+##############################################
+
+function filter_hs(b, b_next, b_next_filt, kernel)
+ if kernel==nothing
+ f = x -> x
+ else
+ f = x -> simple_imfilter(x, kernel; threads=true)
+ end
+
+ # We already filtered b in the previous step (b_next in that step)
+ b_filt = b_next_filt==nothing ? f(b) : b_next_filt
+ b_next_filt = f(b_next)
+
+ return b_filt, b_next_filt
+end
+
+end # Module