add ROF ADMM solver

src/models.jl

@@ -2,21 +2,26 @@ module Models
 
 using ImageBase
 using ImageBase.ImageCore.MappedArrays: of_eltype
-using ImageBase.FiniteDiff
+using ImageBase.FiniteDiff: fdiv!, fdiff!, fgradient!
 
 # Introduced in ColorVectorSpace v0.9.3
 # https://github.com/JuliaGraphics/ColorVectorSpace.jl/pull/172
 using ImageBase.ImageCore.ColorVectorSpace.Future: abs2
 
+using ..ImageFiltering: ffteltype, freqkernel
+using FFTW
+
 """
 This submodule provides predefined image-related models and its solvers that can be reused
 by many image processing tasks.
 
 - solve the Rudin Osher Fatemi (ROF) model using the primal-dual method: [`solve_ROF_PD`](@ref) and [`solve_ROF_PD!`](@ref)
+- solve the Rudin Osher Fatemi (ROF) model using the ADMM method: [`solve_ROF_ADMM`](@ref) and [`solve_ROF_ADMM!`](@ref)
 """
 Models
 
-export solve_ROF_PD, solve_ROF_PD!
+export solve_ROF_PD, solve_ROF_PD!,
+    solve_ROF_ADMM, solve_ROF_ADMM!
 
 
 ##### implementation details
@@ -169,5 +174,129 @@ function _l2norm_vec!(out, Vs::Tuple)
     return out
 end
 
