PredictPDPS: comparison src/AlgorithmBothNL.jl

--1:000000000000
+:a55e35d20336
+######################################################################
+# Predictive online PDPS for optical flow with unknown velocity field
+######################################################################
+__precompile__()
+module AlgorithmBothNL
+identifier = "pdps_unknown_nl"
+using Printf
+using AlgTools.Util
+import AlgTools.Iterate
+using ImageTools.Gradient
+using ..OpticalFlow: ImageSize,
+Image,
+Gradient,
+DisplacementConstant,
+DisplacementFull,
+pdflow!,
+pointwise_gradiprod_2d!,
+pointwise_gradiprod_2dᵀ!,
+filter_hs
+using ..Algorithm: step_lengths
+#############
+# Data types
+#############
+struct Primal{DisplacementT}
+x :: Image
+u :: DisplacementT
+end
+function Base.similar(x::Primal{DisplacementT}) where DisplacementT
+return Primal{DisplacementT}(Base.similar(x.x), Base.similar(x.u))
+end
+function Base.copy(x::Primal{DisplacementT}) where DisplacementT
+return Primal{DisplacementT}(Base.copy(x.x), Base.copy(x.u))
+end
+struct Dual
+tv :: Gradient
+flow :: Image
+end
+function Base.similar(y::Dual)
+return Dual(Base.similar(y.tv), Base.similar(y.flow))
+end
+function Base.copy(y::Dual)
+return Dual(Base.copy(y.tv), Base.copy(y.flow))
+end
+#########################
+# Iterate initialisation
+#########################
+function init_primal(xinit::Image, ::Type{DisplacementConstant})
+return Primal{DisplacementConstant}(xinit, zeros(2))
+end
+function init_primal(xinit::Image, ::Type{DisplacementFull})
+return Primal{DisplacementFull}(xinit, zeros(2, size(xinit)...))
+end
+function init_rest(x::Primal{DisplacementT}) where DisplacementT
+imdim=size(x.x)
+y = Dual(zeros(2, imdim...), zeros(imdim))
+Δx = copy(x)
+Δy = copy(y)
+x̄ = copy(x)
+return x, y, Δx, Δy, x̄
+end
+function init_iterates( :: Type{DisplacementT},
+xinit::Primal{DisplacementT}) where DisplacementT
+return init_rest(copy(xinit))
+end
+function init_iterates( :: Type{DisplacementT}, xinit::Image) where DisplacementT
+return init_rest(init_primal(copy(xinit), DisplacementT))
+end
+function init_iterates( :: Type{DisplacementT}, dim::ImageSize) where DisplacementT
+return init_rest(init_primal(zeros(dim...), DisplacementT))
+end
+##############################################
+# Weighting for different displacements types
+##############################################
+norm²weight( :: Type{DisplacementConstant}, sz ) = prod(sz)
+norm²weight( :: Type{DisplacementFull}, sz ) = 1
+############
+# Algorithm
+############
+function solve( :: Type{DisplacementT};
+dim :: ImageSize,
+iterate = AlgTools.simple_iterate,
+params::NamedTuple) where DisplacementT
+######################
+# Initialise iterates
+######################
+x, y, Δx, Δy, x̄ = init_iterates(DisplacementT, dim)
+init_data = (params.init == :data)
+# … for tracking cumulative movement
+if DisplacementT == DisplacementConstant
+ucumul = [0.0, 0.0]
+else
+ucumul = [NaN, NaN]
+end
+#############################################
+# Extract parameters and set up step lengths
+#############################################
+α, ρ, λ, θ, T = params.α, params.ρ, params.λ, params.θ, params.timestep
+R_K² = max(∇₂_norm₂₂_est², ∇₂_norm₂∞_est²*params.dynrange^2)
+γ = min(1, λ*norm²weight(DisplacementT, size(x.x)))
+τ, σ, σ̃, ρ̃ = step_lengths(params, γ, R_K²)
+kernel = params.kernel
+####################
+# Run the algorithm
+####################
+b_next_filt=nothing
+v = iterate(params) do verbose :: Function,
+b :: Image,
+🚫unused_v_known :: DisplacementT,
+b_next :: Image
+####################################
+# Smooth data for Horn–Schunck term
+####################################
+b_filt, b_next_filt = filter_hs(b, b_next, b_next_filt, kernel)
+############################
+# Construct K for this step
+############################
+K! = (yʹ, xʹ) -> begin
+# Optical flow part
+@. yʹ.flow = b_filt
+flow!(yʹ.flow, Δx.x, xʹ.u)
+@. yʹ.flow = yʹ.flow - b_next_filt
+# TV part
+∇₂!(yʹ.tv, xʹ.x)
+end
+Kᵀ! = (xʹ, yʹ) -> begin
+# Optical flow part: ∇b_k ⋅ y
+#
+# TODO: This really should depend x.u, but x.u is zero.
+#
+pointwise_gradiprod_2dᵀ!(xʹ.u, yʹ.flow, b_filt)
+# TV part
+∇₂ᵀ!(xʹ.x, yʹ.tv)
+end
+##################
+# Prediction step
+##################
+if init_data
+x .= b
+init_data = false
+end
+pdflow!(x.x, Δx.x, y.tv, Δy.tv, y.flow, Δy.flow, x.u, params.dual_flow)
+# Predict zero displacement
+x.u .= 0
+if params.prox_predict
+K!(Δy, x)
+@. y.tv = (y.tv + σ̃*Δy.tv)/(1 + σ̃*(ρ̃+ρ/α))
+proj_norm₂₁ball!(y.tv, α)
+@. y.flow = (y.flow+σ̃*Δy.flow)/(1+σ̃*(ρ̃+1/θ))
+end
+############
+# PDPS step
+#
+# NOTE: For DisplacementConstant, the x.u update is supposed to be with
+# respect to the 𝟙^*𝟙 norm/inner product that makes the norm equivalent
+# to full-space norm when restricted to constant displacements. Since
+# `OpticalFlow.pointwise_gradiprod_2dᵀ!` already uses this inner product,
+# and the λ-weighted term in the problem is with respect to this norm,
+# all the norm weights disappear in this update.
+############
+Kᵀ!(Δx, y)                             # primal step:
+@. x̄.x = x.x                           # |  save old x for over-relax
+@. x̄.u = x.u                           # |
+@. x.x = (x.x-τ*(Δx.x-b))/(1+τ)        # |  prox
+@. x.u = (x.u-τ*Δx.u)/(1+τ*λ)          # |
+@. x̄.x = 2x.x - x̄.x                    # over-relax
+@. x̄.u = 2x.u - x̄.u                    # |
+K!(Δy, x̄)                              # dual step: y
+@. y.tv = (y.tv + σ*Δy.tv)/(1 + σ*ρ/α) # |
+proj_norm₂₁ball!(y.tv, α)              # |  prox
+@. y.flow = (y.flow+σ*Δy.flow)/(1+σ/θ)
+if DisplacementT == DisplacementConstant
+ucumul .+= x.u
+end
+########################################################
+# Give function value and cumulative movement if needed
+########################################################
+v = verbose() do
+K!(Δy, x)
+value = (norm₂²(b-x.x)/2 + θ*norm₂²(Δy.flow)
++ λ*norm₂²(x.u)/2 + α*γnorm₂₁(Δy.tv, ρ))
+value, x.x, ucumul, nothing
+end
+return v
+end
+return x, y, v
+end
+end # Module

Mercurial > repos > PredictPDPS / file comparison

comparison: src/AlgorithmBothNL.jl

src/AlgorithmBothNL.jl