implement generic ROF model using Chambolle04 primal-dual method

Author: johnnychen94

Project.toml

@@ -17,11 +17,13 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ImageFiltering = "6a3955dd-da59-5b1f-98d4-e7296123deb5"
 ImageIO = "82e4d734-157c-48bb-816b-45c225c6df19"
 ImageMagick = "6218d12a-5da1-5696-b52f-db25d2ecc6d1"
+ImageQualityIndexes = "2996bd0c-7a13-11e9-2da2-2f5ce47296a9"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StackViews = "cae243ae-269e-4f55-b966-ac2d0dc13c15"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestImages = "5e47fb64-e119-507b-a336-dd2b206d9990"
 
 [targets]
-test = ["Aqua", "Documenter", "Test", "ImageFiltering", "ImageIO", "ImageMagick", "OffsetArrays", "Statistics", "StackViews", "TestImages"]
+test = ["Aqua", "Documenter", "Test", "ImageFiltering", "ImageIO", "ImageMagick", "ImageQualityIndexes", "OffsetArrays", "Random", "Statistics", "StackViews", "TestImages"]

src/ImageBase.jl

@@ -26,6 +26,9 @@ include("diff.jl")
 include("restrict.jl")
 include("utils.jl")
 include("statistics.jl")
+
+include("models.jl")
+
 include("compat.jl")
 include("deprecated.jl")
 

src/models.jl

@@ -0,0 +1,113 @@
+module Models
+
+using ImageCore
+using ImageCore.MappedArrays: of_eltype
+using ..FiniteDiff
+
+# Introduced in ColorVectorSpace v0.9.3
+# https://github.com/JuliaGraphics/ColorVectorSpace.jl/pull/172
+using ImageCore.ColorVectorSpace.Future: abs2
+
+"""
+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)
+"""
+Models
+
+export solve_ROF_PD
+
+
+##### implementation details
+
+"""
+    solve_ROF_PD(img：：AbstractArray, λ; kwargs...)
+
+Perform Rudin-Osher-Fatemi (ROF) filtering, more commonly known as Total Variation (TV)
+denoising or TV regularization.
+
+This function applies to generic n-dimensional colorant array.
+
+# Arguments
+
+- `img`: the input image, usually is a noisy image.
+- `λ`: the regularization coefficient. Larger `λ` would produce more smooth image.
+
+# Parameters
+
+- `num_iters::Int`: The number of iterations before stopping.
+
+# Examples
+
+```julia
+using ImageBase
+using ImageBase.Models: solve_ROF_PD
+using ImageQualityIndexes
+using TestImages
+
+img_ori = float.(testimage("cameraman"))
+img_noisy = img_ori .+ 0.1 .* randn(size(img_ori))
+assess_psnr(img_noisy, img_ori) # ~20 dB
+
+img_smoothed = solve_ROF_PD(img_noisy, 0.015, 50)
+assess_psnr(img_smoothed, img_ori) # ~27 dB
+
+# larger λ produces over-smoothed result
+img_smoothed = solve_ROF_PD(img_noisy, 5, 50)
+assess_psnr(img_smoothed, img_ori) # ~21 dB
+```
+
+# Extended help
+
+Mathematically, this function solves the following ROF model using the primal-dual method:
+
+```math
+\\min_u \\lVert u - g \\rVert^2 + \\lambda\\lvert\\nabla u\\rvert
+```
+
+# References
+
+- [1] Chambolle, A. (2004). "An algorithm for total variation minimization and applications". _Journal of Mathematical Imaging and Vision_. 20: 89–97
+- [2] https://en.wikipedia.org/wiki/Total_variation_denoising
+"""
+function solve_ROF_PD(img::AbstractArray, λ::Number, iterations::Integer)
+    # Total Variation regularized image denoising using the primal dual algorithm
+    # Implement according to reference [1]
+    τ = 1/4   # see 2nd remark after proof of Theorem 3.1.
+
+    g = of_eltype(floattype(eltype(img)), img) # use the same symbol in the paper
+    u = similar(g)
+    p = fgradient(g)
+    div_p = similar(g)
+    ∇u = map(similar, p)
+    ∇u_mag = similar(g, eltype(eltype(g)))
+
+    # This iterates Eq. (9) of [1]
+    # TODO(johnnychen94): set better stop criterion
+    for _ in 1:iterations
+        fdiv!(div_p, p)
+        # multiply term inside ∇ by -λ. Thm. 3.1 relates this to `u` via Eq. 7.
+        @. u = g - λ*div_p
+        fgradient!(∇u, u)
+        _l2norm_vec!(∇u_mag, ∇u) # |∇(g - λdiv p)|
+        # Eq. (9): update p
+        for i in 1:length(p)
+            @. p[i] = (p[i] - (τ/λ)*∇u[i])/(1 + (τ/λ) * ∇u_mag)
+        end
+    end
+    return u
+end
+
+
+function _l2norm_vec!(out, Vs::Tuple)
+    all(v->axes(out) == axes(v), Vs) || throw(ArgumentError("All axes of input data should be the same."))
+    @. out = abs2(Vs[1])
+    for v in Vs[2:end]
+        @. out += abs2(v)
+    end
+    return out
+end
+
+
+end # module

