PredictPDPS: comparison src/OpticalFlow.jl

--1:000000000000
+:a55e35d20336
+################################
+# 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-ũ|^2 + (θ/2)|b⁺-b+<<u-ŭ,∇b>>|^2 + (λ/2)|u-ŭ|^2,
+#
+# construct matrix M₀ and vector z such that we can solve u from
+#
+# (2) (I/τ+M₀)u = M₀ŭ + ũ/τ - z
+#
+# Note that the problem
+#
+#    argmin_u 1/(2τ)|u-ũ|^2 + (θ/2)|b⁺-b+<<u-ŭ,∇b>>|^2 + (λ/2)|u-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₀ŭ + M₀ʹuʹ + ũ/τ - z + zʹ,
+#
+# where the primed variables correspond to (2) for (1) for uʹ in place of u:
+#
+#    argmin_uʹ 1/(2τ)|uʹ-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

Mercurial > repos > PredictPDPS / file comparison

comparison: src/OpticalFlow.jl

src/OpticalFlow.jl