diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp index 100c362007..e6538e18ea 100644 --- a/cmd/dwidenoise.cpp +++ b/cmd/dwidenoise.cpp @@ -26,6 +26,10 @@ using namespace App; const std::vector dtypes = {"float32", "float64"}; const std::vector estimators = {"exp1", "exp2"}; +const std::vector shapes = {"cuboid", "sphere"}; +enum class shape_type { CUBOID, SPHERE }; +constexpr default_type sphere_multiplier_default = 1.1; + // clang-format off void usage() { @@ -49,7 +53,29 @@ void usage() { + "Note that this function does not correct for non-Gaussian noise biases" " present in magnitude-reconstructed MRI images." " If available, including the MRI phase data can reduce such non-Gaussian biases," - " and the command now supports complex input data."; + " and the command now supports complex input data." + + + "The sliding spatial window behaves differently at the edges of the image FoV " + "depending on the shape / size selected for that window. " + "The default behaviour is to use a spherical kernel centred at the voxel of interest, " + "whose size is some multiple of the number of input volumes; " + "where some such voxels lie outside of the image FoV, " + "the radius of the kernel will be increased until the requisite number of voxels are used. " + "For a spherical kernel of a fixed radius, " + "no such expansion will occur, " + "and so for voxels near the image edge a reduced number of voxels will be present in the kernel. " + "For a cuboid kernel, " + "the centre of the kernel will be offset from the voxel being processed " + "such that the entire volume of the kernel resides within the image FoV." + + + "The size of the default spherical kernel is set to select a number of voxels that is " + "1.1 times the number of volumes in the input series. " + "If a cuboid kernel is requested, " + "but the -extent option is not specified, " + "the command will select the smallest isotropic patch size " + "that exceeds the number of DW images in the input data; " + "e.g., 5x5x5 for data with <= 125 DWI volumes, " + "7x7x7 for data with <= 343 DWI volumes, etc."; AUTHOR = "Daan Christiaens (daan.christiaens@kcl.ac.uk)" " and Jelle Veraart (jelle.veraart@nyumc.org)" @@ -73,19 +99,25 @@ void usage() { + Argument("out", "the output denoised DWI image.").type_image_out(); OPTIONS + + OptionGroup("Options for modifying the application of PCA denoising") + Option("mask", "Only process voxels within the specified binary brain mask image.") + Argument("image").type_image_in() + + Option("datatype", + "Datatype for the eigenvalue decomposition" + " (single or double precision). " + "For complex input data," + " this will select complex float32 or complex float64 datatypes.") + + Argument("float32/float64").type_choice(dtypes) + + Option("estimator", + "Select the noise level estimator" + " (default = Exp2)," + " either: \n" + "* Exp1: the original estimator used in Veraart et al. (2016), or \n" + "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019).") + + Argument("Exp1/Exp2").type_choice(estimators) - + Option("extent", - "Set the patch size of the denoising filter." - " By default, the command will select the smallest isotropic patch size" - " that exceeds the number of DW images in the input data," - " e.g., 5x5x5 for data with <= 125 DWI volumes," - " 7x7x7 for data with <= 343 DWI volumes," - " etc.") - + Argument("window").type_sequence_int() - + + OptionGroup("Options for exporting additional data regarding PCA behaviour") + Option("noise", "The output noise map," " i.e., the estimated noise level 'sigma' in the data." @@ -93,25 +125,33 @@ void usage() { " this will be the total noise level across real and imaginary channels," " so a scale factor sqrt(2) applies.") + Argument("level").type_image_out() - + Option("rank", "The selected signal rank of the output denoised image.") + Argument("cutoff").type_image_out() - + Option("datatype", - "Datatype for the eigenvalue decomposition" - " (single or double precision). " - "For complex input data," - " this will select complex float32 or complex float64 datatypes.") - + Argument("float32/float64").type_choice(dtypes) + + OptionGroup("Options for controlling the sliding spatial window") + + Option("shape", + "Set the shape of the sliding spatial window. " + "Options are: " + join(shapes, ",") + "; default: sphere") + + Argument("choice").type_choice(shapes) + + Option("radius_mm", + "Set an absolute spherical kernel radius in mm") + + Argument("value").type_float(0.0) + + Option("radius_ratio", + "Set the spherical kernel radius as a ratio of number of input volumes " + "(default: 1.1)") + + Argument("value").type_float(0.0) + + + Option("extent", + "Set the patch size of the cuboid filter; " + "can be either a single odd integer or a comma-separated triplet of odd integers") + + Argument("window").type_sequence_int(); + + + + + - + Option("estimator", - "Select the noise level estimator" - " (default = Exp2)," - " either: \n" - "* Exp1: the original estimator used in Veraart et al. (2016), or \n" - "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019).") - + Argument("Exp1/Exp2").type_choice(estimators); COPYRIGHT = "Copyright (c) 2016 New York University, University of Antwerp, and the MRtrix3 contributors \n \n" @@ -160,6 +200,7 @@ template class KernelBase { public: KernelBase() : pos({-1, -1, -1}) {} KernelBase(const KernelBase &) : pos({-1, -1, -1}) {} + virtual ssize_t minimum_size() const = 0; protected: // Store / restore position of image before / after data loading @@ -185,7 +226,7 @@ template class KernelCube : public KernelBase { } KernelCube(const KernelCube &) = default; template void operator()(ImageType &image, KernelData &data) { - assert(data.X.cols() == size()); + assert(data.X.cols() == minimum_size()); KernelBase::stash_pos(image); size_t k = 0; for (int z = -half_extent[2]; z <= half_extent[2]; z++) { @@ -199,10 +240,12 @@ template class KernelCube : public KernelBase { } } KernelBase::restore_pos(image); - data.voxel_count = size(); - data.centre_index = size() / 2; + data.voxel_count = minimum_size(); + data.centre_index = minimum_size() / 2; + } + ssize_t minimum_size() const override { + return (2 * half_extent[0] + 1) * (2 * half_extent[1] + 1) * (2 * half_extent[2] + 1); } - size_t size() const { return (2 * half_extent[0] + 1) * (2 * half_extent[1] + 1) * (2 * half_extent[2] + 1); } private: const std::vector half_extent; @@ -218,6 +261,106 @@ template class KernelCube : public KernelBase { } }; +template class KernelSphereBase : public KernelBase { +public: + KernelSphereBase(const Header &voxel_grid, const default_type max_radius) + : shared(new Shared(voxel_grid, max_radius)) {} + +protected: + class Shared { + public: + using MapType = std::multimap>; + Shared(const Header &voxel_grid, const default_type max_radius) { + const default_type max_radius_sq = Math::pow2(max_radius); + const std::array half_extents({int(std::ceil(max_radius / voxel_grid.spacing(0))), // + int(std::ceil(max_radius / voxel_grid.spacing(1))), // + int(std::ceil(max_radius / voxel_grid.spacing(2)))}); // + // Build the searchlight + std::array offset; + for (offset[2] = -half_extents[2]; offset[2] <= half_extents[2]; ++offset[2]) { + for (offset[1] = -half_extents[1]; offset[1] <= half_extents[1]; ++offset[1]) { + for (offset[0] = -half_extents[0]; offset[0] <= half_extents[0]; ++offset[0]) { + const default_type squared_distance = Math::pow2(offset[0] * voxel_grid.spacing(0)) // + + Math::pow2(offset[1] * voxel_grid.spacing(1)) // + + Math::pow2(offset[2] * voxel_grid.spacing(2)); // + if (squared_distance <= max_radius_sq) + data.insert({squared_distance, offset}); + } + } + } + } + MapType::const_iterator begin() const { return data.begin(); } + MapType::const_iterator end() const { return data.end(); } + + private: + MapType data; + }; + std::shared_ptr shared; +}; + +template class KernelSphereRatio : public KernelSphereBase { +public: + KernelSphereRatio(const Header &voxel_grid, const default_type min_ratio) + : KernelSphereBase(voxel_grid, compute_max_radius(voxel_grid, min_ratio)), + min_size(std::ceil(voxel_grid.size(3) * min_ratio)) {} + template void operator()(ImageType &image, KernelData &data) { + KernelBase::stash_pos(image); + data.voxel_count = 0; + default_type prev_distance = -std::numeric_limits::infinity(); + auto map_it = KernelSphereBase::shared->begin(); + while (map_it != KernelSphereBase::shared->end()) { + // If there's a tie in distances, want to include all such offsets in the kernel, + // even if the size of the utilised kernel extends beyond the minimum size + if (map_it->first != prev_distance && data.voxel_count >= min_size) + break; + for (size_t axis = 0; axis != 3; ++axis) + image.index(axis) = KernelBase::pos[axis] + map_it->second[axis]; + if (!is_out_of_bounds(image, 0, 3)) { + // Is this larger than any kernel this thread has previously encountered? + // If so, try to project what the final size is going to be, + // based on the set of voxels with identical distance to this one + // all getting included in the kernel + if (data.voxel_count == data.X.cols()) { + size_t extra_cols = 1; + auto forward_search = map_it; + for (++forward_search; + forward_search != KernelSphereBase::shared->end() && forward_search->first == map_it->first; + ++forward_search) + ++extra_cols; + data.X.conservativeResize(data.X.rows(), data.voxel_count + extra_cols); + } + data.X.col(data.voxel_count) = image.row(3); + prev_distance = map_it->first; + ++data.voxel_count; + } + ++map_it; + } + if (map_it == KernelSphereBase::shared->end()) + throw Exception("Inadequate spherical kernel initialisation"); + KernelBase::restore_pos(image); + data.centre_index = 0; + } + ssize_t minimum_size() const override { return min_size; } + +private: + size_t min_size; + // Determine an appropriate bounding box from which to generate the search table + // Find the radius for which 7/8 of the sphere will contain the minimum number of voxels, then round up + // This is only for setting the maximal radius for generation of the lookup table + default_type compute_max_radius(const Header &voxel_grid, const default_type min_ratio) const { + const size_t num_volumes = voxel_grid.size(3); + const default_type voxel_volume = voxel_grid.spacing(0) * voxel_grid.spacing(1) * voxel_grid.spacing(2); + const default_type sphere_volume = 8.0 * num_volumes * min_ratio * voxel_volume; + const default_type approx_radius = std::sqrt(sphere_volume * 0.75 / Math::pi); + const std::array half_extents({int(std::ceil(approx_radius / voxel_grid.spacing(0))), // + int(std::ceil(approx_radius / voxel_grid.spacing(1))), // + int(std::ceil(approx_radius / voxel_grid.spacing(2)))}); // + return std::max({half_extents[0] * voxel_grid.spacing(0), + half_extents[1] * voxel_grid.spacing(1), + half_extents[2] * voxel_grid.spacing(2)}); + } +}; + template class DenoisingFunctor { public: @@ -226,16 +369,13 @@ template class DenoisingFunctor { DenoisingFunctor( int ndwi, KernelType &kernel, Image &mask, Image &noise, Image &rank, bool exp1) - : data(ndwi, kernel.size()), + : data(ndwi, kernel.minimum_size()), kernel(kernel), m(ndwi), - n(kernel.size()), - r(std::min(m, n)), - q(std::max(m, n)), exp1(exp1), - XtX(r, r), - eig(r), - s(r), + XtX(std::min(m, kernel.minimum_size()), std::min(m, kernel.minimum_size())), + eig(std::min(m, kernel.minimum_size())), + s(std::min(m, kernel.minimum_size())), mask(mask), noise(noise), rankmap(rank) {} @@ -250,15 +390,28 @@ template class DenoisingFunctor { // Load data in local window kernel(dwi, data); + auto X = data.X.leftCols(data.voxel_count); + + const ssize_t n = data.voxel_count; + const ssize_t r = std::min(m, n); + const ssize_t q = std::max(m, n); + + if (r > XtX.rows()) { + XtX.resize(r, r); + s.resize(r); + } + + // TODO Fill matrices with NaN when in debug mode; + // make sure results from one voxel are not creeping into another // Compute Eigendecomposition: if (m <= n) - XtX.template triangularView() = data.X * data.X.adjoint(); + XtX.topLeftCorner(r, r).template triangularView() = X * X.adjoint(); else - XtX.template triangularView() = data.X.adjoint() * data.X; - eig.compute(XtX); + XtX.topLeftCorner(r, r).template triangularView() = X.adjoint() * X; + eig.compute(XtX.topLeftCorner(r, r)); // eigenvalues sorted in increasing order: - s = eig.eigenvalues().template cast(); + s.head(r) = eig.eigenvalues().template cast(); // Marchenko-Pastur optimal threshold const double lam_r = std::max(s[0], 0.0) / q; @@ -282,20 +435,18 @@ template class DenoisingFunctor { if (cutoff_p > 0) { // recombine data using only eigenvectors above threshold: s.head(cutoff_p).setZero(); - s.tail(r - cutoff_p).setOnes(); + s.segment(cutoff_p, r - cutoff_p).setOnes(); if (m <= n) - data.X.col(data.centre_index) = - eig.eigenvectors() * - (s.cast().asDiagonal() * (eig.eigenvectors().adjoint() * data.X.col(data.centre_index))); + X.col(data.centre_index) = eig.eigenvectors() * (s.head(r).cast().asDiagonal() * + (eig.eigenvectors().adjoint() * X.col(data.centre_index))); else - data.X.col(data.centre_index) = - data.X * - (eig.eigenvectors() * (s.cast().asDiagonal() * eig.eigenvectors().adjoint().col(data.centre_index))); + X.col(data.centre_index) = X * (eig.eigenvectors() * (s.head(r).cast().asDiagonal() * + eig.eigenvectors().adjoint().col(data.centre_index))); } // Store output assign_pos_of(dwi).to(out); - out.row(3) = data.X.col(data.centre_index); + out.row(3) = X.col(data.centre_index); // store noise map if requested: if (noise.valid()) { @@ -312,7 +463,7 @@ template class DenoisingFunctor { private: KernelData data; KernelType kernel; - const ssize_t m, n, r, q; + const ssize_t m; const bool exp1; MatrixType XtX; Eigen::SelfAdjointEigenSolver eig; @@ -348,40 +499,60 @@ void make_kernel(Header &data, Image &rank, const std::string &output_name, bool exp1) { - using KernelType = KernelCube>; + using MatrixType = Eigen::Matrix; - auto opt = get_options("extent"); - std::vector extent; - if (!opt.empty()) { - extent = parse_ints(opt[0][0]); - if (extent.size() == 1) - extent = {extent[0], extent[0], extent[0]}; - if (extent.size() != 3) - throw Exception("-extent must be either a scalar or a list of length 3"); - for (int i = 0; i < 3; i++) { - if (!(extent[i] & 1)) - throw Exception("-extent must be a (list of) odd numbers"); - if (extent[i] > data.size(i)) - throw Exception("-extent must not exceed the image dimensions"); + auto opt = get_options("shape"); + const shape_type shape = opt.empty() ? shape_type::SPHERE : shape_type((int)(opt[0][0])); + + switch (shape) { + case shape_type::SPHERE: { + // TODO Could infer that user wants a cuboid kernel if -extent is used, even if -shape is not + if (!get_options("extent").empty()) { + throw Exception("-extent option does not apply to spherical kernel"); } - } else { - uint32_t e = 1; - while (Math::pow3(e) < data.size(3)) - e += 2; - extent = {std::min(e, uint32_t(data.size(0))), // - std::min(e, uint32_t(data.size(1))), // - std::min(e, uint32_t(data.size(2)))}; // + const default_type min_ratio = get_option_value("-radius_ratio", sphere_multiplier_default); + KernelSphereRatio kernel(data, min_ratio); + process_image>(data, mask, noise, rank, output_name, kernel, exp1); + return; } - INFO("selected patch size: " + str(extent[0]) + " x " + str(extent[1]) + " x " + str(extent[2]) + "."); + case shape_type::CUBOID: { + opt = get_options("extent"); + std::vector extent; + if (!opt.empty()) { + extent = parse_ints(opt[0][0]); + if (extent.size() == 1) + extent = {extent[0], extent[0], extent[0]}; + if (extent.size() != 3) + throw Exception("-extent must be either a scalar or a list of length 3"); + for (int i = 0; i < 3; i++) { + if (!(extent[i] & 1)) + throw Exception("-extent must be a (list of) odd numbers"); + if (extent[i] > data.size(i)) + throw Exception("-extent must not exceed the image dimensions"); + } + } else { + uint32_t e = 1; + while (Math::pow3(e) < data.size(3)) + e += 2; + extent = {std::min(e, uint32_t(data.size(0))), // + std::min(e, uint32_t(data.size(1))), // + std::min(e, uint32_t(data.size(2)))}; // + } + INFO("selected patch size: " + str(extent[0]) + " x " + str(extent[1]) + " x " + str(extent[2]) + "."); - if (std::min(data.size(3), extent[0] * extent[1] * extent[2]) < 15) { - WARN("The number of volumes or the patch size is small. " - "This may lead to discretisation effects in the noise level " - "and cause inconsistent denoising between adjacent voxels."); - } + if (std::min(data.size(3), extent[0] * extent[1] * extent[2]) < 15) { + WARN("The number of volumes or the patch size is small. " + "This may lead to discretisation effects in the noise level " + "and cause inconsistent denoising between adjacent voxels."); + } - KernelType kernel(extent); - process_image(data, mask, noise, rank, output_name, kernel, exp1); + KernelCube kernel(extent); + process_image>(data, mask, noise, rank, output_name, kernel, exp1); + return; + } + default: + assert(false); + } } void run() {