test/models.jl

@@ -0,0 +1,54 @@
+using ImageBase.Models
+
+@testset "solve_ROF_PD" begin
+    # Note: random seed really matters a lot
+
+    @testset "Numerical" begin
+        # 2D Gray
+        img = restrict(testimage("cameraman"))
+        img_noisy = img .+ 0.05randn(MersenneTwister(0), size(img))
+        img_smoothed = solve_ROF_PD(img_noisy, 0.01, 20)
+        @test ndims(img_smoothed) == 2
+        @test eltype(img_smoothed) <: Gray
+        @test assess_psnr(img_smoothed, img) > 28.70
+        @test assess_ssim(img_smoothed, img) > 0.845
+
+        # 2D RGB
+        img = restrict(testimage("mandril"))
+        img_noisy = img .+ colorview(RGB, ntuple(i->0.05.*randn(MersenneTwister(i), size(img)), 3)...)
+        img_smoothed = solve_ROF_PD(img_noisy, 0.001, 20)
+        @test ndims(img_smoothed) == 2
+        @test eltype(img_smoothed) <: RGB
+        @test assess_psnr(img_smoothed, img) > 28.82
+        @test assess_ssim(img_smoothed, img) > 0.849
+
+        # 3D Gray
+        img = restrict(testimage("mri"), (1, 2))
+        img_noisy = img .+ 0.05randn(MersenneTwister(0), size(img))
+        img_smoothed = solve_ROF_PD(img_noisy, 0.005, 20)
+        @test ndims(img_smoothed) == 3
+        @test eltype(img_smoothed) <: Gray
+        @test assess_psnr(img_smoothed, img) > 27.92
+        @test assess_ssim(img_smoothed, img) > 0.66
+
+        # 3D RGB
+        img = RGB.(restrict(testimage("mri"), (1, 2)))
+        img_noisy = img .+ colorview(RGB, ntuple(i->0.05.*randn(MersenneTwister(i), size(img)), 3)...)
+        img_smoothed = solve_ROF_PD(img_noisy, 0.001, 20)
+        @test ndims(img_smoothed) == 3
+        @test eltype(img_smoothed) <: RGB
+        @test assess_psnr(img_smoothed, img) > 30.73
+        @test assess_ssim(img_smoothed, img) > 0.79
+    end
+
+    @testset "FixedPointNumbers" begin
+        A = rand(N0f8, 20, 20)
+        @test solve_ROF_PD(A, 0.01, 5) ≈ solve_ROF_PD(float32.(A), 0.01, 5)
+    end
+
+    @testset "OffsetArray" begin
+        Ao = OffsetArray(A, -1, -1)
+        out = solve_ROF_PD(Ao, 0.01, 5)
+        @test axes(out) == axes(Ao)
+    end
+end

test/runtests.jl

@@ -1,6 +1,8 @@
 using ImageBase, OffsetArrays, StackViews
 using ImageFiltering
+using ImageQualityIndexes
 using Test, TestImages, Aqua, Documenter
+using Random
 
 using OffsetArrays: IdentityUnitRange
 include("testutils.jl")
@@ -24,6 +26,7 @@ include("testutils.jl")
     include("diff.jl")
     include("restrict.jl")
     include("statistics.jl")
+    include("models.jl")
 
     @info "deprecations are expected"
     include("deprecated.jl")