implement generic ROF model using Chambolle04 primal-dual method

Author: johnnychen94

.gitignore

@@ -4,8 +4,7 @@
 .vscode
 docs/build
 docs/site
-docs/Manifest.toml
 docs/src/democards
-/Manifest.toml
+Manifest.toml
 /.benchmarkci
 /benchmark/*.json

Project.toml

@@ -9,6 +9,7 @@ ComputationalResources = "ed09eef8-17a6-5b46-8889-db040fac31e3"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 FFTViews = "4f61f5a4-77b1-5117-aa51-3ab5ef4ef0cd"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+ImageBase = "c817782e-172a-44cc-b673-b171935fbb9e"
 ImageCore = "a09fc81d-aa75-5fe9-8630-4744c3626534"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
@@ -24,6 +25,7 @@ ComputationalResources = "0.3"
 DataStructures = "0.17.7, 0.18"
 FFTViews = "0.3"
 FFTW = "0.3, 1"
+ImageBase = "0.1.5"
 ImageCore = "0.9"
 OffsetArrays = "1.9"
 Reexport = "1.1"
@@ -33,10 +35,14 @@ julia = "1"
 
 [extras]
 AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
+ImageIO = "82e4d734-157c-48bb-816b-45c225c6df19"
+ImageMagick = "6218d12a-5da1-5696-b52f-db25d2ecc6d1"
 ImageMetadata = "bc367c6b-8a6b-528e-b4bd-a4b897500b49"
+ImageQualityIndexes = "2996bd0c-7a13-11e9-2da2-2f5ce47296a9"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TestImages = "5e47fb64-e119-507b-a336-dd2b206d9990"
 
 [targets]
-test = ["AxisArrays", "ImageMetadata", "Logging", "Random", "Test"]
+test = ["AxisArrays", "ImageIO", "ImageMagick", "ImageMetadata", "ImageQualityIndexes", "Logging", "Random", "Test", "TestImages"]

docs/src/function_reference.md

@@ -72,6 +72,12 @@ Algorithm.IIR
 Algorithm.Mixed
 ```
 
+# Solvers for predefined models
+
+```@autodocs
+Modules = [Models]
+```
+
 # Internal machinery
 
 ```@docs

src/ImageFiltering.jl

@@ -92,6 +92,8 @@ include("mapwindow.jl")
 using .MapWindow
 include("extrema.jl")
 
+include("models.jl")
+
 function __init__()
     # See ComputationalResources README for explanation
     push!(LOAD_PATH, dirname(@__FILE__))

src/models.jl

@@ -0,0 +1,114 @@
+module Models
+
+using ImageBase
+using ImageBase.ImageCore.MappedArrays: of_eltype
+using ImageBase.FiniteDiff
+
+# Introduced in ColorVectorSpace v0.9.3
+# https://github.com/JuliaGraphics/ColorVectorSpace.jl/pull/172
+using ImageBase.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 and is also CUDA-compatible.
+
+# 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 ImageFiltering
+using ImageFiltering.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, λ::Real, num_iters::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:num_iters
+        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
+    @. out = sqrt(out)
+    return out
+end
+
+
+end # module

test/cuda/Project.toml

@@ -0,0 +1,10 @@
+[deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ImageBase = "c817782e-172a-44cc-b673-b171935fbb9e"
+ImageFiltering = "6a3955dd-da59-5b1f-98d4-e7296123deb5"
+ImageIO = "82e4d734-157c-48bb-816b-45c225c6df19"
+ImageMagick = "6218d12a-5da1-5696-b52f-db25d2ecc6d1"
+ImageQualityIndexes = "2996bd0c-7a13-11e9-2da2-2f5ce47296a9"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TestImages = "5e47fb64-e119-507b-a336-dd2b206d9990"

test/cuda/models.jl

@@ -0,0 +1,67 @@
+using ImageFiltering.Models
+
+@testset "solve_ROF_PD" begin
+    # This testset is modified from its CPU version
+
+    @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.05, 20)
+        @test ndims(img_smoothed) == 2
+        @test eltype(img_smoothed) <: Gray
+        @test assess_psnr(img_smoothed, img) > 31.67
+        @test assess_ssim(img_smoothed, img) > 0.90
+
+        img_noisy_cu = CuArray(float32.(img_noisy))
+        img_smoothed_cu = solve_ROF_PD(img_noisy_cu, 0.05, 20)
+        @test img_smoothed_cu isa CuArray
+        @test eltype(eltype(img_smoothed_cu)) == Float32
+        @test Array(img_smoothed_cu) ≈ img_smoothed
+
+        # 2D RGB
+        img = restrict(testimage("lighthouse"))
+        img_noisy = img .+ colorview(RGB, ntuple(i->0.05.*randn(MersenneTwister(i), size(img)), 3)...)
+        img_smoothed = solve_ROF_PD(img_noisy, 0.03, 20)
+        @test ndims(img_smoothed) == 2
+        @test eltype(img_smoothed) <: RGB
+        @test assess_psnr(img_smoothed, img) > 32.15
+        @test assess_ssim(img_smoothed, img) > 0.90
+
+        img_noisy_cu = CuArray(float32.(img_noisy))
+        img_smoothed_cu = solve_ROF_PD(img_noisy_cu, 0.03, 20)
+        @test img_smoothed_cu isa CuArray
+        @test eltype(eltype(img_smoothed_cu)) == Float32
+        @test Array(img_smoothed_cu) ≈ img_smoothed
+
+        # 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.02, 20)
+        @test ndims(img_smoothed) == 3
+        @test eltype(img_smoothed) <: Gray
+        @test assess_psnr(img_smoothed, img) > 31.78
+        @test assess_ssim(img_smoothed, img) > 0.85
+
+        img_noisy_cu = CuArray(float32.(img_noisy))
+        img_smoothed_cu = solve_ROF_PD(img_noisy_cu, 0.02, 20)
+        @test img_smoothed_cu isa CuArray
+        @test eltype(eltype(img_smoothed_cu)) == Float32
+        @test Array(img_smoothed_cu) ≈ img_smoothed
+
+        # 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.02, 20)
+        @test ndims(img_smoothed) == 3
+        @test eltype(img_smoothed) <: RGB
+        @test assess_psnr(img_smoothed, img) > 31.17
+        @test assess_ssim(img_smoothed, img) > 0.79
+
+        img_noisy_cu = CuArray(float32.(img_noisy))
+        img_smoothed_cu = solve_ROF_PD(img_noisy_cu, 0.02, 20)
+        @test img_smoothed_cu isa CuArray
+        @test eltype(eltype(img_smoothed_cu)) == Float32
+        @test Array(img_smoothed_cu) ≈ img_smoothed
+    end
+end

test/cuda/runtests.jl

@@ -0,0 +1,17 @@
+# This file is maintained in a way to support CUDA-only test via
+# `julia --project=test/cuda -e 'include("runtests.jl")'`
+using ImageFiltering
+using CUDA
+using TestImages
+using ImageBase
+using ImageQualityIndexes
+using Test
+using Random
+
+CUDA.allowscalar(false)
+
+@testset "ImageFiltering" begin
+    if CUDA.functional()
+        include("models.jl")
+    end
+end

test/models.jl

@@ -0,0 +1,54 @@
+using ImageFiltering.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.05, 20)
+        @test ndims(img_smoothed) == 2
+        @test eltype(img_smoothed) <: Gray
+        @test assess_psnr(img_smoothed, img) > 31.67
+        @test assess_ssim(img_smoothed, img) > 0.90
+
+        # 2D RGB
+        img = restrict(testimage("lighthouse"))
+        img_noisy = img .+ colorview(RGB, ntuple(i->0.05.*randn(MersenneTwister(i), size(img)), 3)...)
+        img_smoothed = solve_ROF_PD(img_noisy, 0.03, 20)
+        @test ndims(img_smoothed) == 2
+        @test eltype(img_smoothed) <: RGB
+        @test assess_psnr(img_smoothed, img) > 32.15
+        @test assess_ssim(img_smoothed, img) > 0.90
+
+        # 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.02, 20)
+        @test ndims(img_smoothed) == 3
+        @test eltype(img_smoothed) <: Gray
+        @test assess_psnr(img_smoothed, img) > 31.78
+        @test assess_ssim(img_smoothed, img) > 0.85
+
+        # 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.02, 20)
+        @test ndims(img_smoothed) == 3
+        @test eltype(img_smoothed) <: RGB
+        @test assess_psnr(img_smoothed, img) > 31.17
+        @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(rand(N0f8, 20, 20), -1, -1)
+        out = solve_ROF_PD(Ao, 0.01, 5)
+        @test axes(out) == axes(Ao)
+    end
+end

test/runtests.jl

@@ -1,5 +1,8 @@
-using ImageFiltering, ImageCore
+using ImageFiltering, ImageCore, ImageBase
+using OffsetArrays
 using Test
+using TestImages
+using ImageQualityIndexes
 import StaticArrays
 using Random
 
@@ -31,5 +34,30 @@ include("mapwindow.jl")
 include("extrema.jl")
 include("basic.jl")
 include("gabor.jl")
+include("models.jl")
 
+
+CUDA_INSTALLED = false
+try
+    global CUDA_INSTALLED
+    # This errors with `IOError` when nvidia driver is not available,
+    # in which case we don't even need to try `using CUDA`
+    run(pipeline(`nvidia-smi`, stdout=devnull, stderr=devnull))
+    push!(LOAD_PATH, "@v#.#") # force using global CUDA installation
+
+    @eval using CUDA
+    CUDA.allowscalar(false)
+    CUDA_INSTALLED = true
+catch e
+    e isa Base.IOError || @warn e LOAD_PATH
+end
+CUDA_FUNCTIONAL = CUDA_INSTALLED && CUDA.functional()
+if CUDA_FUNCTIONAL
+    @info "CUDA test: enabled"
+    @testset "CUDA" begin
+        include("cuda/runtests.jl")
+    end
+else
+    @warn "CUDA test: disabled"
+end
 nothing