diff -r 4bbb246f4cac -r e4ad8f7ce671 src/PET/OpticalFlow.jl
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/src/PET/OpticalFlow.jl	Fri Apr 19 17:00:37 2024 +0300
@@ -0,0 +1,431 @@
+################################
+# Code relevant to optical flow
+################################
+
+__precompile__()
+
+module OpticalFlow
+
+using AlgTools.Util
+using ImageTools.Gradient
+import ImageTools.Translate
+using ImageTools.ImFilter
+
+# using ImageTransformations
+# using Images, CoordinateTransformations, Rotations, OffsetArrays
+# using Interpolations
+
+import Images: center, warp
+import CoordinateTransformations: recenter
+import Rotations: RotMatrix
+import Interpolations: Flat
+
+##########
+# 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,
+       petpdflow!
+
+###############################################
+# 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
+
+# Additional method for Greedy
+function pdflow!(x, Δx, y, Δy, u; threads=:none)
+    @assert(size(u)==(2,))  
+    Δx .= x
+    Δy .= y
+    flow!(x, u; threads=(threads==:inner))
+    Dxx = similar(Δy)
+    DΔx = similar(Δy)
+    ∇₂!(Dxx, x)
+    ∇₂!(DΔx, Δx)
+    inds = abs.(Dxx) .≤ 1e-1
+    Dxx[inds] .= 1
+    DΔx[inds] .= 1
+    y .= y.* DΔx ./ Dxx  
+end
+
+# Additional method for Rotation
+function pdflow!(x, Δx, y, u; threads=:none) 
+    @assert(size(u)==(2,))
+    Δx .= x
+    flow!(x, u; threads=(threads==:inner))
+
+    (m,n) = size(x)
+    dx = similar(y)
+    dx_banana = similar(y)
+    ∇₂!(dx, Δx)
+    ∇₂!(dx_banana, x)
+
+    for i=1:m
+        for j=1:n
+            ndx = @views sum(dx[:, i, j].^2)
+            ndx_banana = @views sum(dx_banana[:, i, j].^2)
+            if ndx > 1e-4 && ndx_banana > 1e-4
+                A = dx[:, i, j]
+                B = dx_banana[:, i, j]
+                theta = atan(B[1] * A[2] - B[2] * A[1], B[1] * A[1] + B[2] * A[2]) # Oriented angle from A to B
+                cos_theta = cos(theta)
+                sin_theta = sin(theta)
+                a = cos_theta * y[1, i, j] - sin_theta * y[2, i, j]
+                b = sin_theta * y[1, i, j] + cos_theta * y[2, i, j]
+                y[1, i, j] = a
+                y[2, i, j] = b
+            end
+        end
+    end
+end
+
+# Additional method for Dual Scaling
+function pdflow!(x, y, u; threads=:none)
+    @assert(size(u)==(2,))
+    oldx = copy(x)
+    flow!(x, u; threads=(threads==:inner))
+    C = similar(y)
+    cc = abs.(x-oldx)
+    cm = max(1e-12,maximum(cc))
+    c = 1 .* (1 .- cc./ cm) .^(10)
+    C[1,:,:] .= c
+    C[2,:,:] .= c
+    y .= C.*y 
+end
+
+
+##########################
+# PET
+##########################
+function petflow_interp!(x::AbstractImage, tmp::AbstractImage, u::DisplacementConstant, theta_known::DisplacementConstant;
+    threads = false)
+    tmp .= x
+    center_point = center(x) .+ u
+    tform = recenter(RotMatrix(theta_known[1]), center_point)
+    tmp = warp(x, tform, axes(x), fillvalue=Flat())
+    x .= tmp
+end
+
+petflow! = petflow_interp!
+
+function petpdflow!(x, Δx, y, Δy, u, theta_known, dual_flow; threads=:none)
+    if dual_flow
+        @backgroundif (threads==:outer) begin
+        petflow!(x, Δx, u, theta_known; threads=(threads==:inner))
+    end begin
+        petflow!(@view(y[1, :, :]), @view(Δy[1, :, :]), u, theta_known; threads=(threads==:inner))
+        petflow!(@view(y[2, :, :]), @view(Δy[2, :, :]), u, theta_known; threads=(threads==:inner))
+        end
+    else
+    petflow!(x, Δx, u, theta_known)
+    end
+end
+
+# Method for greedy predictor
+function petpdflow!(x, Δx, y, Δy, u, theta_known, dual_flow, β; threads=:none)
+    oldx = copy(x)
+    center_point = center(x) .+ u
+    tform = recenter(RotMatrix(theta_known[1]), center_point)
+    Δx = warp(x, tform, axes(x), fillvalue=Flat())
+    @. x = Δx
+    @. Δy = y
+    if dual_flow
+        Dxx = copy(Δy)
+        DΔx = copy(Δy)
+        ∇₂!(Dxx, x)
+        ∇₂!(DΔx, oldx)
+        inds = abs.(Dxx) .≤ β
+        Dxx[inds] .= 1
+        DΔx[inds] .= 1
+        y .= y.* DΔx ./ Dxx            
+    end
+end
+
+# Method for affine predictor
+function petpdflow!(x, Δx, y, u, theta_known, dual_flow; threads=:none)
+    oldx = copy(x)
+    center_point = center(x) .+ u
+    tform = recenter(RotMatrix(theta_known[1]), center_point)
+    Δx = warp(x, tform, axes(x), fillvalue=Flat())
+    @. x = Δx
+    C = similar(y)
+    cc = abs.(x-oldx)
+    if dual_flow
+        cm = max(1e-12,maximum(cc))
+        c = 1 .* (1 .- cc./ cm) .^(10)
+        C[1,:,:] .= c
+        C[2,:,:] .= c
+        y .= C.*y 
+    end
+end
+
+# Method for rotation prediction (exploiting property of inverse rotation)
+function petpdflow!(x, Δx, y, Δy, u, theta_known, dual_flow, β₁, β₂; threads=:none)
+    if dual_flow
+        @backgroundif (threads==:outer) begin
+            petflow!(x, Δx, u, theta_known; threads=(threads==:inner))
+        end begin
+            petflow!(@view(y[1, :, :]), @view(Δy[1, :, :]), u, -theta_known; threads=(threads==:inner))
+            petflow!(@view(y[2, :, :]), @view(Δy[2, :, :]), u, -theta_known; threads=(threads==:inner))
+        end
+    else
+        petflow!(x, Δx, u, theta_known)
+    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