+"""
+    solve_ROF_ADMM([T], img::AbstractArray, λ; kwargs...)
+
+Return a smoothed version of `img`, using Rudin-Osher-Fatemi (ROF) filtering, more commonly
+known as Total Variation (TV) denoising or TV regularization. This algorithm is based on the
+alternating direction method of multipliers (ADMM) method; it is also called the split
+bregman method.
+
+See also [`solve_ROF_PD`](@ref) for the primal dual method.
+
+# References
+
+- [1] Goldstein, Tom, and Stanley Osher. "The split Bregman method for L1-regularized problems." _SIAM journal on imaging sciences_ 2.2 (2009): 323-343.
+- [2] Getreuer, Pascal. "Rudin-Osher-Fatemi total variation denoising using split Bregman." _Image Processing On Line 2_ (2012): 74-95.
+"""
+solve_ROF_ADMM(img::AbstractArray{T}, args...; kwargs...) where T = solve_ROF_ADMM(float32(T), img, args...; kwargs...)
+
+function solve_ROF_ADMM(::Type{T}, img, args...; kwargs...) where T
+    # TODO(johnnychen94): support generic images
+    out = similar(img, T)
+    buffer = preallocate_solve_ROF_ADMM(T, img)
+    solve_ROF_ADMM!(out, buffer, img, args...; kwargs...)
+end
+
+# non-exported helper
+preallocate_solve_ROF_ADMM(img::AbstractArray{T}) where T = preallocate_solve_ROF_ADMM(float32(T), img)
+function preallocate_solve_ROF_ADMM(::Type{T}, img) where T
+    # Use similar to allow construct CuArray when `img` is a CuArray
+
+    # split variable for ∇
+    d = ntuple(i->fill!(similar(img, T), zero(T)), ndims(img))
+    # dual variable
+    b = map(copy, d)
+
+    # precomputed the negative Laplacian kernel in frequency space
+    DTD = _negative_flaplacian_freqkernel(img)
+
+    # buffer for u subproblem
+    #   minᵤ 0.5μ||u-f||₂² + 0.5∑ᵢ(λ||dᵢ - ∇ᵢu - bᵢ||₂²)
+    # For generality to arbitrary dimension, fft-based solver is applied
+    #   u = real.(ifft(fft(RHS)./fft(LHS)))
+    # where LHS is actually a convolution operator thus we can apply convolution theorem.
+    Δbd = similar(first(b))
+    LHS = similar(first(d)) # left-hand side: μI + Δ
+    RHS = similar(LHS) # right-hand side: μf + λ∑ᵢ(∇ᵢ'(dᵢ - bᵢ))
+    fft_tmp = similar(LHS, Complex{eltype(T)})
+
+    return d, b, DTD, Δbd, LHS, RHS, fft_tmp
+end
+
+"""
+    solve_ROF_ADMM!(out, buffer, img, λ, μ; kwargs...)
+
+The in-place version of [`solve_ROF_ADMM`](@ref).
+"""
+function solve_ROF_ADMM!(out::AbstractArray{T}, buffer, img, λ, μ; kwargs...) where T
+    # seperate a stub method to reduce latency
+    FT = float32(T)
+    if FT == T
+        _solve_ROF_ADMM_anisotropic!(out, buffer, img, Float32(λ), Float32(μ); kwargs...)
+    else
+        _solve_ROF_ADMM_anisotropic!(out, buffer, FT.(img), Float32(λ), Float32(μ); kwargs...)
+    end
+end
+
+function _solve_ROF_ADMM_anisotropic!(
+        out::AbstractArray,
+        (d, b, DTD, Δbd, LHS, RHS, fft_tmp)::Tuple,
+        img::AbstractArray,
+        λ::Float32, μ::Float32; num_iters::Integer=200)
+    # The notation and algorithm follows reference [1] with one modification:
+    u, f = out, img
+
+    # u-subproblem is solved using direct fft-based method instead of using Gauss–Seidel
+    # method. This is because Gauss–Seidel method itself is a solver for 2-dimensional case,
+    # while FFT-based solver support arbitrary dimensional. However, we might still want to
+    # provide a specialization solver for the 2D case using Gauss-Seidel method and see if
+    # it helps speaking of performance.
+
+    # soft-thresholding: d-subproblem is a simplified lasso problem
+    S(x, λ) = sign(x) * max(abs(x)-λ, zero(x))
+
+    # initialization
+    foreach(x->fill!(x, zero(eltype(x))), d)
+    foreach(x->fill!(x, zero(eltype(x))), b)
+    @. LHS = λ + μ * DTD
+    # apply the algorithm below Eq. (4.2)
+    for _ in 1:num_iters
+        @. RHS = λ/μ * f
+        for i in 1:length(d)
+            # negative adjoint gradient is the reverse finite difference
+            @. Δbd = b[i] - d[i]
+            # reuse d here as buffer as it will be updated later in d-subproblem
+            RHS .+= fdiff!(d[i], Δbd; dims=i, rev=true)
+        end
+        # TODO(johnnychen94): optimize the memory allocation
+        fft_tmp .= μ .* fft(RHS)./LHS
+        u .= real.(ifft(fft_tmp))
+
+        # for anisotropic problem, the update of d[i] is decoupled
+        du = RHS # reuse RHS as buffer for d-subproblem
+        for i in 1:length(d)
+            fdiff!(du, u; dims=i)
+            # 2) solve d subproblem
+            @. d[i] = S(du + b[i], 1/μ)
+            # 3) update dual variable b
+            @. b[i] = b[i] + du - d[i]
+        end
+    end
+    return u
+end
+
+# generates -fft(Δ) = -∑fft(∇ᵢ')(∇ᵢ) of size `size(X)` in the frequency domain
+function _negative_flaplacian_freqkernel(X::AbstractArray{T}) where T<:Real
+    FΔ = fill!(similar(X), zero(ffteltype(T)))
+    for i in 1:ndims(X)
+        # forward diff
+        ∇ᵢ = reshape([1, -1], ntuple(j->j==i ? 2 : 1, ndims(X)))
+        # TODO(johnnychen94): reuse the allocated memory
+        F∇ᵢ = freqkernel(∇ᵢ, size(X))
+        @. FΔ += Base.abs2(F∇ᵢ)
+    end
+    return FΔ
+end
 
 end # module