Tests/src/Tests.jl

##################
# Denoise testing
##################

module Tests

# Our exports
export test_denoise,
       test_deblur

# Dependencies
using Printf
using FileIO
using ColorTypes: Gray
# ColorVectorSpace is only needed to ensure that conversions
# between different ColorTypes are defined.
import ColorVectorSpace
import TestImages

using AlgTools.Util
using AlgTools.LinkedLists
using ImageTools.Denoise
using ImageTools.Visualise
using ImageTools.TVRecon
using ImageTools.ImFilter

# Parameters
const default_save_prefix="result_"

const default_params = (
    ρ = 0,
    noise_level = 0.5,
    verbose_iter = 10,
    maxiter = 1000,
    save_results = false,
    image_name = "lighthouse",
    save_iterations = false
)

const denoise_params = (
    α = 2,
    # FISTA
    τ₀ = 0.9,
    accel = true,
)

const deblur_params = (
    α = 0.5,
    # PDPS
    τ₀ = 5,
    σ₀ = 0.99/5,
    accel = false,
)

function save_results(params, b, b_data, x)
    if params.save_results
        perffile = params.save_prefix * ".txt"
        println("Saving " * perffile)
        write_log(perffile, st.log, "# params = $(params)\n")
        fn = (t, ext) -> "$(params.save_prefix)_$(t).$(ext)"
        save(File(format"PNG", fn("true", "png")), grayimg(b))
        save(File(format"PNG", fn("data", "png")), grayimg(b_data))
        save(File(format"PNG", fn("reco", "png")), grayimg(x))
    end
end

###############
# Denoise test
###############

function test_denoise(;
                      visualise=true,
                      save_prefix=default_save_prefix,
                      kwargs...)

    # Parameters for this experiment
    params = default_params ⬿ denoise_params ⬿ kwargs
    params = params ⬿ (save_prefix = save_prefix * "denoise_" * params.image_name,)

    # Load image and add noise
    b = Float64.(Gray{Float64}.(TestImages.testimage(params.image_name)))
    b_noisy = b .+ params.noise_level.*randn(size(b)...)

    # Launch (background) visualiser
    st, iterate = initialise_visualisation(visualise)

    # Run algorithm
    x, y, st = denoise_fista(b_noisy; iterate=iterate, params=params)

    save_results(params, b, b_noisy, x)

    # Exit background visualiser
    finalise_visualisation(st)
end

##############
# Deblur test
##############

function test_deblur(;
                     visualise=true,
                     save_prefix=default_save_prefix,
                     kwargs...)

    # Parameters for this experiment
    params = default_params ⬿ deblur_params ⬿ kwargs
    params = params ⬿ (save_prefix = save_prefix * "deblur_" * params.image_name,)

    # Load image and add noise
    b = Float64.(Gray{Float64}.(TestImages.testimage(params.image_name)))
    op = FilterKernel(gaussian(2, [7,7]))
    b_data = op(b .+ params.noise_level.*randn(size(b)...))

    # Launch (background) visualiser
    st, iterate = initialise_visualisation(visualise)

    # Run algorithm
    x, y, st = recon_pdps(b_data, op; xinit=zeros(size(b)...), iterate=iterate, params=params)

    save_results(params, b, b_data, x)

    # Exit background visualiser
    finalise_visualisation(st)
end

end