implement generic ROF model using Chambolle04 primal-dual method

[johnnychen94]

Images.imROF only supports 2D Gray image, with the RGB image processed per channel. This version implements an RGB-native version (by fusing RGB into the inner calculation) and is also dimension-agnostic because of #22

Comparison to Images.imROF:

using Images
using ImageBase.Models
using ImageQualityIndexes
using TestImages
using Random

# 2d Gray
img = testimage("cameraman")
img_noisy = img .+ 0.1*randn(MersenneTwister(0), size(img))
R = CartesianIndices(img)[5:end-5, 5:end-5]

smoothed_img = @btime solve_ROF_PD(img_noisy, 0.2, 30)
# 78.293 ms (134 allocations: 134.00 MiB)
assess_psnr(smoothed_img[R], img[R]) # 28.3089

smoothed_img_legacy = @btime Images.imROF(img_noisy, 0.2, 30)
# master (type unstable): 3.309996 seconds (78.15 M allocations: 2.398 GiB, 13.47% gc time)
# master (fixed): 124.702 ms (1594 allocations: 546.04 MiB)
assess_psnr(smoothed_img_legacy[R], img[R]) # 28.2979

# 2d RGB
img = float32.(testimage("lighthouse"))
img_noisy = img .+ colorview(RGB, ntuple(i->0.05.*randn(MersenneTwister(i), size(img)), 3)...)
R = CartesianIndices(img)[5:end-5, 5:end-5]

smoothed_img = @btime solve_ROF_PD(img_noisy, 0.2, 30)
# 447.726 ms (134 allocations: 597.00 MiB)
assess_psnr(smoothed_img[R], img[R]) # 25.1099

smoothed_img_legacy = @btime Images.imROF(img_noisy, 0.2, 30)
# master (type unstable): 13.527543 seconds (387.26 M allocations: 10.285 GiB, 11.06% gc time)
# master (fixed): 865.158 ms (4790 allocations: 2.42 GiB)
assess_psnr(smoothed_img_legacy[R], img[R]) # 25.1021

See the tests for 3d cases.

Todo:

• CUDA test

[codecov]

Codecov Report

Merging #24 (d1bb8fb) into master (5687005) will increase coverage by 0.81%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master      #24      +/-   ##
==========================================
+ Coverage   92.07%   92.88%   +0.81%     
==========================================
  Files           5        6       +1     
  Lines         227      253      +26     
==========================================
+ Hits          209      235      +26     
  Misses         18       18              

Impacted Files	Coverage Δ
src/models.jl	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 5687005...d1bb8fb. Read the comment docs.

[johnnychen94]

This PR and #22 are specially written in a vectorized manner so that it also works for GPU arrays:

using ImageBase
using ImageBase.Models
using TestImages
using Random

using CUDA
CUDA.allowscalar(false)

# 2d Gray
img = testimage("cameraman")

img_noisy = img .+ 0.1f0*randn(MersenneTwister(0), Float32, size(img))
@btime solve_ROF_PD($img_noisy, 0.2, 30); 
# 101.396 ms (134 allocations: 134.00 MiB)

img_noisy_cu = CuArray(Float32.(img_noisy))
@btime solve_ROF_PD($img_noisy_cu, 0.2, 30); # GTX 3090
# 4.060 ms (8793 allocations: 754.84 KiB)

# 2d RGB
img = testimage("lighthouse")

img_noisy = img .+ colorview(RGB, ntuple(i->0.05.*randn(MersenneTwister(i), size(img)), 3)...)
@btime solve_ROF_PD($img_noisy, 0.2, 30); 
# 610.880 ms (134 allocations: 597.00 MiB)

img_noisy_cu = CuArray(float32.(img_noisy))
@btime solve_ROF_PD($img_noisy_cu, 0.2, 30); # GTX 3090
# 7.176 ms (8433 allocations: 729.53 KiB)

@timholy how do we test CUDA codes in JuliaImages? I have some special scripts in BlockMatching.jl to detect if Nvidia driver is installed, and if it is then it uses the CUDA package in the root environment and runs related tests. But since we don't have GPU servers we can't set up CI here.

[johnnychen94]

Also ping @JKay0327 since this PR takes a very similar spirit to what I and @timholy thought when we reviewed JuliaImages/ImageSegmentation.jl#84

[johnnychen94]

About the code organization: in the near future I'm going to add the Split Bregman solver for the ROF model, which requires freqkernel and spacekernel (FFT-based) from ImageFiltering.jl. Maybe I should just implement this primal-dual solver in ImageFiltering so that we don't get confused on where to put the codes. Thoughts?

[timholy]

@timholy how do we test CUDA codes in JuliaImages?

Nothing has been set up yet. We might want to ask the JuliaGPU folks what's involved in setting this up.

Maybe I should just implement this primal-dual solver in ImageFiltering so that we don't get confused on where to put the codes. Thoughts?

That sounds like a good idea to me.

[timholy]

Looks good. My main concern is whether this belongs here or in a more specialized package. I thought ImageBase is supposed to have basic, widely-used algorithms. Or are you foreseeing using this widely across several other JuliaImages packages?

[timholy]

To avoid forcing things to Float64 would it be better to just divide by 4 somewhere?

[timholy]

I have my fingers crossed for JuliaLang/julia#41915

[johnnychen94]

Nothing has been set up yet. We might want to ask the JuliaGPU folks what's involved in setting this up.

From slack, it seems that we just need to apply https://github.com/JuliaGPU/buildkite.

[johnnychen94]

close in favor of JuliaImages/ImageFiltering.jl#233