onnx_inpainting.cc File Reference

#include <algorithm>
#include <exception>
#include <functional>
#include <iostream>
#include <random>
#include <vector>
#include <ctime>
#include <catch2/catch_all.hpp>
#include "sopt/imaging_forward_backward.h"
#include "sopt/real_indicator.h"
#include "sopt/logging.h"
#include "sopt/maths.h"
#include "sopt/relative_variation.h"
#include "sopt/sampling.h"
#include "sopt/types.h"
#include "sopt/utilities.h"
#include "sopt/ort_session.h"
#include "sopt/onnx_differentiable_func.h"
#include "tools_for_tests/directories.h"
#include "tools_for_tests/tiffwrappers.h"

Include dependency graph for onnx_inpainting.cc:

Go to the source code of this file.

Typedefs
using	Scalar = double
using	Vector = sopt::Vector< Scalar >
using	Matrix = sopt::Matrix< Scalar >
using	Image = sopt::Image< Scalar >
using	LinearTransform = sopt::LinearTransform< Vector >

Functions
	TEST_CASE ("Inpainting")

Typedef Documentation

Image

using Image = sopt::Image<Scalar>

Definition at line 32 of file onnx_inpainting.cc.

LinearTransform

using LinearTransform = sopt::LinearTransform<Vector>

Definition at line 33 of file onnx_inpainting.cc.

Matrix

using Matrix = sopt::Matrix<Scalar>

Definition at line 31 of file onnx_inpainting.cc.

Scalar

using Scalar = double

Definition at line 29 of file onnx_inpainting.cc.

Vector

using Vector = sopt::Vector<Scalar>

Definition at line 30 of file onnx_inpainting.cc.

Function Documentation

TEST_CASE()

TEST_CASE

(

"Inpainting"

)

Definition at line 35 of file onnx_inpainting.cc.

                       {
 
  // black magic?
  double lambda = 5e4;
  double mu = 20;
 
  extern std::unique_ptr<std::mt19937_64> mersenne;
  std::string const input = "cameraman256";
 
 
  Image const image = sopt::tools::read_standard_tiff(input);
 
  sopt::t_uint nmeasure = std::floor(0.5 * image.size());
  LinearTransform const sampling =
      sopt::linear_transform<Scalar>(sopt::Sampling(image.size(), nmeasure, *mersenne));
 
  Vector const y0 = sampling * Vector::Map(image.data(), image.size());
  auto constexpr snr = 30.0;
  auto const sigma = y0.stableNorm() / std::sqrt(y0.size()) * std::pow(10.0, -(snr / 20.0));
  auto const epsilon = std::sqrt(nmeasure + 2 * std::sqrt(y0.size())) * sigma;
 
  // set the model function and gradient
  std::string const prior_path = std::string(sopt::tools::models_directory() + "/example_cost_dynamic_CRR_sigma_5_t_5.onnx");
  std::string const prior_gradient_path = std::string(sopt::tools::models_directory() + "/example_grad_dynamic_CRR_sigma_5_t_5.onnx");
  std::shared_ptr<sopt::ONNXDifferentiableFunc<Scalar>> diff_function = std::make_shared<sopt::ONNXDifferentiableFunc<Scalar>>(prior_path, prior_gradient_path, sigma, mu, lambda, sampling);
 
  std::normal_distribution<> gaussian_dist(0, sigma);
  Vector y(y0.size());
  for (sopt::t_int i = 0; i < y0.size(); i++) y(i) = y0(i) + gaussian_dist(*mersenne);
 
  Eigen::VectorXd dirty_image = sampling.adjoint() * y;
  Eigen::VectorXd init_res = ((sampling * dirty_image) - y);
  auto init_res_norm = init_res.array().abs().sum();
  SOPT_HIGH_LOG("Initial residual norm: {}", init_res_norm);
 
  sopt::t_real constexpr regulariser_strength = 18;
  sopt::t_real const beta = sigma * sigma * 0.5;
 
  // Arbitrary (absolute) tolerance level to produce a reasonable image which converges
  sopt::RelativeVariation<Scalar> scalvar(0.4,
                                          "Convergence function");
  std::function<bool(const Vector &, const Vector &)> convergence = [&scalvar](const Vector &x, const Vector &residual)
  {
    return scalvar(x.array());
  };
 
  auto fb = sopt::algorithm::ImagingForwardBackward<Scalar>(y);
  fb.itermax(500)
    .step_size(beta)    // stepsize
    .sigma(sigma)  // sigma
    .regulariser_strength(regulariser_strength)  // regularisation paramater
    .relative_variation(1e-3)
    .residual_tolerance(0)
    .tight_frame(true)
    .Phi(sampling)
    .is_converged(convergence);
  
  //fb.set_f_gradient(f_gradient);
  fb.f_function(diff_function);
 
  // Create a shared pointer to the real indicator (non differentiable) function
  auto non_diff_func = std::make_shared<sopt::algorithm::RealIndicator<Scalar>>();
 
  // Inject it into the ImagingForwardBackward object
  fb.g_function(non_diff_func);
 
  auto const diagnostic = fb();
 
  // CHECK(diagnostic.good);
  // CHECK(diagnostic.niters < 500);
 
  // compare input image to cleaned output image
  // calculate mean squared error sum_i ( ( x_true(i) - x_est(i) ) **2 )
  // check this is less than the number of pixels * 0.01
  //sopt::utilities::write_tiff(Matrix::Map(diagnostic.x.data(), image.rows(), image.cols()),
  //                              "onnx_reconstruction.tiff");
  //sopt::utilities::write_tiff(Matrix::Map(dirty_image.data(), image.rows(), image.cols()),
  //                              "dirty.tiff");
  Eigen::Map<const Eigen::VectorXd> flat_image(image.data(), image.size());
  auto mse = (flat_image - diagnostic.x).array().square().sum() / image.size();
  CAPTURE(mse);
  SOPT_HIGH_LOG("MSE: {}", mse);
  CHECK(mse < 0.01);
}

References sopt::LinearTransform< VECTOR >::adjoint(), sopt::epsilon(), mersenne(), sopt::tools::models_directory(), sopt::tools::read_standard_tiff(), sopt::sigma(), and SOPT_HIGH_LOG.