|
| 1 | +module Models |
| 2 | + |
| 3 | +using ImageBase |
| 4 | +using ImageBase.ImageCore.MappedArrays: of_eltype |
| 5 | +using ImageBase.FiniteDiff |
| 6 | + |
| 7 | +# Introduced in ColorVectorSpace v0.9.3 |
| 8 | +# https://github.com/JuliaGraphics/ColorVectorSpace.jl/pull/172 |
| 9 | +using ImageBase.ImageCore.ColorVectorSpace.Future: abs2 |
| 10 | + |
| 11 | +""" |
| 12 | +This submodule provides predefined image-related models and its solvers that can be reused |
| 13 | +by many image processing tasks. |
| 14 | +
|
| 15 | +- solve the Rudin Osher Fatemi (ROF) model using the primal-dual method: [`solve_ROF_PD`](@ref) and [`solve_ROF_PD!`](@ref) |
| 16 | +""" |
| 17 | +Models |
| 18 | + |
| 19 | +export solve_ROF_PD, solve_ROF_PD! |
| 20 | + |
| 21 | + |
| 22 | +##### implementation details |
| 23 | + |
| 24 | +""" |
| 25 | + solve_ROF_PD([T], img::AbstractArray, λ; kwargs...) |
| 26 | +
|
| 27 | +Return a smoothed version of `img`, using Rudin-Osher-Fatemi (ROF) filtering, more commonly |
| 28 | +known as Total Variation (TV) denoising or TV regularization. This algorithm is based on the |
| 29 | +primal-dual method. |
| 30 | +
|
| 31 | +This function applies to generic N-dimensional colorant array and is also CUDA-compatible. |
| 32 | +See also [`solve_ROF_PD!`](@ref) for the in-place version. |
| 33 | +
|
| 34 | +# Arguments |
| 35 | +
|
| 36 | +- `T`: the output element type. By default it is `float32(eltype(img))`. |
| 37 | +- `img`: the input image, usually a noisy image. |
| 38 | +- `λ`: the regularization coefficient. Larger `λ` results in more smoothing. |
| 39 | +
|
| 40 | +# Parameters |
| 41 | +
|
| 42 | +- `num_iters::Int`: The number of iterations before stopping. |
| 43 | +
|
| 44 | +# Examples |
| 45 | +
|
| 46 | +```julia |
| 47 | +using ImageFiltering |
| 48 | +using ImageFiltering.Models: solve_ROF_PD |
| 49 | +using ImageQualityIndexes |
| 50 | +using TestImages |
| 51 | +
|
| 52 | +img_ori = float.(testimage("cameraman")) |
| 53 | +img_noisy = img_ori .+ 0.1 .* randn(size(img_ori)) |
| 54 | +assess_psnr(img_noisy, img_ori) # ~20 dB |
| 55 | +
|
| 56 | +img_smoothed = solve_ROF_PD(img_noisy, 0.015, 50) |
| 57 | +assess_psnr(img_smoothed, img_ori) # ~27 dB |
| 58 | +
|
| 59 | +# larger λ produces over-smoothed result |
| 60 | +img_smoothed = solve_ROF_PD(img_noisy, 5, 50) |
| 61 | +assess_psnr(img_smoothed, img_ori) # ~21 dB |
| 62 | +``` |
| 63 | +
|
| 64 | +# Extended help |
| 65 | +
|
| 66 | +Mathematically, this function solves the following ROF model using the primal-dual method: |
| 67 | +
|
| 68 | +```math |
| 69 | +\\min_u \\lVert u - g \\rVert^2 + \\lambda\\lvert\\nabla u\\rvert |
| 70 | +``` |
| 71 | +
|
| 72 | +# References |
| 73 | +
|
| 74 | +- [1] Chambolle, A. (2004). "An algorithm for total variation minimization and applications". _Journal of Mathematical Imaging and Vision_. 20: 89–97 |
| 75 | +- [2] https://en.wikipedia.org/wiki/Total_variation_denoising |
| 76 | +""" |
| 77 | +solve_ROF_PD(img::AbstractArray{T}, args...) where T = solve_ROF_PD(float32(T), img, args...) |
| 78 | +function solve_ROF_PD(::Type{T}, img::AbstractArray, args...) where T |
| 79 | + u = similar(img, T) |
| 80 | + buffer = preallocate_solve_ROF_PD(T, img) |
| 81 | + solve_ROF_PD!(u, buffer, img, args...) |
| 82 | +end |
| 83 | + |
| 84 | +# non-exported helper |
| 85 | +preallocate_solve_ROF_PD(img::AbstractArray{T}) where T = preallocate_solve_ROF_PD(float32(T), img) |
| 86 | +function preallocate_solve_ROF_PD(::Type{T}, img) where T |
| 87 | + div_p = similar(img, T) |
| 88 | + p = ntuple(i->similar(img, T), ndims(img)) |
| 89 | + ∇u = ntuple(i->similar(img, T), ndims(img)) |
| 90 | + ∇u_mag = similar(img, eltype(T)) |
| 91 | + return div_p, p, ∇u, ∇u_mag |
| 92 | +end |
| 93 | + |
| 94 | +""" |
| 95 | + solve_ROF_PD!(out, buffer, img, λ, num_iters) |
| 96 | +
|
| 97 | +The in-place version of [`solve_ROF_PD`](@ref). |
| 98 | +
|
| 99 | +It is not uncommon to use ROF solver in a higher-level loop, in which case it makes sense to |
| 100 | +preallocate the output and intermediate arrays to make it faster. |
| 101 | +
|
| 102 | +!!! note "Buffer" |
| 103 | + The content and meaning of `buffer` might change without any notice if the internal |
| 104 | + implementation is changed. Use `preallocate_solve_ROF_PD` helper function to avoid |
| 105 | + potential changes. |
| 106 | +
|
| 107 | +# Examples |
| 108 | +
|
| 109 | +```julia |
| 110 | +using ImageFiltering.Models: preallocate_solve_ROF_PD |
| 111 | +
|
| 112 | +out = similar(img) |
| 113 | +buffer = preallocate_solve_ROF_PD(img) |
| 114 | +solve_ROF_PD!(out, buffer, img, 0.2, 30) |
| 115 | +``` |
| 116 | +
|
| 117 | +""" |
| 118 | +function solve_ROF_PD!( |
| 119 | + out::AbstractArray{T}, |
| 120 | + buffer::Tuple, |
| 121 | + img::AbstractArray, |
| 122 | + λ::Real, |
| 123 | + num_iters::Integer) where T |
| 124 | + # seperate a stub method to reduce latency |
| 125 | + FT = float32(T) |
| 126 | + if FT == T |
| 127 | + solve_ROF_PD!(out, buffer, img, Float32(λ), Int(num_iters)) |
| 128 | + else |
| 129 | + solve_ROF_PD!(out, buffer, FT.(img), Float32(λ), Int(num_iters)) |
| 130 | + end |
| 131 | +end |
| 132 | +function solve_ROF_PD!( |
| 133 | + out::AbstractArray, |
| 134 | + (div_p, p, ∇u, ∇u_mag)::Tuple, |
| 135 | + img::AbstractArray, |
| 136 | + λ::Float32, |
| 137 | + num_iters::Int) |
| 138 | + # Total Variation regularized image denoising using the primal dual algorithm |
| 139 | + # Implement according to reference [1] |
| 140 | + τ = 1//4 # see 2nd remark after proof of Theorem 3.1. |
| 141 | + |
| 142 | + # use the same symbol in the paper |
| 143 | + u, g = out, img |
| 144 | + |
| 145 | + fgradient!(p, g) |
| 146 | + # This iterates Eq. (9) of [1] |
| 147 | + # TODO(johnnychen94): set better stop criterion |
| 148 | + for _ in 1:num_iters |
| 149 | + fdiv!(div_p, p) |
| 150 | + # multiply term inside ∇ by -λ. Thm. 3.1 relates this to `u` via Eq. 7. |
| 151 | + @. u = g - λ*div_p |
| 152 | + fgradient!(∇u, u) |
| 153 | + _l2norm_vec!(∇u_mag, ∇u) # |∇(g - λdiv p)| |
| 154 | + # Eq. (9): update p |
| 155 | + for i in 1:length(p) |
| 156 | + @. p[i] = (p[i] - (τ/λ)*∇u[i])/(1 + (τ/λ) * ∇u_mag) |
| 157 | + end |
| 158 | + end |
| 159 | + return u |
| 160 | +end |
| 161 | + |
| 162 | +function _l2norm_vec!(out, Vs::Tuple) |
| 163 | + all(v->axes(out) == axes(v), Vs) || throw(ArgumentError("All axes of input data should be the same.")) |
| 164 | + @. out = abs2(Vs[1]) |
| 165 | + for v in Vs[2:end] |
| 166 | + @. out += abs2(v) |
| 167 | + end |
| 168 | + @. out = sqrt(out) |
| 169 | + return out |
| 170 | +end |
| 171 | + |
| 172 | + |
| 173 | +end # module |
0 commit comments