From 3c8b48ef5b24f7f34ff196b4fafaef383c26224f Mon Sep 17 00:00:00 2001 From: Charles Poirier Date: Wed, 6 Oct 2021 16:03:21 -0400 Subject: [PATCH 01/20] Implement bilateral filtering with multivar Gaussian --- requirements.txt | 1 + scilpy/denoise/asym_enhancement.py | 132 ++++++++++++++----- scripts/scil_execute_asymmetric_filtering.py | 33 +++-- 3 files changed, 124 insertions(+), 42 deletions(-) diff --git a/requirements.txt b/requirements.txt index a494e2559..0ab3f6c99 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,6 +11,7 @@ matplotlib==2.2.* nibabel==3.0.* nilearn==0.6.* numpy==1.20.* +numba==0.54.* Pillow==8.2.* bids-validator==1.6.0 pybids==0.10.* diff --git a/scilpy/denoise/asym_enhancement.py b/scilpy/denoise/asym_enhancement.py index 6b061cf7b..4ce88c6eb 100644 --- a/scilpy/denoise/asym_enhancement.py +++ b/scilpy/denoise/asym_enhancement.py @@ -5,12 +5,18 @@ from dipy.data import get_sphere from dipy.core.sphere import Sphere from scipy.ndimage import correlate +from scipy.stats import multivariate_normal +from numba import jit -def local_asym_filtering(in_sh, sh_order=8, sh_basis='descoteaux07', - in_full_basis=False, out_full_basis=True, - dot_sharpness=1.0, sphere_str='repulsion724', - sigma=1.0): + +def multivariate_bilateral_filtering(in_sh, sh_order=8, + sh_basis='descoteaux07', + in_full_basis=False, + out_full_basis=True, + sphere_str='repulsion724', + var_cov=np.eye(2), + sigma_range=0.5): """Average the SH projected on a sphere using a first-neighbor gaussian blur and a dot product weight between sphere directions and the direction to neighborhood voxels, forcing to 0 negative values and thus performing @@ -33,8 +39,9 @@ def local_asym_filtering(in_sh, sh_order=8, sh_basis='descoteaux07', are not weighted by the dot product. sphere_str: str, optional Name of the sphere used to project SH coefficients to SF. - sigma: float, optional + sigma_spatial: float, optional Sigma for the Gaussian. + sigma_range: float, optional Returns ------- @@ -45,7 +52,7 @@ def local_asym_filtering(in_sh, sh_order=8, sh_basis='descoteaux07', sphere = get_sphere(sphere_str) # Normalized filter for each sf direction - weights = _get_weights(sphere, dot_sharpness, sigma) + weights = _get_weights_multivariate(sphere, var_cov) nb_sf = len(sphere.vertices) mean_sf = np.zeros(np.append(in_sh.shape[:-1], nb_sf)) @@ -62,24 +69,44 @@ def local_asym_filtering(in_sh, sh_order=8, sh_basis='descoteaux07', for sf_i in range(nb_sf): w_filter = weights[..., sf_i] - # Calculate contribution of center voxel + # Generate 1-channel images for directions u and -u current_sf = np.dot(in_sh, B[:, sf_i]) - mean_sf[..., sf_i] = w_filter[1, 1, 1] * current_sf + opposite_sf = np.dot(in_sh, neg_B[:, sf_i]) - # Add contributions of neighbors using opposite hemispheres - current_sf = np.dot(in_sh, neg_B[:, sf_i]) - w_filter[1, 1, 1] = 0.0 - mean_sf[..., sf_i] += correlate(current_sf, w_filter, mode="constant") + mean_sf[..., sf_i] = correlate_spatial(current_sf, + opposite_sf, + w_filter, + sigma_range) # Convert back to SH coefficients _, B_inv = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, full_basis=out_full_basis) - out_sh = np.array([np.dot(i, B_inv) for i in mean_sf], dtype=in_sh.dtype) + out_norm_L1 = np.sum(np.abs(out_sh), axis=-1) + in_norm_L1 = np.sum(np.abs(in_sh), axis=-1) + out_sh[out_norm_L1 > 0.] /= out_norm_L1[out_norm_L1 > 0.][..., None] + out_sh *= in_norm_L1[..., None] return out_sh -def _get_weights(sphere, dot_sharpness, sigma): +@jit(nopython=True) +def evaluate_gaussian_dist(x, sigma): + assert sigma > 0.0, "Sigma must be greater than 0." + cnorm = 1.0 / sigma / np.sqrt(2.0*np.pi) + return cnorm * np.exp(-x**2/2/sigma**2) + + +def evaluate_multivariate_gaussian(x, cov): + x_shape = x.shape + x = x.reshape((-1, 2)) + flat_out = multivariate_normal.pdf(x, mean=[0, 0], cov=cov) + + fx = flat_out.reshape(x_shape[:-1]) + + return fx + + +def _get_weights_multivariate(sphere, cov): """ Get neighbors weight in respect to the direction to a voxel. @@ -94,32 +121,75 @@ def _get_weights(sphere, dot_sharpness, sigma): Returns ------- - weights: dictionary + weights: ndarray Vertices weights with respect to voxel directions. """ - directions = np.zeros((3, 3, 3, 3)) - for x in range(3): - for y in range(3): - for z in range(3): - directions[x, y, z, 0] = x - 1 - directions[x, y, z, 1] = y - 1 - directions[x, y, z, 2] = z - 1 - - non_zero_dir = np.ones((3, 3, 3), dtype=bool) - non_zero_dir[1, 1, 1] = False + win_size = np.ceil(6.0 * cov[0, 0] + 1.0).astype(int) + if win_size % 2 == 0: + win_size += 1 + half_size = win_size // 2 + directions = np.zeros((win_size, win_size, win_size, 3)) + for x in range(win_size): + for y in range(win_size): + for z in range(win_size): + directions[x, y, z, 0] = x - half_size + directions[x, y, z, 1] = y - half_size + directions[x, y, z, 2] = z - half_size + + non_zero_dir = np.ones(directions.shape[:-1], dtype=bool) + non_zero_dir[half_size, half_size, half_size] = False # normalize dir dir_norm = np.linalg.norm(directions, axis=-1, keepdims=True) directions[non_zero_dir] /= dir_norm[non_zero_dir] - g_weights = np.exp(-dir_norm**2 / (2 * sigma**2)) - d_weights = np.dot(directions, sphere.vertices.T) + # angle in the range [0, pi] + angle = np.arccos(np.dot(directions, sphere.vertices.T)) + angle[half_size, half_size, half_size, :] = 0.0 - d_weights = np.where(d_weights > 0.0, d_weights**dot_sharpness, 0.0) - weights = d_weights * g_weights - weights[1, 1, 1, :] = 1.0 + dir_norm = np.broadcast_to(dir_norm, angle.shape) + norm_angles = np.stack([dir_norm, angle], axis=-1) + + weights = evaluate_multivariate_gaussian(norm_angles, cov) - # Normalize filters so that all sphere directions weights sum to 1 weights /= weights.reshape((-1, weights.shape[-1])).sum(axis=0) return weights + + +@jit(nopython=True) +def correlate_spatial(image_u, image_neg_u, h_filter, sigma_range): + """ + Implementation of correlate function. + """ + h_w, h_h, h_d = h_filter.shape[:3] + half_w, half_h, half_d = h_w // 2, h_h // 2, h_d // 2 + pad_img = np.zeros((image_neg_u.shape[0] + 2*half_w, + image_neg_u.shape[1] + 2*half_h, + image_neg_u.shape[2] + 2*half_d)) + pad_img[half_w:-half_w, half_h:-half_h, half_d:-half_d] = image_neg_u + + out_im = np.zeros_like(image_u) + for ii in range(image_u.shape[0]): + for jj in range(image_u.shape[1]): + for kk in range(image_u.shape[2]): + x = pad_img[ii:ii+h_w, jj:jj+h_h, kk:kk+h_d]\ + - image_u[ii, jj, kk] + range_filter = evaluate_gaussian_dist(x, sigma_range) + + res_filter = range_filter * h_filter + # Divide the filter into two filters; + # One for the current sphere direction and + # the other for the opposite direction. + res_filter_sum = np.sum(res_filter) + center_val = res_filter[half_w, half_h, half_d] + res_filter[half_w, half_h, half_d] = 0.0 + + out_im[ii, jj, kk] = image_u[ii, jj, kk] * center_val + out_im[ii, jj, kk] += np.sum(pad_img[ii:ii+h_w, + jj:jj+h_h, + kk:kk+h_d] + * res_filter) + out_im[ii, jj, kk] /= res_filter_sum + + return out_im diff --git a/scripts/scil_execute_asymmetric_filtering.py b/scripts/scil_execute_asymmetric_filtering.py index 000eec3c5..0a64de815 100755 --- a/scripts/scil_execute_asymmetric_filtering.py +++ b/scripts/scil_execute_asymmetric_filtering.py @@ -38,7 +38,7 @@ add_sh_basis_args, assert_outputs_exist) -from scilpy.denoise.asym_enhancement import local_asym_filtering +from scilpy.denoise.asym_enhancement import multivariate_bilateral_filtering def _build_arg_parser(): @@ -53,20 +53,28 @@ def _build_arg_parser(): add_sh_basis_args(p) + p.add_argument('--out_sym', action='store_true', + help='If set, saves output in symmetric SH basis.') + p.add_argument('--sphere', default='repulsion724', choices=sorted(SPHERE_FILES.keys()), help='Sphere used for the SH to SF projection. ' '[%(default)s]') - p.add_argument('--sharpness', default=1.0, type=float, - help='Specify sharpness factor to use for weighted average.' + p.add_argument('--sigma_angular', default=1.0, type=float, + help='Standard deviation for angular distance.' ' [%(default)s]') - p.add_argument('--sigma', default=1.0, type=float, - help='Sigma of the gaussian to use. [%(default)s]') + p.add_argument('--sigma_spatial', default=1.0, type=float, + help='Standard deviation for spatial distance.' + ' [%(default)s]') - p.add_argument('--out_sym', action='store_true', - help='If set, saves output in symmetric SH basis.') + p.add_argument('--sigma_range', default=1.0, type=float, + help='Standard deviation for intensity range.' + ' [%(default)s]') + + p.add_argument('--covariance', default=0.0, type=float, + help='Covariance of angular and spatial value.') add_verbose_arg(p) add_overwrite_arg(p) @@ -90,15 +98,18 @@ def main(): sh_order, full_basis = get_sh_order_and_fullness(data.shape[-1]) - logging.info('Executing local asymmetric filtering.') - filtered_sh = local_asym_filtering( + var_cov = np.array([[args.sigma_spatial**2, args.covariance], + [args.covariance, args.sigma_angular**2]]) + + logging.info('Executing asymmetric filtering.') + filtered_sh = multivariate_bilateral_filtering( data, sh_order=sh_order, sh_basis=args.sh_basis, in_full_basis=full_basis, out_full_basis=not(args.out_sym), sphere_str=args.sphere, - dot_sharpness=args.sharpness, - sigma=args.sigma) + var_cov=var_cov, + sigma_range=args.sigma_range) logging.info('Saving filtered SH to file {0}.'.format(args.out_sh)) nib.save(nib.Nifti1Image(filtered_sh, sh_img.affine), args.out_sh) From eb5febe81b677ff5e21b43d405f1cba18c15a6b2 Mon Sep 17 00:00:00 2001 From: Charles Poirier Date: Fri, 15 Oct 2021 09:38:00 -0400 Subject: [PATCH 02/20] Remove numba --- requirements.txt | 1 - scilpy/denoise/asym_enhancement.py | 5 ----- 2 files changed, 6 deletions(-) diff --git a/requirements.txt b/requirements.txt index 0ab3f6c99..a494e2559 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,7 +11,6 @@ matplotlib==2.2.* nibabel==3.0.* nilearn==0.6.* numpy==1.20.* -numba==0.54.* Pillow==8.2.* bids-validator==1.6.0 pybids==0.10.* diff --git a/scilpy/denoise/asym_enhancement.py b/scilpy/denoise/asym_enhancement.py index 4ce88c6eb..c42dd2f99 100644 --- a/scilpy/denoise/asym_enhancement.py +++ b/scilpy/denoise/asym_enhancement.py @@ -4,11 +4,8 @@ from dipy.reconst.shm import sh_to_sf_matrix from dipy.data import get_sphere from dipy.core.sphere import Sphere -from scipy.ndimage import correlate from scipy.stats import multivariate_normal -from numba import jit - def multivariate_bilateral_filtering(in_sh, sh_order=8, sh_basis='descoteaux07', @@ -89,7 +86,6 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, return out_sh -@jit(nopython=True) def evaluate_gaussian_dist(x, sigma): assert sigma > 0.0, "Sigma must be greater than 0." cnorm = 1.0 / sigma / np.sqrt(2.0*np.pi) @@ -157,7 +153,6 @@ def _get_weights_multivariate(sphere, cov): return weights -@jit(nopython=True) def correlate_spatial(image_u, image_neg_u, h_filter, sigma_range): """ Implementation of correlate function. From b2bfdef255f66d2defc4a68de2d0b354ce6ec861 Mon Sep 17 00:00:00 2001 From: Charles Poirier Date: Fri, 15 Oct 2021 09:42:24 -0400 Subject: [PATCH 03/20] Remove weird L1-normalization --- scilpy/denoise/asym_enhancement.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/scilpy/denoise/asym_enhancement.py b/scilpy/denoise/asym_enhancement.py index c42dd2f99..8b32f4999 100644 --- a/scilpy/denoise/asym_enhancement.py +++ b/scilpy/denoise/asym_enhancement.py @@ -79,10 +79,6 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, _, B_inv = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, full_basis=out_full_basis) out_sh = np.array([np.dot(i, B_inv) for i in mean_sf], dtype=in_sh.dtype) - out_norm_L1 = np.sum(np.abs(out_sh), axis=-1) - in_norm_L1 = np.sum(np.abs(in_sh), axis=-1) - out_sh[out_norm_L1 > 0.] /= out_norm_L1[out_norm_L1 > 0.][..., None] - out_sh *= in_norm_L1[..., None] return out_sh From 3bf0ddfecb01ee8d1bd5f5296c1d2783b7179aa6 Mon Sep 17 00:00:00 2001 From: Charles Poirier Date: Fri, 15 Oct 2021 12:03:00 -0400 Subject: [PATCH 04/20] Multiprocessing for filtering --- scilpy/denoise/asym_enhancement.py | 97 +++++++++++++++++--- scripts/scil_execute_asymmetric_filtering.py | 5 +- 2 files changed, 87 insertions(+), 15 deletions(-) diff --git a/scilpy/denoise/asym_enhancement.py b/scilpy/denoise/asym_enhancement.py index 8b32f4999..fff4379b7 100644 --- a/scilpy/denoise/asym_enhancement.py +++ b/scilpy/denoise/asym_enhancement.py @@ -1,6 +1,8 @@ # -*- coding: utf-8 -*- import numpy as np +import multiprocessing +import itertools from dipy.reconst.shm import sh_to_sf_matrix from dipy.data import get_sphere from dipy.core.sphere import Sphere @@ -13,7 +15,8 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, out_full_basis=True, sphere_str='repulsion724', var_cov=np.eye(2), - sigma_range=0.5): + sigma_range=0.5, + nbr_processes=1): """Average the SH projected on a sphere using a first-neighbor gaussian blur and a dot product weight between sphere directions and the direction to neighborhood voxels, forcing to 0 negative values and thus performing @@ -52,7 +55,6 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, weights = _get_weights_multivariate(sphere, var_cov) nb_sf = len(sphere.vertices) - mean_sf = np.zeros(np.append(in_sh.shape[:-1], nb_sf)) B = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, return_inv=False, full_basis=in_full_basis) @@ -62,18 +64,33 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, basis_type=sh_basis, return_inv=False, full_basis=in_full_basis) - # Apply filter to each sphere vertice - for sf_i in range(nb_sf): - w_filter = weights[..., sf_i] - - # Generate 1-channel images for directions u and -u - current_sf = np.dot(in_sh, B[:, sf_i]) - opposite_sf = np.dot(in_sh, neg_B[:, sf_i]) - - mean_sf[..., sf_i] = correlate_spatial(current_sf, - opposite_sf, - w_filter, - sigma_range) + if nbr_processes > 1: + # Apply filter to each sphere vertice in parallel + pool = multiprocessing.Pool(nbr_processes) + + # divide the sphere directions among the processes + base_chunk_size = int(nb_sf / nbr_processes + 0.5) + first_ids = np.arange(0, nb_sf, base_chunk_size) + residuals = nb_sf - first_ids + chunk_sizes = np.where(residuals < base_chunk_size, + residuals, base_chunk_size) + res = pool.map(_process_subset_directions, + zip(itertools.repeat(weights), + itertools.repeat(in_sh), + first_ids, + chunk_sizes, + itertools.repeat(B), + itertools.repeat(neg_B), + itertools.repeat(sigma_range))) + pool.close() + pool.join() + + # Patch chunks together. + mean_sf = np.concatenate(res, axis=-1) + else: + args = [weights, in_sh, 0, nb_sf, + B, neg_B, sigma_range] + mean_sf = _process_subset_directions(args) # Convert back to SH coefficients _, B_inv = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, @@ -149,6 +166,58 @@ def _get_weights_multivariate(sphere, cov): return weights +def _process_subset_directions(args): + """ + Filter a subset of all sphere directions. + + Parameters + ---------- + args: List + args[0]: weights, ndarray + Filter weights per direction. + args[1]: in_sh, ndarray + Input SH coefficients array. + args[2]: first_dir_id, int + ID of first sphere direction. + args[3]: chunk_size, int + Number of sphere directions in chunk. + args[4]: B, ndarray + SH to SF matrix for current sphere directions. + args[5]: neg_B, ndarray + SH to SF matrix for opposite sphere directions. + args[6]: sigma_range, int + Sigma of the Gaussian use for range filtering. + + Returns + ------- + out_sf: ndarray + SF array for subset directions. + """ + weights = args[0] + in_sh = args[1] + first_dir_id = args[2] + chunk_size = args[3] + B = args[4] + neg_B = args[5] + sigma_range = args[6] + + out_sf = np.zeros(in_sh.shape[:-1] + (chunk_size,)) + # Apply filter to each sphere vertice + for offset_i in range(chunk_size): + sph_id = first_dir_id + offset_i + w_filter = weights[..., sph_id] + + # Generate 1-channel images for directions u and -u + current_sf = np.dot(in_sh, B[:, sph_id]) + opposite_sf = np.dot(in_sh, neg_B[:, sph_id]) + + out_sf[..., offset_i] = correlate_spatial(current_sf, + opposite_sf, + w_filter, + sigma_range) + return out_sf + + def correlate_spatial(image_u, image_neg_u, h_filter, sigma_range): """ Implementation of correlate function. diff --git a/scripts/scil_execute_asymmetric_filtering.py b/scripts/scil_execute_asymmetric_filtering.py index 0a64de815..4e6416761 100755 --- a/scripts/scil_execute_asymmetric_filtering.py +++ b/scripts/scil_execute_asymmetric_filtering.py @@ -33,6 +33,7 @@ from dipy.data import SPHERE_FILES from scilpy.io.utils import (add_overwrite_arg, + add_processes_arg, add_verbose_arg, assert_inputs_exist, add_sh_basis_args, @@ -78,6 +79,7 @@ def _build_arg_parser(): add_verbose_arg(p) add_overwrite_arg(p) + add_processes_arg(p) return p @@ -109,7 +111,8 @@ def main(): out_full_basis=not(args.out_sym), sphere_str=args.sphere, var_cov=var_cov, - sigma_range=args.sigma_range) + sigma_range=args.sigma_range, + nbr_processes=args.nbr_processes) logging.info('Saving filtered SH to file {0}.'.format(args.out_sh)) nib.save(nib.Nifti1Image(filtered_sh, sh_img.affine), args.out_sh) From 87bf21049852880690cc97f1eb366c32457c3a5c Mon Sep 17 00:00:00 2001 From: "charles.poirier" Date: Wed, 27 Oct 2021 13:51:52 -0400 Subject: [PATCH 05/20] Split asymmetric averaging and bilateral filtering --- scilpy/denoise/asym_enhancement.py | 208 +++------------ scilpy/denoise/bilateral.py | 264 +++++++++++++++++++ scripts/scil_execute_asymmetric_filtering.py | 36 +-- scripts/scil_execute_bilateral_filtering.py | 133 ++++++++++ 4 files changed, 447 insertions(+), 194 deletions(-) create mode 100644 scilpy/denoise/bilateral.py create mode 100755 scripts/scil_execute_bilateral_filtering.py diff --git a/scilpy/denoise/asym_enhancement.py b/scilpy/denoise/asym_enhancement.py index fff4379b7..6b061cf7b 100644 --- a/scilpy/denoise/asym_enhancement.py +++ b/scilpy/denoise/asym_enhancement.py @@ -1,22 +1,16 @@ # -*- coding: utf-8 -*- import numpy as np -import multiprocessing -import itertools from dipy.reconst.shm import sh_to_sf_matrix from dipy.data import get_sphere from dipy.core.sphere import Sphere -from scipy.stats import multivariate_normal +from scipy.ndimage import correlate -def multivariate_bilateral_filtering(in_sh, sh_order=8, - sh_basis='descoteaux07', - in_full_basis=False, - out_full_basis=True, - sphere_str='repulsion724', - var_cov=np.eye(2), - sigma_range=0.5, - nbr_processes=1): +def local_asym_filtering(in_sh, sh_order=8, sh_basis='descoteaux07', + in_full_basis=False, out_full_basis=True, + dot_sharpness=1.0, sphere_str='repulsion724', + sigma=1.0): """Average the SH projected on a sphere using a first-neighbor gaussian blur and a dot product weight between sphere directions and the direction to neighborhood voxels, forcing to 0 negative values and thus performing @@ -39,9 +33,8 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, are not weighted by the dot product. sphere_str: str, optional Name of the sphere used to project SH coefficients to SF. - sigma_spatial: float, optional + sigma: float, optional Sigma for the Gaussian. - sigma_range: float, optional Returns ------- @@ -52,9 +45,10 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, sphere = get_sphere(sphere_str) # Normalized filter for each sf direction - weights = _get_weights_multivariate(sphere, var_cov) + weights = _get_weights(sphere, dot_sharpness, sigma) nb_sf = len(sphere.vertices) + mean_sf = np.zeros(np.append(in_sh.shape[:-1], nb_sf)) B = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, return_inv=False, full_basis=in_full_basis) @@ -64,58 +58,28 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, basis_type=sh_basis, return_inv=False, full_basis=in_full_basis) - if nbr_processes > 1: - # Apply filter to each sphere vertice in parallel - pool = multiprocessing.Pool(nbr_processes) - - # divide the sphere directions among the processes - base_chunk_size = int(nb_sf / nbr_processes + 0.5) - first_ids = np.arange(0, nb_sf, base_chunk_size) - residuals = nb_sf - first_ids - chunk_sizes = np.where(residuals < base_chunk_size, - residuals, base_chunk_size) - res = pool.map(_process_subset_directions, - zip(itertools.repeat(weights), - itertools.repeat(in_sh), - first_ids, - chunk_sizes, - itertools.repeat(B), - itertools.repeat(neg_B), - itertools.repeat(sigma_range))) - pool.close() - pool.join() - - # Patch chunks together. - mean_sf = np.concatenate(res, axis=-1) - else: - args = [weights, in_sh, 0, nb_sf, - B, neg_B, sigma_range] - mean_sf = _process_subset_directions(args) + # Apply filter to each sphere vertice + for sf_i in range(nb_sf): + w_filter = weights[..., sf_i] + + # Calculate contribution of center voxel + current_sf = np.dot(in_sh, B[:, sf_i]) + mean_sf[..., sf_i] = w_filter[1, 1, 1] * current_sf + + # Add contributions of neighbors using opposite hemispheres + current_sf = np.dot(in_sh, neg_B[:, sf_i]) + w_filter[1, 1, 1] = 0.0 + mean_sf[..., sf_i] += correlate(current_sf, w_filter, mode="constant") # Convert back to SH coefficients _, B_inv = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, full_basis=out_full_basis) + out_sh = np.array([np.dot(i, B_inv) for i in mean_sf], dtype=in_sh.dtype) return out_sh -def evaluate_gaussian_dist(x, sigma): - assert sigma > 0.0, "Sigma must be greater than 0." - cnorm = 1.0 / sigma / np.sqrt(2.0*np.pi) - return cnorm * np.exp(-x**2/2/sigma**2) - - -def evaluate_multivariate_gaussian(x, cov): - x_shape = x.shape - x = x.reshape((-1, 2)) - flat_out = multivariate_normal.pdf(x, mean=[0, 0], cov=cov) - - fx = flat_out.reshape(x_shape[:-1]) - - return fx - - -def _get_weights_multivariate(sphere, cov): +def _get_weights(sphere, dot_sharpness, sigma): """ Get neighbors weight in respect to the direction to a voxel. @@ -130,126 +94,32 @@ def _get_weights_multivariate(sphere, cov): Returns ------- - weights: ndarray + weights: dictionary Vertices weights with respect to voxel directions. """ - win_size = np.ceil(6.0 * cov[0, 0] + 1.0).astype(int) - if win_size % 2 == 0: - win_size += 1 - half_size = win_size // 2 - directions = np.zeros((win_size, win_size, win_size, 3)) - for x in range(win_size): - for y in range(win_size): - for z in range(win_size): - directions[x, y, z, 0] = x - half_size - directions[x, y, z, 1] = y - half_size - directions[x, y, z, 2] = z - half_size - - non_zero_dir = np.ones(directions.shape[:-1], dtype=bool) - non_zero_dir[half_size, half_size, half_size] = False + directions = np.zeros((3, 3, 3, 3)) + for x in range(3): + for y in range(3): + for z in range(3): + directions[x, y, z, 0] = x - 1 + directions[x, y, z, 1] = y - 1 + directions[x, y, z, 2] = z - 1 + + non_zero_dir = np.ones((3, 3, 3), dtype=bool) + non_zero_dir[1, 1, 1] = False # normalize dir dir_norm = np.linalg.norm(directions, axis=-1, keepdims=True) directions[non_zero_dir] /= dir_norm[non_zero_dir] - # angle in the range [0, pi] - angle = np.arccos(np.dot(directions, sphere.vertices.T)) - angle[half_size, half_size, half_size, :] = 0.0 - - dir_norm = np.broadcast_to(dir_norm, angle.shape) - norm_angles = np.stack([dir_norm, angle], axis=-1) + g_weights = np.exp(-dir_norm**2 / (2 * sigma**2)) + d_weights = np.dot(directions, sphere.vertices.T) - weights = evaluate_multivariate_gaussian(norm_angles, cov) + d_weights = np.where(d_weights > 0.0, d_weights**dot_sharpness, 0.0) + weights = d_weights * g_weights + weights[1, 1, 1, :] = 1.0 + # Normalize filters so that all sphere directions weights sum to 1 weights /= weights.reshape((-1, weights.shape[-1])).sum(axis=0) return weights - - -def _process_subset_directions(args): - """ - Filter a subset of all sphere directions. - - Parameters - ---------- - args: List - args[0]: weights, ndarray - Filter weights per direction. - args[1]: in_sh, ndarray - Input SH coefficients array. - args[2]: first_dir_id, int - ID of first sphere direction. - args[3]: chunk_size, int - Number of sphere directions in chunk. - args[4]: B, ndarray - SH to SF matrix for current sphere directions. - args[5]: neg_B, ndarray - SH to SF matrix for opposite sphere directions. - args[6]: sigma_range, int - Sigma of the Gaussian use for range filtering. - - Returns - ------- - out_sf: ndarray - SF array for subset directions. - """ - weights = args[0] - in_sh = args[1] - first_dir_id = args[2] - chunk_size = args[3] - B = args[4] - neg_B = args[5] - sigma_range = args[6] - - out_sf = np.zeros(in_sh.shape[:-1] + (chunk_size,)) - # Apply filter to each sphere vertice - for offset_i in range(chunk_size): - sph_id = first_dir_id + offset_i - w_filter = weights[..., sph_id] - - # Generate 1-channel images for directions u and -u - current_sf = np.dot(in_sh, B[:, sph_id]) - opposite_sf = np.dot(in_sh, neg_B[:, sph_id]) - - out_sf[..., offset_i] = correlate_spatial(current_sf, - opposite_sf, - w_filter, - sigma_range) - return out_sf - - -def correlate_spatial(image_u, image_neg_u, h_filter, sigma_range): - """ - Implementation of correlate function. - """ - h_w, h_h, h_d = h_filter.shape[:3] - half_w, half_h, half_d = h_w // 2, h_h // 2, h_d // 2 - pad_img = np.zeros((image_neg_u.shape[0] + 2*half_w, - image_neg_u.shape[1] + 2*half_h, - image_neg_u.shape[2] + 2*half_d)) - pad_img[half_w:-half_w, half_h:-half_h, half_d:-half_d] = image_neg_u - - out_im = np.zeros_like(image_u) - for ii in range(image_u.shape[0]): - for jj in range(image_u.shape[1]): - for kk in range(image_u.shape[2]): - x = pad_img[ii:ii+h_w, jj:jj+h_h, kk:kk+h_d]\ - - image_u[ii, jj, kk] - range_filter = evaluate_gaussian_dist(x, sigma_range) - - res_filter = range_filter * h_filter - # Divide the filter into two filters; - # One for the current sphere direction and - # the other for the opposite direction. - res_filter_sum = np.sum(res_filter) - center_val = res_filter[half_w, half_h, half_d] - res_filter[half_w, half_h, half_d] = 0.0 - - out_im[ii, jj, kk] = image_u[ii, jj, kk] * center_val - out_im[ii, jj, kk] += np.sum(pad_img[ii:ii+h_w, - jj:jj+h_h, - kk:kk+h_d] - * res_filter) - out_im[ii, jj, kk] /= res_filter_sum - - return out_im diff --git a/scilpy/denoise/bilateral.py b/scilpy/denoise/bilateral.py new file mode 100644 index 000000000..06b9b43ef --- /dev/null +++ b/scilpy/denoise/bilateral.py @@ -0,0 +1,264 @@ +# -*- coding: utf-8 -*- + +import numpy as np +import multiprocessing +import itertools +from dipy.reconst.shm import sh_to_sf_matrix +from dipy.data import get_sphere +from dipy.core.sphere import Sphere +from scipy.stats import multivariate_normal + + +def multivariate_bilateral_filtering(in_sh, sh_order=8, + sh_basis='descoteaux07', + in_full_basis=False, + return_sym=False, + sphere_str='repulsion724', + var_cov=np.eye(2), + sigma_range=0.5, + nbr_processes=1): + """Average the SH projected on a sphere using a first-neighbor gaussian + blur and a dot product weight between sphere directions and the direction + to neighborhood voxels, forcing to 0 negative values and thus performing + asymmetric hemisphere-aware filtering. + + Parameters + ---------- + in_sh: ndarray (x, y, z, n_coeffs) + Input SH coefficients array + sh_order: int, optional + Maximum order of the SH series. + sh_basis: {'descoteaux07', 'tournier07'}, optional + SH basis of the input signal. + in_full_basis: bool, optional + True if the input is in full SH basis. + out_full_basis: bool, optional + If True, save output SH using full SH basis. + dot_sharpness: float, optional + Exponent of the dot product. When set to 0.0, directions + are not weighted by the dot product. + sphere_str: str, optional + Name of the sphere used to project SH coefficients to SF. + sigma_spatial: float, optional + Sigma for the Gaussian. + sigma_range: float, optional + + Returns + ------- + out_sh: ndarray (x, y, z, n_coeffs) + Filtered signal as SH coefficients. + """ + # Load the sphere used for projection of SH + sphere = get_sphere(sphere_str) + + # Normalized filter for each sf direction + weights = _get_weights_multivariate(sphere, var_cov) + + nb_sf = len(sphere.vertices) + B = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, + return_inv=False, full_basis=in_full_basis) + + # We want a B matrix to project on an inverse sphere to have the sf on + # the opposite hemisphere for a given vertice + neg_B = sh_to_sf_matrix(Sphere(xyz=-sphere.vertices), sh_order=sh_order, + basis_type=sh_basis, return_inv=False, + full_basis=in_full_basis) + + if nbr_processes > 1: + # Apply filter to each sphere vertice in parallel + pool = multiprocessing.Pool(nbr_processes) + + # divide the sphere directions among the processes + base_chunk_size = int(nb_sf / nbr_processes + 0.5) + first_ids = np.arange(0, nb_sf, base_chunk_size) + residuals = nb_sf - first_ids + chunk_sizes = np.where(residuals < base_chunk_size, + residuals, base_chunk_size) + res = pool.map(_process_subset_directions, + zip(itertools.repeat(weights), + itertools.repeat(in_sh), + first_ids, + chunk_sizes, + itertools.repeat(B), + itertools.repeat(neg_B), + itertools.repeat(sigma_range))) + pool.close() + pool.join() + + # Patch chunks together. + mean_sf = np.concatenate(res, axis=-1) + else: + args = [weights, in_sh, 0, nb_sf, + B, neg_B, sigma_range] + mean_sf = _process_subset_directions(args) + + # Convert back to SH coefficients + _, B_inv = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, + full_basis=True) + out_sh = np.array([np.dot(i, B_inv) for i in mean_sf], dtype=in_sh.dtype) + if return_sym: + _, B_inv = sh_to_sf_matrix(sphere, sh_order=sh_order, + basis_type=sh_basis, + full_basis=False) + out_sh_sym = np.array([np.dot(i, B_inv) for i in mean_sf], + dtype=in_sh.dtype) + return out_sh, out_sh_sym + + # By default, return only asymmetric SH + return out_sh, None + + +def evaluate_gaussian_dist(x, sigma): + assert sigma > 0.0, "Sigma must be greater than 0." + cnorm = 1.0 / sigma / np.sqrt(2.0*np.pi) + return cnorm * np.exp(-x**2/2/sigma**2) + + +def evaluate_multivariate_gaussian(x, cov): + x_shape = x.shape + x = x.reshape((-1, 2)) + flat_out = multivariate_normal.pdf(x, mean=[0, 0], cov=cov) + + fx = flat_out.reshape(x_shape[:-1]) + + return fx + + +def _get_weights_multivariate(sphere, cov): + """ + Get neighbors weight in respect to the direction to a voxel. + + Parameters + ---------- + sphere: Sphere + Sphere used for SF reconstruction. + dot_sharpness: float + Dot product exponent. + sigma: float + Variance of the gaussian used for weighting neighbors. + + Returns + ------- + weights: ndarray + Vertices weights with respect to voxel directions. + """ + win_size = np.ceil(6.0 * cov[0, 0] + 1.0).astype(int) + if win_size % 2 == 0: + win_size += 1 + half_size = win_size // 2 + directions = np.zeros((win_size, win_size, win_size, 3)) + for x in range(win_size): + for y in range(win_size): + for z in range(win_size): + directions[x, y, z, 0] = x - half_size + directions[x, y, z, 1] = y - half_size + directions[x, y, z, 2] = z - half_size + + non_zero_dir = np.ones(directions.shape[:-1], dtype=bool) + non_zero_dir[half_size, half_size, half_size] = False + + # normalize dir + dir_norm = np.linalg.norm(directions, axis=-1, keepdims=True) + directions[non_zero_dir] /= dir_norm[non_zero_dir] + + # angle in the range [0, pi] + angle = np.arccos(np.dot(directions, sphere.vertices.T)) + angle[half_size, half_size, half_size, :] = 0.0 + + dir_norm = np.broadcast_to(dir_norm, angle.shape) + norm_angles = np.stack([dir_norm, angle], axis=-1) + + weights = evaluate_multivariate_gaussian(norm_angles, cov) + + weights /= weights.reshape((-1, weights.shape[-1])).sum(axis=0) + + return weights + + +def _process_subset_directions(args): + """ + Filter a subset of all sphere directions. + + Parameters + ---------- + args: List + args[0]: weights, ndarray + Filter weights per direction. + args[1]: in_sh, ndarray + Input SH coefficients array. + args[2]: first_dir_id, int + ID of first sphere direction. + args[3]: chunk_size, int + Number of sphere directions in chunk. + args[4]: B, ndarray + SH to SF matrix for current sphere directions. + args[5]: neg_B, ndarray + SH to SF matrix for opposite sphere directions. + args[6]: sigma_range, int + Sigma of the Gaussian use for range filtering. + + Returns + ------- + out_sf: ndarray + SF array for subset directions. + """ + weights = args[0] + in_sh = args[1] + first_dir_id = args[2] + chunk_size = args[3] + B = args[4] + neg_B = args[5] + sigma_range = args[6] + + out_sf = np.zeros(in_sh.shape[:-1] + (chunk_size,)) + # Apply filter to each sphere vertice + for offset_i in range(chunk_size): + sph_id = first_dir_id + offset_i + w_filter = weights[..., sph_id] + + # Generate 1-channel images for directions u and -u + current_sf = np.dot(in_sh, B[:, sph_id]) + opposite_sf = np.dot(in_sh, neg_B[:, sph_id]) + + out_sf[..., offset_i] = correlate_spatial(current_sf, + opposite_sf, + w_filter, + sigma_range) + return out_sf + + +def correlate_spatial(image_u, image_neg_u, h_filter, sigma_range): + """ + Implementation of correlate function. + """ + h_w, h_h, h_d = h_filter.shape[:3] + half_w, half_h, half_d = h_w // 2, h_h // 2, h_d // 2 + pad_img = np.zeros((image_neg_u.shape[0] + 2*half_w, + image_neg_u.shape[1] + 2*half_h, + image_neg_u.shape[2] + 2*half_d)) + pad_img[half_w:-half_w, half_h:-half_h, half_d:-half_d] = image_neg_u + + out_im = np.zeros_like(image_u) + for ii in range(image_u.shape[0]): + for jj in range(image_u.shape[1]): + for kk in range(image_u.shape[2]): + x = pad_img[ii:ii+h_w, jj:jj+h_h, kk:kk+h_d]\ + - image_u[ii, jj, kk] + range_filter = evaluate_gaussian_dist(x, sigma_range) + + res_filter = range_filter * h_filter + # Divide the filter into two filters; + # One for the current sphere direction and + # the other for the opposite direction. + res_filter_sum = np.sum(res_filter) + center_val = res_filter[half_w, half_h, half_d] + res_filter[half_w, half_h, half_d] = 0.0 + + out_im[ii, jj, kk] = image_u[ii, jj, kk] * center_val + out_im[ii, jj, kk] += np.sum(pad_img[ii:ii+h_w, + jj:jj+h_h, + kk:kk+h_d] + * res_filter) + out_im[ii, jj, kk] /= res_filter_sum + + return out_im diff --git a/scripts/scil_execute_asymmetric_filtering.py b/scripts/scil_execute_asymmetric_filtering.py index 4e6416761..000eec3c5 100755 --- a/scripts/scil_execute_asymmetric_filtering.py +++ b/scripts/scil_execute_asymmetric_filtering.py @@ -33,13 +33,12 @@ from dipy.data import SPHERE_FILES from scilpy.io.utils import (add_overwrite_arg, - add_processes_arg, add_verbose_arg, assert_inputs_exist, add_sh_basis_args, assert_outputs_exist) -from scilpy.denoise.asym_enhancement import multivariate_bilateral_filtering +from scilpy.denoise.asym_enhancement import local_asym_filtering def _build_arg_parser(): @@ -54,32 +53,23 @@ def _build_arg_parser(): add_sh_basis_args(p) - p.add_argument('--out_sym', action='store_true', - help='If set, saves output in symmetric SH basis.') - p.add_argument('--sphere', default='repulsion724', choices=sorted(SPHERE_FILES.keys()), help='Sphere used for the SH to SF projection. ' '[%(default)s]') - p.add_argument('--sigma_angular', default=1.0, type=float, - help='Standard deviation for angular distance.' + p.add_argument('--sharpness', default=1.0, type=float, + help='Specify sharpness factor to use for weighted average.' ' [%(default)s]') - p.add_argument('--sigma_spatial', default=1.0, type=float, - help='Standard deviation for spatial distance.' - ' [%(default)s]') + p.add_argument('--sigma', default=1.0, type=float, + help='Sigma of the gaussian to use. [%(default)s]') - p.add_argument('--sigma_range', default=1.0, type=float, - help='Standard deviation for intensity range.' - ' [%(default)s]') - - p.add_argument('--covariance', default=0.0, type=float, - help='Covariance of angular and spatial value.') + p.add_argument('--out_sym', action='store_true', + help='If set, saves output in symmetric SH basis.') add_verbose_arg(p) add_overwrite_arg(p) - add_processes_arg(p) return p @@ -100,19 +90,15 @@ def main(): sh_order, full_basis = get_sh_order_and_fullness(data.shape[-1]) - var_cov = np.array([[args.sigma_spatial**2, args.covariance], - [args.covariance, args.sigma_angular**2]]) - - logging.info('Executing asymmetric filtering.') - filtered_sh = multivariate_bilateral_filtering( + logging.info('Executing local asymmetric filtering.') + filtered_sh = local_asym_filtering( data, sh_order=sh_order, sh_basis=args.sh_basis, in_full_basis=full_basis, out_full_basis=not(args.out_sym), sphere_str=args.sphere, - var_cov=var_cov, - sigma_range=args.sigma_range, - nbr_processes=args.nbr_processes) + dot_sharpness=args.sharpness, + sigma=args.sigma) logging.info('Saving filtered SH to file {0}.'.format(args.out_sh)) nib.save(nib.Nifti1Image(filtered_sh, sh_img.affine), args.out_sh) diff --git a/scripts/scil_execute_bilateral_filtering.py b/scripts/scil_execute_bilateral_filtering.py new file mode 100755 index 000000000..768b9080e --- /dev/null +++ b/scripts/scil_execute_bilateral_filtering.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +Script to compute per-vertices hemisphere-aware (asymmetric) +filtering of spherical functions (SF) given an array of spherical harmonics +(SH) coefficients. SF are filtered using a first-neighbor Gaussian filter. +Sphere directions are also weighted by their dot product with the direction +to the center of each neighbor, clipping to 0 negative values. + +The argument `sigma` controls the standard deviation of the Gaussian. The +argument `sharpness` controls the exponent of the cosine weights. The higher it +is, the faster the weights of misaligned sphere directions decrease. A +sharpness of 0 gives the same weight to all sphere directions in an hemisphere. +Both `sharpness` and `sigma` must be positive. + +The resulting SF can be expressed using a full SH basis or a symmetric SH basis +(where the effect of the filtering is a simple denoising). + +Using default parameters, the script completes in about 15-20 minutes for a +HCP subject fiber ODF processed with tractoflow. Also note the bigger the +sphere used for SH to SF projection, the higher the RAM consumption and +compute time. +""" + +import argparse +import logging +from scilpy.reconst.utils import get_sh_order_and_fullness + +import nibabel as nib +import numpy as np + +from dipy.data import SPHERE_FILES + +from scilpy.io.utils import (add_overwrite_arg, + add_processes_arg, + add_verbose_arg, + assert_inputs_exist, + add_sh_basis_args, + assert_outputs_exist, + validate_nbr_processes) + +from scilpy.denoise.asym_enhancement import multivariate_bilateral_filtering + + +def _build_arg_parser(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawTextHelpFormatter) + + p.add_argument('in_sh', + help='Path to the input file.') + + p.add_argument('out_sh', + help='File name for averaged signal.') + + add_sh_basis_args(p) + + p.add_argument('--out_sym', default=None, + help='If set, saves additional output ' + 'in symmetric SH basis.') + + p.add_argument('--sphere', default='repulsion724', + choices=sorted(SPHERE_FILES.keys()), + help='Sphere used for the SH to SF projection. ' + '[%(default)s]') + + p.add_argument('--sigma_angular', default=1.0, type=float, + help='Standard deviation for angular distance.' + ' [%(default)s]') + + p.add_argument('--sigma_spatial', default=1.0, type=float, + help='Standard deviation for spatial distance.' + ' [%(default)s]') + + p.add_argument('--sigma_range', default=1.0, type=float, + help='Standard deviation for intensity range.' + ' [%(default)s]') + + p.add_argument('--covariance', default=0.0, type=float, + help='Covariance of angular and spatial value.') + + add_verbose_arg(p) + add_overwrite_arg(p) + add_processes_arg(p) + + return p + + +def main(): + parser = _build_arg_parser() + args = parser.parse_args() + if args.verbose: + logging.basicConfig(level=logging.INFO) + + # Checking args + outputs = [args.out_sh] + if args.out_sym: + outputs.append(args.out_sym) + assert_outputs_exist(parser, args, outputs) + assert_inputs_exist(parser, args.in_sh) + + validate_nbr_processes(parser, args) + + # Prepare data + sh_img = nib.load(args.in_sh) + data = sh_img.get_fdata(dtype=np.float32) + + sh_order, full_basis = get_sh_order_and_fullness(data.shape[-1]) + + var_cov = np.array([[args.sigma_spatial**2, args.covariance], + [args.covariance, args.sigma_angular**2]]) + + logging.info('Executing asymmetric filtering.') + asym_sh, sym_sh = multivariate_bilateral_filtering( + data, sh_order=sh_order, + sh_basis=args.sh_basis, + in_full_basis=full_basis, + return_sym=args.out_sym is not None, + sphere_str=args.sphere, + var_cov=var_cov, + sigma_range=args.sigma_range, + nbr_processes=args.nbr_processes) + + logging.info('Saving filtered SH to file {0}.'.format(args.out_sh)) + nib.save(nib.Nifti1Image(asym_sh, sh_img.affine), args.out_sh) + + if args.out_sym: + logging.info('Saving symmetric SH to file {0}.'.format(args.out_sym)) + nib.save(nib.Nifti1Image(sym_sh, sh_img.affine), args.out_sym) + + +if __name__ == "__main__": + main() From 136c331ed31f3fde88072f2b0a78dedbbcdf67ad Mon Sep 17 00:00:00 2001 From: "charles.poirier" Date: Wed, 27 Oct 2021 14:01:07 -0400 Subject: [PATCH 06/20] Rename filtering modules --- scilpy/denoise/{asym_enhancement.py => asym_averaging.py} | 0 scilpy/denoise/{bilateral.py => bilateral_filtering.py} | 0 scripts/scil_execute_asymmetric_filtering.py | 2 +- scripts/scil_execute_bilateral_filtering.py | 2 +- 4 files changed, 2 insertions(+), 2 deletions(-) rename scilpy/denoise/{asym_enhancement.py => asym_averaging.py} (100%) rename scilpy/denoise/{bilateral.py => bilateral_filtering.py} (100%) diff --git a/scilpy/denoise/asym_enhancement.py b/scilpy/denoise/asym_averaging.py similarity index 100% rename from scilpy/denoise/asym_enhancement.py rename to scilpy/denoise/asym_averaging.py diff --git a/scilpy/denoise/bilateral.py b/scilpy/denoise/bilateral_filtering.py similarity index 100% rename from scilpy/denoise/bilateral.py rename to scilpy/denoise/bilateral_filtering.py diff --git a/scripts/scil_execute_asymmetric_filtering.py b/scripts/scil_execute_asymmetric_filtering.py index 000eec3c5..1713de87c 100755 --- a/scripts/scil_execute_asymmetric_filtering.py +++ b/scripts/scil_execute_asymmetric_filtering.py @@ -38,7 +38,7 @@ add_sh_basis_args, assert_outputs_exist) -from scilpy.denoise.asym_enhancement import local_asym_filtering +from scilpy.denoise.asym_averaging import local_asym_filtering def _build_arg_parser(): diff --git a/scripts/scil_execute_bilateral_filtering.py b/scripts/scil_execute_bilateral_filtering.py index 768b9080e..9dd034119 100755 --- a/scripts/scil_execute_bilateral_filtering.py +++ b/scripts/scil_execute_bilateral_filtering.py @@ -40,7 +40,7 @@ assert_outputs_exist, validate_nbr_processes) -from scilpy.denoise.asym_enhancement import multivariate_bilateral_filtering +from scilpy.denoise.bilateral_filtering import multivariate_bilateral_filtering def _build_arg_parser(): From 89098853ffa8759d030668866978755549ab1eee Mon Sep 17 00:00:00 2001 From: "charles.poirier" Date: Fri, 5 Nov 2021 15:44:12 -0400 Subject: [PATCH 07/20] Get symmetric coefficients from full basis --- scilpy/denoise/bilateral_filtering.py | 29 ++++++--------------- scripts/scil_execute_bilateral_filtering.py | 8 +++--- 2 files changed, 13 insertions(+), 24 deletions(-) diff --git a/scilpy/denoise/bilateral_filtering.py b/scilpy/denoise/bilateral_filtering.py index 06b9b43ef..02c1d2b6b 100644 --- a/scilpy/denoise/bilateral_filtering.py +++ b/scilpy/denoise/bilateral_filtering.py @@ -12,15 +12,12 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, sh_basis='descoteaux07', in_full_basis=False, - return_sym=False, sphere_str='repulsion724', var_cov=np.eye(2), sigma_range=0.5, nbr_processes=1): - """Average the SH projected on a sphere using a first-neighbor gaussian - blur and a dot product weight between sphere directions and the direction - to neighborhood voxels, forcing to 0 negative values and thus performing - asymmetric hemisphere-aware filtering. + """ + Multivariate bilateral filtering. Parameters ---------- @@ -32,16 +29,14 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, SH basis of the input signal. in_full_basis: bool, optional True if the input is in full SH basis. - out_full_basis: bool, optional - If True, save output SH using full SH basis. - dot_sharpness: float, optional - Exponent of the dot product. When set to 0.0, directions - are not weighted by the dot product. sphere_str: str, optional Name of the sphere used to project SH coefficients to SF. - sigma_spatial: float, optional - Sigma for the Gaussian. + var_cov: ndarray (2, 2), optional + Variance-covariance matrix for spatio-augular distribution. sigma_range: float, optional + Variance of the gaussian used for weighting intensities. + nbr_processes: int, optional + Number of processes to use. Returns ------- @@ -96,16 +91,8 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, _, B_inv = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, full_basis=True) out_sh = np.array([np.dot(i, B_inv) for i in mean_sf], dtype=in_sh.dtype) - if return_sym: - _, B_inv = sh_to_sf_matrix(sphere, sh_order=sh_order, - basis_type=sh_basis, - full_basis=False) - out_sh_sym = np.array([np.dot(i, B_inv) for i in mean_sf], - dtype=in_sh.dtype) - return out_sh, out_sh_sym - # By default, return only asymmetric SH - return out_sh, None + return out_sh def evaluate_gaussian_dist(x, sigma): diff --git a/scripts/scil_execute_bilateral_filtering.py b/scripts/scil_execute_bilateral_filtering.py index 9dd034119..c320f1e5c 100755 --- a/scripts/scil_execute_bilateral_filtering.py +++ b/scripts/scil_execute_bilateral_filtering.py @@ -31,6 +31,7 @@ import numpy as np from dipy.data import SPHERE_FILES +from dipy.reconst.shm import sph_harm_ind_list from scilpy.io.utils import (add_overwrite_arg, add_processes_arg, @@ -111,11 +112,10 @@ def main(): [args.covariance, args.sigma_angular**2]]) logging.info('Executing asymmetric filtering.') - asym_sh, sym_sh = multivariate_bilateral_filtering( + asym_sh = multivariate_bilateral_filtering( data, sh_order=sh_order, sh_basis=args.sh_basis, in_full_basis=full_basis, - return_sym=args.out_sym is not None, sphere_str=args.sphere, var_cov=var_cov, sigma_range=args.sigma_range, @@ -125,8 +125,10 @@ def main(): nib.save(nib.Nifti1Image(asym_sh, sh_img.affine), args.out_sh) if args.out_sym: + _, orders = sph_harm_ind_list(sh_order, full_basis=True) logging.info('Saving symmetric SH to file {0}.'.format(args.out_sym)) - nib.save(nib.Nifti1Image(sym_sh, sh_img.affine), args.out_sym) + nib.save(nib.Nifti1Image(asym_sh[..., orders % 2 == 0], sh_img.affine), + args.out_sym) if __name__ == "__main__": From 5712fe207fe58af94297f524a8d3346d8d9d54a0 Mon Sep 17 00:00:00 2001 From: "charles.poirier" Date: Fri, 5 Nov 2021 15:44:40 -0400 Subject: [PATCH 08/20] Define correlate_sf_from_sh function --- scilpy/denoise/correlation.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 scilpy/denoise/correlation.py diff --git a/scilpy/denoise/correlation.py b/scilpy/denoise/correlation.py new file mode 100644 index 000000000..7f0a0bc5a --- /dev/null +++ b/scilpy/denoise/correlation.py @@ -0,0 +1,25 @@ +# -*- coding: utf-8 -*- + + +def correlate_sf_from_sh(in_sh, sphere, get_filter_func, sigma_range=None): + """ + Correlation operation on SF amplitudes from SH coefficients array. + + Parameters + ---------- + in_sh: ndarray (X, Y, Z, n_coeffs) + SH coefficients image to correlate. + sphere: dipy.sphere + Sphere used for SH to SF projection. + get_filter_func: Callable(direction) + Callback taking a sphere direction as input + and returning filter weights for this direction. + sigma_range: float or None, optional + If given, a range kernel is applied on-the-fly for edges preservation. + + Returns + ------- + out_sh: ndarray (X, Y, Z, n_coeffs_full) + Output SH image saved in full SH basis to preserve asymmetries. + """ + pass From 7d69d53ac13a5c9ccbfb75e313107008810f7e9b Mon Sep 17 00:00:00 2001 From: "charles.poirier" Date: Wed, 10 Nov 2021 17:37:21 -0500 Subject: [PATCH 09/20] GPU convolution using pyopencl (WIP) --- requirements.txt | 1 + scilpy/denoise/bilateral_filtering.py | 148 ++++++++++++++++++++ scilpy/denoise/opencl_utils.py | 66 +++++++++ scripts/scil_execute_bilateral_filtering.py | 20 +-- 4 files changed, 216 insertions(+), 19 deletions(-) create mode 100644 scilpy/denoise/opencl_utils.py diff --git a/requirements.txt b/requirements.txt index a494e2559..afd514457 100644 --- a/requirements.txt +++ b/requirements.txt @@ -34,3 +34,4 @@ statsmodels==0.11.* dmri-commit==1.4.* openpyxl==2.6.* cvxpy==1.1.* +pyopencl==2021.2.* diff --git a/scilpy/denoise/bilateral_filtering.py b/scilpy/denoise/bilateral_filtering.py index 02c1d2b6b..04f69911e 100644 --- a/scilpy/denoise/bilateral_filtering.py +++ b/scilpy/denoise/bilateral_filtering.py @@ -7,6 +7,115 @@ from dipy.data import get_sphere from dipy.core.sphere import Sphere from scipy.stats import multivariate_normal +from scilpy.denoise.opencl_utils import CLManager + + +def angle_aware_bilateral_filtering(in_sh, sh_order=8, + sh_basis='descoteaux07', + in_full_basis=False, + sphere_str='repulsion724', + sigma_spatial=1.0, + sigma_angular=1.0, + sigma_range=0.5): + """ + Angle-aware bilateral filtering. + """ + h_weights = _get_spatial_weights(sigma_spatial) + h_weights = h_weights.astype(np.float32) + h_half_width = len(h_weights) // 2 + + sphere = get_sphere(sphere_str) + a_weights = _get_angular_weights(h_weights.shape, sphere, sigma_angular) + + # range filtering needs to be done directly on the GPU + # because it depends on the window. + + clmanager = CLManager() + clmanager.add_program( + """ + __constant int IM_X_DIM = {0}; + __constant int IM_Y_DIM = {1}; + __constant int IM_Z_DIM = {2}; + __constant int N_COEFFS = {3}; + + __constant int H_X_DIM = {4}; + __constant int H_Y_DIM = {5}; + __constant int H_Z_DIM = {6}; + + __constant int PAD_WIDTH = {7}; + + int get_flat_index_image(const int x, const int y, + const int z, const int w, + const int padding) + {{ + return x + + y * (IM_X_DIM + 2 * padding) + + z * (IM_X_DIM + 2 * padding) + * (IM_Y_DIM + 2 * padding) + + w * (IM_X_DIM + 2 * padding) + * (IM_Y_DIM + 2 * padding) + * (IM_Z_DIM + 2 * padding); + }} + + int get_flat_index_weights(const int hx, + const int hy, + const int hz) + {{ + return hx + hy * H_X_DIM + hz * H_X_DIM * H_Y_DIM; + }} + + __kernel void correlate( + __global const float *sh_buffer, + __global const float *h_weights, + __global float *out_sh_buffer) + {{ + const int idx = get_global_id(0); + const int idy = get_global_id(1); + const int idz = get_global_id(2); + + // Example declaration for float array + float sf_coefficients[724]; + float pix_val = 0.0f; + for(int i = 0; i < N_COEFFS; ++i) + {{ + for(int hi = 0; hi < H_X_DIM; ++hi) + {{ + for(int hj = 0; hj < H_Y_DIM; ++hj) + {{ + for(int hk = 0; hk < H_Z_DIM; ++hk) + {{ + const int h_index = get_flat_index_weights(hi, + hj, + hk); + const int im_index = get_flat_index_image(idx + hi, + idy + hj, + idz + hk, + i, + PAD_WIDTH); + pix_val += h_weights[h_index] * sh_buffer[im_index]; + }} + }} + }} + }} + pix_val /= (float)N_COEFFS; + const int out_index = get_flat_index_image(idx, idy, idz, 0, 0); + out_sh_buffer[out_index] = pix_val; + }} + """.format(*in_sh.shape, *h_weights.shape, h_half_width), + 'correlate') + + v_shape = in_sh.shape + in_sh = np.pad(in_sh, ((h_half_width, h_half_width), + (h_half_width, h_half_width), + (h_half_width, h_half_width), + (0, 0))) + + clmanager.add_input_buffer(in_sh) + clmanager.add_input_buffer(h_weights) + clmanager.add_output_buffer(v_shape[:3], np.float32) + + outputs = clmanager.run(v_shape[:3]) + return outputs[0] def multivariate_bilateral_filtering(in_sh, sh_order=8, @@ -111,6 +220,45 @@ def evaluate_multivariate_gaussian(x, cov): return fx +def _get_window_directions(shape): + grid = np.indices(shape) + grid = np.moveaxis(grid, 0, -1) + grid = grid - np.asarray(shape) // 2 + return grid + + +def _get_spatial_weights(sigma_spatial): + shape = int(6 * sigma_spatial) + if shape % 2 == 0: + shape += 1 + shape = (shape, shape, shape) + + grid = _get_window_directions(shape) + + distances = np.linalg.norm(grid, axis=-1) + spatial_weights = evaluate_gaussian_dist(distances, sigma_spatial) + + # normalize filter + spatial_weights /= np.sum(spatial_weights) + return spatial_weights + + +def _get_angular_weights(shape, sphere, sigma_angular): + grid_dirs = _get_window_directions(shape).astype(np.float32) + dir_norms = np.linalg.norm(grid_dirs, axis=-1) + + # normalized grid directions + grid_dirs[dir_norms > 0] /= dir_norms[dir_norms > 0][:, None] + angles = np.arccos(np.dot(grid_dirs, sphere.vertices.T)) + angles[np.logical_not(dir_norms > 0), :] = 0.0 + + angular_weights = evaluate_gaussian_dist(angles, sigma_angular) + + # normalize filter per direction + angular_weights /= np.sum(angular_weights, axis=(0, 1, 2)) + return angular_weights + + def _get_weights_multivariate(sphere, cov): """ Get neighbors weight in respect to the direction to a voxel. diff --git a/scilpy/denoise/opencl_utils.py b/scilpy/denoise/opencl_utils.py new file mode 100644 index 000000000..05c4be2e2 --- /dev/null +++ b/scilpy/denoise/opencl_utils.py @@ -0,0 +1,66 @@ +# -*- coding: utf-8 -*- +import numpy as np +import pyopencl as cl +from pyopencl import mem_flags as mf + + +class CLManager(object): + class OutBuffer(object): + def __init__(self, buf, shape, dtype): + self.buf = buf + self.shape = shape + self.dtype = dtype + + def __init__(self): + self.context = cl.create_some_context(interactive=False) + self.queue = cl.CommandQueue(self.context) + self.input_buffers = [] + self.output_buffers = [] + + def add_input_buffer(self, arr, dtype=np.float32): + # convert to fortran ordered, float32 array + arr = np.asfortranarray(arr, dtype=dtype) + buf = cl.Buffer(self.context, + mf.READ_ONLY | mf.COPY_HOST_PTR, + hostbuf=arr) + self.input_buffers.append(buf) + + def add_output_buffer(self, shape, dtype): + buf = cl.Buffer(self.context, mf.WRITE_ONLY, + np.prod(shape) * np.dtype(dtype).itemsize) + self.output_buffers.append(self.OutBuffer(buf, shape, dtype)) + + def add_program(self, code_str, kernel_name): + program = cl.Program(self.context, code_str).build() + self.kernel = cl.Kernel(program, kernel_name) + + def run(self, global_size, local_size=None): + wait_event = self.kernel(self.queue, + global_size, + local_size, + *self.input_buffers, + *[out.buf for out in self.output_buffers]) + outputs = [] + for output in self.output_buffers: + out_arr = np.empty(output.shape, dtype=output.dtype, order='F') + cl.enqueue_copy(self.queue, out_arr, output.buf, + wait_for=[wait_event]) + outputs.append(out_arr) + return outputs + + +def angle_aware_bilateral_filtering_cl(): + """ + 1. convert arrays to float32 + 2. convert arrays to fortran ordering + 3. generate buffers + 4. kernel code + 4.1 kernel code includes range filtering + 4.2 kernel code processes all sphere directions for a voxel + 4.3 conversion to full basis is done on the GPU + 5. return output in full SH basis + + OPTIM: Do not pad, return 0 when outside + image dimensions directly in kernel. + """ + pass diff --git a/scripts/scil_execute_bilateral_filtering.py b/scripts/scil_execute_bilateral_filtering.py index c320f1e5c..247845105 100755 --- a/scripts/scil_execute_bilateral_filtering.py +++ b/scripts/scil_execute_bilateral_filtering.py @@ -2,25 +2,7 @@ # -*- coding: utf-8 -*- """ -Script to compute per-vertices hemisphere-aware (asymmetric) -filtering of spherical functions (SF) given an array of spherical harmonics -(SH) coefficients. SF are filtered using a first-neighbor Gaussian filter. -Sphere directions are also weighted by their dot product with the direction -to the center of each neighbor, clipping to 0 negative values. - -The argument `sigma` controls the standard deviation of the Gaussian. The -argument `sharpness` controls the exponent of the cosine weights. The higher it -is, the faster the weights of misaligned sphere directions decrease. A -sharpness of 0 gives the same weight to all sphere directions in an hemisphere. -Both `sharpness` and `sigma` must be positive. - -The resulting SF can be expressed using a full SH basis or a symmetric SH basis -(where the effect of the filtering is a simple denoising). - -Using default parameters, the script completes in about 15-20 minutes for a -HCP subject fiber ODF processed with tractoflow. Also note the bigger the -sphere used for SH to SF projection, the higher the RAM consumption and -compute time. +Script to compute angle-aware bilateral filtering. """ import argparse From 6c380c5f2f8dea7eb27cb056cf6a8028cd853db0 Mon Sep 17 00:00:00 2001 From: "charles.poirier" Date: Wed, 10 Nov 2021 19:43:56 -0500 Subject: [PATCH 10/20] Utilities for creating opencl kernel --- scilpy/denoise/opencl_utils.py | 58 ++++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/scilpy/denoise/opencl_utils.py b/scilpy/denoise/opencl_utils.py index 05c4be2e2..af82513aa 100644 --- a/scilpy/denoise/opencl_utils.py +++ b/scilpy/denoise/opencl_utils.py @@ -4,6 +4,24 @@ from pyopencl import mem_flags as mf +FLAT_INDEX_CL_CODE = """ +int get_flat_index(const int x, const int y, + const int z, const int w, + const int xLen, + const int yLen, + const int zLen) +{{ + return x + + y * xLen + + z * (xLen) + * (yLen) + + w * (yLen) + * (yLen) + * (zLen); +}} +""" + + class CLManager(object): class OutBuffer(object): def __init__(self, buf, shape, dtype): @@ -49,6 +67,46 @@ def run(self, global_size, local_size=None): return outputs +class CLKernel(object): + class KernelConstVar(object): + def __init__(self, ctype, value): + self.ctype = ctype + self.value = value + + def __init__(self): + self.code = "" + self.entrypoint = "" + self.constants = {} + + def add_constant(self, var_name, ctype, value): + var_name = var_name.capitalize() + if var_name in self.constants: + raise ValueError('Constant {0} already defined in kernel.' + .format(var_name)) + self.constants[var_name] = self.KernelConstVar(ctype, value) + + def set_kernel_code(self, code_str, entrypoint): + self.code = code_str + self.entrypoint = entrypoint + + def __str__(self): + code_str = "" + + # write constant values + for cname in self.constants: + const_var = self.constants[cname] + code_str += "__constant {0} {1} = {2};\n".format(const_var.ctype, + cname, + const_var.value) + + # add helper functions + code_str += FLAT_INDEX_CL_CODE + "\n" + + # write actual kernel code + code_str += self.code + return code_str + + def angle_aware_bilateral_filtering_cl(): """ 1. convert arrays to float32 From e175ce7547efb9613991e5f0d45568adf03091e6 Mon Sep 17 00:00:00 2001 From: Charles Poirier Date: Fri, 12 Nov 2021 14:59:37 -0500 Subject: [PATCH 11/20] Add optional opencl implementation for bilateral --- scilpy/denoise/angle_aware_bilateral.cl | 129 ++++++++ scilpy/denoise/bilateral_filtering.py | 336 ++++++++++---------- scilpy/denoise/opencl_utils.py | 122 +++---- scripts/scil_execute_bilateral_filtering.py | 23 +- 4 files changed, 351 insertions(+), 259 deletions(-) create mode 100644 scilpy/denoise/angle_aware_bilateral.cl diff --git a/scilpy/denoise/angle_aware_bilateral.cl b/scilpy/denoise/angle_aware_bilateral.cl new file mode 100644 index 000000000..6301ccdc6 --- /dev/null +++ b/scilpy/denoise/angle_aware_bilateral.cl @@ -0,0 +1,129 @@ +/* +OpenCL kernel code for computing angle-aware bilateral filtering. +*/ +#define IN_N_COEFFS 0 // PLACEHOLDER VALUE +#define OUT_N_COEFFS 0 // PLACEHOLDER VALUE +#define N_DIRS 0 // PLACEHOLDER VALUE +#define SIGMA_RANGE 0 // PLACEHOLDER VALUE +#define IM_X_DIM 0 // PLACEHOLDER VALUE +#define IM_Y_DIM 0 // PLACEHOLDER VALUE +#define IM_Z_DIM 0 // PLACEHOLDER VALUE +#define H_X_DIM 0 // PLACEHOLDER VALUE +#define H_Y_DIM 0 // PLACEHOLDER VALUE +#define H_Z_DIM 0 // PLACEHOLDER VALUE + +int get_flat_index(const int x, const int y, + const int z, const int w, + const int xLen, + const int yLen, + const int zLen) +{ + return x + + y * xLen + + z * xLen * yLen + + w * xLen * yLen * zLen; +} + +float sf_for_direction(const int idx, const int idy, const int idz, + const int dir_id, const float* sh_buffer, + const float* sh_to_sf_mat) +{ + float sf_coeff = 0.0f; + for(int i = 0; i < IN_N_COEFFS; ++i) + { + const int im_index = get_flat_index(idx, idy, idz, i, + IM_X_DIM + H_X_DIM - 1, + IM_Y_DIM + H_Y_DIM - 1, + IM_Z_DIM + H_Z_DIM - 1); + const float ylmu = sh_to_sf_mat[get_flat_index(i, dir_id, 0, 0, + IN_N_COEFFS, + N_DIRS, 0)]; + sf_coeff += sh_buffer[im_index] * ylmu; + } + return sf_coeff; +} + +void sf_to_sh(const float* sf_coeffs, const float* sf_to_sh_mat, + float* sh_coeffs) +{ + // although vector operations are supported + // in OpenCL, maximum n value is of 16. + for(int i = 0; i < OUT_N_COEFFS; ++i) + { + sh_coeffs[i] = 0.0f; + for(int u = 0; u < N_DIRS; ++u) + { + const float ylmu = sf_to_sh_mat[get_flat_index(u, i, 0, 0, + N_DIRS, + OUT_N_COEFFS, + 0)]; + sh_coeffs[i] += sf_coeffs[u] * ylmu; + } + } +} + +float range_distribution(const float iVal, const float jVal) +{ + const float x = fabs(iVal - jVal); + return exp(-pow(x, 2)/2.0f/pown((float)SIGMA_RANGE, 2)) + / SIGMA_RANGE / sqrt(2.0f * M_PI); +} + +__kernel void correlate(__global const float *sh_buffer, + __global const float *h_weights, + __global const float *sh_to_sf_mat, + __global const float *sf_to_sh_mat, + __global float *out_sh_buffer) +{ + const int idx = get_global_id(0); + const int idy = get_global_id(1); + const int idz = get_global_id(2); + + float sf_coeffs[N_DIRS]; + float norm_w; + for(int u = 0; u < N_DIRS; ++u) + { + sf_coeffs[u] = 0.0f; + const float sf_center = sf_for_direction(idx + H_X_DIM / 2, + idy + H_Y_DIM / 2, + idz + H_Z_DIM / 2, + u, sh_buffer, + sh_to_sf_mat); + norm_w = 0.0f; + for(int hx = 0; hx < H_X_DIM; ++hx) + { + for(int hy = 0; hy < H_Y_DIM; ++hy) + { + for(int hz = 0; hz < H_Z_DIM; ++hz) + { + const float sf_u = + sf_for_direction(idx + hx, idy + hy, + idz + hz, u, sh_buffer, + sh_to_sf_mat); + + const float range_w = + range_distribution(sf_center, sf_u); + + const float weight = + h_weights[get_flat_index(hx, hy, hz, u, + H_X_DIM, H_Y_DIM, + H_Z_DIM)] * range_w; + + sf_coeffs[u] += sf_u * weight; + norm_w += weight; + } + } + } + sf_coeffs[u] /= norm_w; + } + + float sh_coeffs[OUT_N_COEFFS]; + sf_to_sh(sf_coeffs, sf_to_sh_mat, sh_coeffs); + for(int i = 0; i < OUT_N_COEFFS; ++i) + { + const int out_index = get_flat_index(idx, idy, idz, i, + IM_X_DIM, IM_Y_DIM, + IM_Z_DIM); + out_sh_buffer[out_index] = sh_coeffs[i]; + } +} diff --git a/scilpy/denoise/bilateral_filtering.py b/scilpy/denoise/bilateral_filtering.py index 04f69911e..80f892041 100644 --- a/scilpy/denoise/bilateral_filtering.py +++ b/scilpy/denoise/bilateral_filtering.py @@ -6,8 +6,8 @@ from dipy.reconst.shm import sh_to_sf_matrix from dipy.data import get_sphere from dipy.core.sphere import Sphere -from scipy.stats import multivariate_normal -from scilpy.denoise.opencl_utils import CLManager +from scilpy.denoise.opencl_utils import (have_opencl, CLKernel, CLManager, + get_kernel_path) def angle_aware_bilateral_filtering(in_sh, sh_order=8, @@ -16,147 +16,192 @@ def angle_aware_bilateral_filtering(in_sh, sh_order=8, sphere_str='repulsion724', sigma_spatial=1.0, sigma_angular=1.0, - sigma_range=0.5): + sigma_range=0.5, + use_gpu=True, + nbr_processes=1): """ Angle-aware bilateral filtering. + + Parameters + ---------- + in_sh: ndarray (x, y, z, ncoeffs) + Input SH volume. + sh_order: int, optional + Maximum SH order of input volume. + sh_basis: str, optional + Name of SH basis used. + in_full_basis: bool, optional + True if input is expressed in full SH basis. + sphere_str: str, optional + Name of the DIPY sphere to use for sh to sf projection. + sigma_spatial: float, optional + Standard deviation for spatial filter. + sigma_angular: float, optional + Standard deviation for angular filter. + sigma_range: float, optional + Standard deviation for range filter. + use_gpu: bool, optional + True if GPU should be used. + nbr_processes: int, optional + Number of processes to use. + + Returns + ------- + out_sh: ndarray (x, y, z, ncoeffs) + Output SH coefficient array in full SH basis. + """ + if use_gpu and have_opencl: + return angle_aware_bilateral_filtering_gpu(in_sh, sh_order, + sh_basis, in_full_basis, + sphere_str, sigma_spatial, + sigma_angular, sigma_range) + elif use_gpu and not have_opencl: + raise RuntimeError("OpenCL not available.") + else: + return angle_aware_bilateral_filtering_cpu(in_sh, sh_order, + sh_basis, in_full_basis, + sphere_str, sigma_spatial, + sigma_angular, sigma_range, + nbr_processes) + + +def angle_aware_bilateral_filtering_gpu(in_sh, sh_order=8, + sh_basis='descoteaux07', + in_full_basis=False, + sphere_str='repulsion724', + sigma_spatial=1.0, + sigma_angular=1.0, + sigma_range=0.5): """ - h_weights = _get_spatial_weights(sigma_spatial) - h_weights = h_weights.astype(np.float32) - h_half_width = len(h_weights) // 2 + Angle-aware bilateral filtering using OpenCL for GPU computing. + + Parameters + ---------- + in_sh: ndarray (x, y, z, ncoeffs) + Input SH volume. + sh_order: int, optional + Maximum SH order of input volume. + sh_basis: str, optional + Name of SH basis used. + in_full_basis: bool, optional + True if input is expressed in full SH basis. + sphere_str: str, optional + Name of the DIPY sphere to use for sh to sf projection. + sigma_spatial: float, optional + Standard deviation for spatial filter. + sigma_angular: float, optional + Standard deviation for angular filter. + sigma_range: float, optional + Standard deviation for range filter. + + Returns + ------- + out_sh: ndarray (x, y, z, ncoeffs) + Output SH coefficient array in full SH basis. + """ + s_weights = _get_spatial_weights(sigma_spatial) + h_half_width = len(s_weights) // 2 sphere = get_sphere(sphere_str) - a_weights = _get_angular_weights(h_weights.shape, sphere, sigma_angular) - - # range filtering needs to be done directly on the GPU - # because it depends on the window. - - clmanager = CLManager() - clmanager.add_program( - """ - __constant int IM_X_DIM = {0}; - __constant int IM_Y_DIM = {1}; - __constant int IM_Z_DIM = {2}; - __constant int N_COEFFS = {3}; - - __constant int H_X_DIM = {4}; - __constant int H_Y_DIM = {5}; - __constant int H_Z_DIM = {6}; - - __constant int PAD_WIDTH = {7}; - - int get_flat_index_image(const int x, const int y, - const int z, const int w, - const int padding) - {{ - return x + - y * (IM_X_DIM + 2 * padding) + - z * (IM_X_DIM + 2 * padding) - * (IM_Y_DIM + 2 * padding) + - w * (IM_X_DIM + 2 * padding) - * (IM_Y_DIM + 2 * padding) - * (IM_Z_DIM + 2 * padding); - }} - - int get_flat_index_weights(const int hx, - const int hy, - const int hz) - {{ - return hx + hy * H_X_DIM + hz * H_X_DIM * H_Y_DIM; - }} - - __kernel void correlate( - __global const float *sh_buffer, - __global const float *h_weights, - __global float *out_sh_buffer) - {{ - const int idx = get_global_id(0); - const int idy = get_global_id(1); - const int idz = get_global_id(2); - - // Example declaration for float array - float sf_coefficients[724]; - float pix_val = 0.0f; - for(int i = 0; i < N_COEFFS; ++i) - {{ - for(int hi = 0; hi < H_X_DIM; ++hi) - {{ - for(int hj = 0; hj < H_Y_DIM; ++hj) - {{ - for(int hk = 0; hk < H_Z_DIM; ++hk) - {{ - const int h_index = get_flat_index_weights(hi, - hj, - hk); - const int im_index = get_flat_index_image(idx + hi, - idy + hj, - idz + hk, - i, - PAD_WIDTH); - pix_val += h_weights[h_index] * sh_buffer[im_index]; - }} - }} - }} - }} - pix_val /= (float)N_COEFFS; - const int out_index = get_flat_index_image(idx, idy, idz, 0, 0); - out_sh_buffer[out_index] = pix_val; - }} - """.format(*in_sh.shape, *h_weights.shape, h_half_width), - 'correlate') - - v_shape = in_sh.shape + a_weights = _get_angular_weights(s_weights.shape, sphere, sigma_angular) + + h_weights = s_weights[..., None] * a_weights + h_weights /= np.sum(h_weights, axis=(0, 1, 2)) + + sh_to_sf_mat = sh_to_sf_matrix(sphere, sh_order=sh_order, + basis_type=sh_basis, + full_basis=in_full_basis, + return_inv=False) + + _, sf_to_sh_mat = sh_to_sf_matrix(sphere, sh_order=sh_order, + basis_type=sh_basis, + full_basis=True, + return_inv=True) + + out_n_coeffs = sf_to_sh_mat.shape[1] + n_dirs = len(sphere.vertices) + volume_shape = in_sh.shape in_sh = np.pad(in_sh, ((h_half_width, h_half_width), (h_half_width, h_half_width), (h_half_width, h_half_width), (0, 0))) - clmanager.add_input_buffer(in_sh) - clmanager.add_input_buffer(h_weights) - clmanager.add_output_buffer(v_shape[:3], np.float32) + kernel_path = get_kernel_path('denoise', 'angle_aware_bilateral.cl') + cl_kernel = CLKernel('correlate', kernel_path) + cl_kernel.set_define('IM_X_DIM', volume_shape[0]) + cl_kernel.set_define('IM_Y_DIM', volume_shape[1]) + cl_kernel.set_define('IM_Z_DIM', volume_shape[2]) + + cl_kernel.set_define('H_X_DIM', h_weights.shape[0]) + cl_kernel.set_define('H_Y_DIM', h_weights.shape[1]) + cl_kernel.set_define('H_Z_DIM', h_weights.shape[2]) + + cl_kernel.set_define('SIGMA_RANGE', float(sigma_range)) + + cl_kernel.set_define('IN_N_COEFFS', volume_shape[-1]) + cl_kernel.set_define('OUT_N_COEFFS', out_n_coeffs) + cl_kernel.set_define('N_DIRS', n_dirs) - outputs = clmanager.run(v_shape[:3]) + cl_manager = CLManager(cl_kernel) + cl_manager.add_input_buffer(in_sh) + cl_manager.add_input_buffer(h_weights) + cl_manager.add_input_buffer(sh_to_sf_mat) + cl_manager.add_input_buffer(sf_to_sh_mat) + + cl_manager.add_output_buffer(volume_shape[:3] + (out_n_coeffs,), + np.float32) + + outputs = cl_manager.run(volume_shape[:3]) return outputs[0] -def multivariate_bilateral_filtering(in_sh, sh_order=8, - sh_basis='descoteaux07', - in_full_basis=False, - sphere_str='repulsion724', - var_cov=np.eye(2), - sigma_range=0.5, - nbr_processes=1): +def angle_aware_bilateral_filtering_cpu(in_sh, sh_order=8, + sh_basis='descoteaux07', + in_full_basis=False, + sphere_str='repulsion724', + sigma_spatial=1.0, + sigma_angular=1.0, + sigma_range=0.5, + nbr_processes=1): """ - Multivariate bilateral filtering. + Angle-aware bilateral filtering on the CPU + (optionally using multiple threads). Parameters ---------- - in_sh: ndarray (x, y, z, n_coeffs) - Input SH coefficients array + in_sh: ndarray (x, y, z, ncoeffs) + Input SH volume. sh_order: int, optional - Maximum order of the SH series. - sh_basis: {'descoteaux07', 'tournier07'}, optional - SH basis of the input signal. + Maximum SH order of input volume. + sh_basis: str, optional + Name of SH basis used. in_full_basis: bool, optional - True if the input is in full SH basis. + True if input is expressed in full SH basis. sphere_str: str, optional - Name of the sphere used to project SH coefficients to SF. - var_cov: ndarray (2, 2), optional - Variance-covariance matrix for spatio-augular distribution. + Name of the DIPY sphere to use for sh to sf projection. + sigma_spatial: float, optional + Standard deviation for spatial filter. + sigma_angular: float, optional + Standard deviation for angular filter. sigma_range: float, optional - Variance of the gaussian used for weighting intensities. + Standard deviation for range filter. nbr_processes: int, optional Number of processes to use. Returns ------- - out_sh: ndarray (x, y, z, n_coeffs) - Filtered signal as SH coefficients. + out_sh: ndarray (x, y, z, ncoeffs) + Output SH coefficient array in full SH basis. """ # Load the sphere used for projection of SH sphere = get_sphere(sphere_str) # Normalized filter for each sf direction - weights = _get_weights_multivariate(sphere, var_cov) + s_weights = _get_spatial_weights(sigma_spatial) + a_weights = _get_angular_weights(s_weights.shape, sphere, sigma_angular) + + weights = s_weights[..., None] * a_weights + weights /= np.sum(weights, axis=(0, 1, 2)) nb_sf = len(sphere.vertices) B = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, @@ -204,22 +249,12 @@ def multivariate_bilateral_filtering(in_sh, sh_order=8, return out_sh -def evaluate_gaussian_dist(x, sigma): +def evaluate_gaussian_distribution(x, sigma): assert sigma > 0.0, "Sigma must be greater than 0." cnorm = 1.0 / sigma / np.sqrt(2.0*np.pi) return cnorm * np.exp(-x**2/2/sigma**2) -def evaluate_multivariate_gaussian(x, cov): - x_shape = x.shape - x = x.reshape((-1, 2)) - flat_out = multivariate_normal.pdf(x, mean=[0, 0], cov=cov) - - fx = flat_out.reshape(x_shape[:-1]) - - return fx - - def _get_window_directions(shape): grid = np.indices(shape) grid = np.moveaxis(grid, 0, -1) @@ -236,7 +271,7 @@ def _get_spatial_weights(sigma_spatial): grid = _get_window_directions(shape) distances = np.linalg.norm(grid, axis=-1) - spatial_weights = evaluate_gaussian_dist(distances, sigma_spatial) + spatial_weights = evaluate_gaussian_distribution(distances, sigma_spatial) # normalize filter spatial_weights /= np.sum(spatial_weights) @@ -252,64 +287,13 @@ def _get_angular_weights(shape, sphere, sigma_angular): angles = np.arccos(np.dot(grid_dirs, sphere.vertices.T)) angles[np.logical_not(dir_norms > 0), :] = 0.0 - angular_weights = evaluate_gaussian_dist(angles, sigma_angular) + angular_weights = evaluate_gaussian_distribution(angles, sigma_angular) # normalize filter per direction angular_weights /= np.sum(angular_weights, axis=(0, 1, 2)) return angular_weights -def _get_weights_multivariate(sphere, cov): - """ - Get neighbors weight in respect to the direction to a voxel. - - Parameters - ---------- - sphere: Sphere - Sphere used for SF reconstruction. - dot_sharpness: float - Dot product exponent. - sigma: float - Variance of the gaussian used for weighting neighbors. - - Returns - ------- - weights: ndarray - Vertices weights with respect to voxel directions. - """ - win_size = np.ceil(6.0 * cov[0, 0] + 1.0).astype(int) - if win_size % 2 == 0: - win_size += 1 - half_size = win_size // 2 - directions = np.zeros((win_size, win_size, win_size, 3)) - for x in range(win_size): - for y in range(win_size): - for z in range(win_size): - directions[x, y, z, 0] = x - half_size - directions[x, y, z, 1] = y - half_size - directions[x, y, z, 2] = z - half_size - - non_zero_dir = np.ones(directions.shape[:-1], dtype=bool) - non_zero_dir[half_size, half_size, half_size] = False - - # normalize dir - dir_norm = np.linalg.norm(directions, axis=-1, keepdims=True) - directions[non_zero_dir] /= dir_norm[non_zero_dir] - - # angle in the range [0, pi] - angle = np.arccos(np.dot(directions, sphere.vertices.T)) - angle[half_size, half_size, half_size, :] = 0.0 - - dir_norm = np.broadcast_to(dir_norm, angle.shape) - norm_angles = np.stack([dir_norm, angle], axis=-1) - - weights = evaluate_multivariate_gaussian(norm_angles, cov) - - weights /= weights.reshape((-1, weights.shape[-1])).sum(axis=0) - - return weights - - def _process_subset_directions(args): """ Filter a subset of all sphere directions. @@ -379,7 +363,7 @@ def correlate_spatial(image_u, image_neg_u, h_filter, sigma_range): for kk in range(image_u.shape[2]): x = pad_img[ii:ii+h_w, jj:jj+h_h, kk:kk+h_d]\ - image_u[ii, jj, kk] - range_filter = evaluate_gaussian_dist(x, sigma_range) + range_filter = evaluate_gaussian_distribution(x, sigma_range) res_filter = range_filter * h_filter # Divide the filter into two filters; diff --git a/scilpy/denoise/opencl_utils.py b/scilpy/denoise/opencl_utils.py index af82513aa..7101ac8da 100644 --- a/scilpy/denoise/opencl_utils.py +++ b/scilpy/denoise/opencl_utils.py @@ -1,25 +1,11 @@ # -*- coding: utf-8 -*- import numpy as np -import pyopencl as cl -from pyopencl import mem_flags as mf - - -FLAT_INDEX_CL_CODE = """ -int get_flat_index(const int x, const int y, - const int z, const int w, - const int xLen, - const int yLen, - const int zLen) -{{ - return x + - y * xLen + - z * (xLen) - * (yLen) + - w * (yLen) - * (yLen) - * (zLen); -}} -""" +import inspect +import os +import scilpy + +from dipy.utils.optpkg import optional_package +cl, have_opencl, _ = optional_package('pyopencl') class CLManager(object): @@ -29,29 +15,28 @@ def __init__(self, buf, shape, dtype): self.shape = shape self.dtype = dtype - def __init__(self): - self.context = cl.create_some_context(interactive=False) - self.queue = cl.CommandQueue(self.context) + def __init__(self, cl_kernel): self.input_buffers = [] self.output_buffers = [] + self.context = cl.create_some_context(interactive=False) + self.queue = cl.CommandQueue(self.context) + program = cl.Program(self.context, cl_kernel.code_string).build() + self.kernel = cl.Kernel(program, cl_kernel.entry_point) + def add_input_buffer(self, arr, dtype=np.float32): # convert to fortran ordered, float32 array arr = np.asfortranarray(arr, dtype=dtype) buf = cl.Buffer(self.context, - mf.READ_ONLY | mf.COPY_HOST_PTR, + cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR, hostbuf=arr) self.input_buffers.append(buf) def add_output_buffer(self, shape, dtype): - buf = cl.Buffer(self.context, mf.WRITE_ONLY, + buf = cl.Buffer(self.context, cl.mem_flags.WRITE_ONLY, np.prod(shape) * np.dtype(dtype).itemsize) self.output_buffers.append(self.OutBuffer(buf, shape, dtype)) - def add_program(self, code_str, kernel_name): - program = cl.Program(self.context, code_str).build() - self.kernel = cl.Kernel(program, kernel_name) - def run(self, global_size, local_size=None): wait_event = self.kernel(self.queue, global_size, @@ -73,52 +58,43 @@ def __init__(self, ctype, value): self.ctype = ctype self.value = value - def __init__(self): - self.code = "" - self.entrypoint = "" - self.constants = {} - - def add_constant(self, var_name, ctype, value): - var_name = var_name.capitalize() - if var_name in self.constants: - raise ValueError('Constant {0} already defined in kernel.' - .format(var_name)) - self.constants[var_name] = self.KernelConstVar(ctype, value) - - def set_kernel_code(self, code_str, entrypoint): - self.code = code_str + def __init__(self, entrypoint, path_to_kernel): + f = open(path_to_kernel, 'r') + self.code = f.readlines() self.entrypoint = entrypoint - def __str__(self): - code_str = "" - - # write constant values - for cname in self.constants: - const_var = self.constants[cname] - code_str += "__constant {0} {1} = {2};\n".format(const_var.ctype, - cname, - const_var.value) - - # add helper functions - code_str += FLAT_INDEX_CL_CODE + "\n" - - # write actual kernel code - code_str += self.code + def set_define(self, def_name, value): + # warning! #define are not typed and therefore prone to compilation + # error. They are however faster than accessing a const variable on + # the GPU. + def_name = def_name.upper() + to_find = '#define {}'.format(def_name) + def_line = -1 + for i, line in enumerate(self.code): + if line.find(to_find) != -1: + if def_line != -1: + raise ValueError('Multiple definitions for {0}' + .format(def_name)) + def_line = i + break + if def_line == -1: + raise ValueError('Definition {0} not found in kernel code' + .format(def_name)) + + self.code[def_line] = '#define {0} {1}\n'.format(def_name, value) + + @property + def entry_point(self): + return self.entrypoint + + @property + def code_string(self): + code_str = ''.join(self.code) return code_str -def angle_aware_bilateral_filtering_cl(): - """ - 1. convert arrays to float32 - 2. convert arrays to fortran ordering - 3. generate buffers - 4. kernel code - 4.1 kernel code includes range filtering - 4.2 kernel code processes all sphere directions for a voxel - 4.3 conversion to full basis is done on the GPU - 5. return output in full SH basis - - OPTIM: Do not pad, return 0 when outside - image dimensions directly in kernel. - """ - pass +def get_kernel_path(module, kernel_name): + module_path = inspect.getfile(scilpy) + kernel_path = os.path.join(os.path.dirname(module_path), + module, kernel_name) + return kernel_path diff --git a/scripts/scil_execute_bilateral_filtering.py b/scripts/scil_execute_bilateral_filtering.py index 247845105..58f499cc7 100755 --- a/scripts/scil_execute_bilateral_filtering.py +++ b/scripts/scil_execute_bilateral_filtering.py @@ -7,6 +7,7 @@ import argparse import logging +import time from scilpy.reconst.utils import get_sh_order_and_fullness import nibabel as nib @@ -23,7 +24,7 @@ assert_outputs_exist, validate_nbr_processes) -from scilpy.denoise.bilateral_filtering import multivariate_bilateral_filtering +from scilpy.denoise.bilateral_filtering import angle_aware_bilateral_filtering def _build_arg_parser(): @@ -59,8 +60,8 @@ def _build_arg_parser(): help='Standard deviation for intensity range.' ' [%(default)s]') - p.add_argument('--covariance', default=0.0, type=float, - help='Covariance of angular and spatial value.') + p.add_argument('--use_gpu', action='store_true', + help='Use GPU for computation.') add_verbose_arg(p) add_overwrite_arg(p) @@ -82,7 +83,7 @@ def main(): assert_outputs_exist(parser, args, outputs) assert_inputs_exist(parser, args.in_sh) - validate_nbr_processes(parser, args) + nbr_processes = validate_nbr_processes(parser, args) # Prepare data sh_img = nib.load(args.in_sh) @@ -90,18 +91,20 @@ def main(): sh_order, full_basis = get_sh_order_and_fullness(data.shape[-1]) - var_cov = np.array([[args.sigma_spatial**2, args.covariance], - [args.covariance, args.sigma_angular**2]]) - + t0 = time.perf_counter() logging.info('Executing asymmetric filtering.') - asym_sh = multivariate_bilateral_filtering( + asym_sh = angle_aware_bilateral_filtering( data, sh_order=sh_order, sh_basis=args.sh_basis, in_full_basis=full_basis, sphere_str=args.sphere, - var_cov=var_cov, + sigma_spatial=args.sigma_spatial, + sigma_angular=args.sigma_angular, sigma_range=args.sigma_range, - nbr_processes=args.nbr_processes) + use_gpu=args.use_gpu, + nbr_processes=nbr_processes) + t1 = time.perf_counter() + logging.info('Elapsed time (s): {0}'.format(t1 - t0)) logging.info('Saving filtered SH to file {0}.'.format(args.out_sh)) nib.save(nib.Nifti1Image(asym_sh, sh_img.affine), args.out_sh) From 40ead134cd0c23a46bbe5441df31a0bf03f67d9f Mon Sep 17 00:00:00 2001 From: Charles Poirier Date: Thu, 18 Nov 2021 13:56:50 -0500 Subject: [PATCH 12/20] Better documentation --- requirements.txt | 1 - scilpy/denoise/angle_aware_bilateral.cl | 22 +++++---- scilpy/denoise/asym_averaging.py | 12 ++--- scilpy/denoise/bilateral_filtering.py | 52 ++++++-------------- scilpy/io/utils.py | 4 +- scripts/scil_execute_asymmetric_filtering.py | 18 ++++--- scripts/scil_execute_bilateral_filtering.py | 27 +++++++--- 7 files changed, 67 insertions(+), 69 deletions(-) diff --git a/requirements.txt b/requirements.txt index afd514457..a494e2559 100644 --- a/requirements.txt +++ b/requirements.txt @@ -34,4 +34,3 @@ statsmodels==0.11.* dmri-commit==1.4.* openpyxl==2.6.* cvxpy==1.1.* -pyopencl==2021.2.* diff --git a/scilpy/denoise/angle_aware_bilateral.cl b/scilpy/denoise/angle_aware_bilateral.cl index 6301ccdc6..dbdaa6c10 100644 --- a/scilpy/denoise/angle_aware_bilateral.cl +++ b/scilpy/denoise/angle_aware_bilateral.cl @@ -1,16 +1,18 @@ /* OpenCL kernel code for computing angle-aware bilateral filtering. */ -#define IN_N_COEFFS 0 // PLACEHOLDER VALUE -#define OUT_N_COEFFS 0 // PLACEHOLDER VALUE -#define N_DIRS 0 // PLACEHOLDER VALUE -#define SIGMA_RANGE 0 // PLACEHOLDER VALUE -#define IM_X_DIM 0 // PLACEHOLDER VALUE -#define IM_Y_DIM 0 // PLACEHOLDER VALUE -#define IM_Z_DIM 0 // PLACEHOLDER VALUE -#define H_X_DIM 0 // PLACEHOLDER VALUE -#define H_Y_DIM 0 // PLACEHOLDER VALUE -#define H_Z_DIM 0 // PLACEHOLDER VALUE + +// Compiler definitions with placeholder values +#define IN_N_COEFFS 0 +#define OUT_N_COEFFS 0 +#define N_DIRS 0 +#define SIGMA_RANGE 0 +#define IM_X_DIM 0 +#define IM_Y_DIM 0 +#define IM_Z_DIM 0 +#define H_X_DIM 0 +#define H_Y_DIM 0 +#define H_Z_DIM 0 int get_flat_index(const int x, const int y, const int z, const int w, diff --git a/scilpy/denoise/asym_averaging.py b/scilpy/denoise/asym_averaging.py index 6b061cf7b..27e313161 100644 --- a/scilpy/denoise/asym_averaging.py +++ b/scilpy/denoise/asym_averaging.py @@ -8,9 +8,8 @@ def local_asym_filtering(in_sh, sh_order=8, sh_basis='descoteaux07', - in_full_basis=False, out_full_basis=True, - dot_sharpness=1.0, sphere_str='repulsion724', - sigma=1.0): + in_full_basis=False, dot_sharpness=1.0, + sphere_str='repulsion724', sigma=1.0): """Average the SH projected on a sphere using a first-neighbor gaussian blur and a dot product weight between sphere directions and the direction to neighborhood voxels, forcing to 0 negative values and thus performing @@ -39,7 +38,7 @@ def local_asym_filtering(in_sh, sh_order=8, sh_basis='descoteaux07', Returns ------- out_sh: ndarray (x, y, z, n_coeffs) - Filtered signal as SH coefficients. + Filtered signal as SH coefficients in full SH basis. """ # Load the sphere used for projection of SH sphere = get_sphere(sphere_str) @@ -72,8 +71,9 @@ def local_asym_filtering(in_sh, sh_order=8, sh_basis='descoteaux07', mean_sf[..., sf_i] += correlate(current_sf, w_filter, mode="constant") # Convert back to SH coefficients - _, B_inv = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, - full_basis=out_full_basis) + _, B_inv = sh_to_sf_matrix(sphere, sh_order=sh_order, + basis_type=sh_basis, + full_basis=True) out_sh = np.array([np.dot(i, B_inv) for i in mean_sf], dtype=in_sh.dtype) return out_sh diff --git a/scilpy/denoise/bilateral_filtering.py b/scilpy/denoise/bilateral_filtering.py index 80f892041..dbd3c16e5 100644 --- a/scilpy/denoise/bilateral_filtering.py +++ b/scilpy/denoise/bilateral_filtering.py @@ -5,7 +5,6 @@ import itertools from dipy.reconst.shm import sh_to_sf_matrix from dipy.data import get_sphere -from dipy.core.sphere import Sphere from scilpy.denoise.opencl_utils import (have_opencl, CLKernel, CLManager, get_kernel_path) @@ -56,7 +55,8 @@ def angle_aware_bilateral_filtering(in_sh, sh_order=8, sphere_str, sigma_spatial, sigma_angular, sigma_range) elif use_gpu and not have_opencl: - raise RuntimeError("OpenCL not available.") + raise RuntimeError('Package pyopencl not found. Install pyopencl' + ' or set use_gpu to False.') else: return angle_aware_bilateral_filtering_cpu(in_sh, sh_order, sh_basis, in_full_basis, @@ -207,12 +207,6 @@ def angle_aware_bilateral_filtering_cpu(in_sh, sh_order=8, B = sh_to_sf_matrix(sphere, sh_order=sh_order, basis_type=sh_basis, return_inv=False, full_basis=in_full_basis) - # We want a B matrix to project on an inverse sphere to have the sf on - # the opposite hemisphere for a given vertice - neg_B = sh_to_sf_matrix(Sphere(xyz=-sphere.vertices), sh_order=sh_order, - basis_type=sh_basis, return_inv=False, - full_basis=in_full_basis) - if nbr_processes > 1: # Apply filter to each sphere vertice in parallel pool = multiprocessing.Pool(nbr_processes) @@ -229,7 +223,6 @@ def angle_aware_bilateral_filtering_cpu(in_sh, sh_order=8, first_ids, chunk_sizes, itertools.repeat(B), - itertools.repeat(neg_B), itertools.repeat(sigma_range))) pool.close() pool.join() @@ -238,7 +231,7 @@ def angle_aware_bilateral_filtering_cpu(in_sh, sh_order=8, mean_sf = np.concatenate(res, axis=-1) else: args = [weights, in_sh, 0, nb_sf, - B, neg_B, sigma_range] + B, sigma_range] mean_sf = _process_subset_directions(args) # Convert back to SH coefficients @@ -326,8 +319,7 @@ def _process_subset_directions(args): first_dir_id = args[2] chunk_size = args[3] B = args[4] - neg_B = args[5] - sigma_range = args[6] + sigma_range = args[5] out_sf = np.zeros(in_sh.shape[:-1] + (chunk_size,)) # Apply filter to each sphere vertice @@ -337,47 +329,35 @@ def _process_subset_directions(args): # Generate 1-channel images for directions u and -u current_sf = np.dot(in_sh, B[:, sph_id]) - opposite_sf = np.dot(in_sh, neg_B[:, sph_id]) - out_sf[..., offset_i] = correlate_spatial(current_sf, - opposite_sf, w_filter, sigma_range) return out_sf -def correlate_spatial(image_u, image_neg_u, h_filter, sigma_range): +def correlate_spatial(image_u, h_filter, sigma_range): """ Implementation of correlate function. """ h_w, h_h, h_d = h_filter.shape[:3] half_w, half_h, half_d = h_w // 2, h_h // 2, h_d // 2 - pad_img = np.zeros((image_neg_u.shape[0] + 2*half_w, - image_neg_u.shape[1] + 2*half_h, - image_neg_u.shape[2] + 2*half_d)) - pad_img[half_w:-half_w, half_h:-half_h, half_d:-half_d] = image_neg_u - out_im = np.zeros_like(image_u) - for ii in range(image_u.shape[0]): - for jj in range(image_u.shape[1]): - for kk in range(image_u.shape[2]): - x = pad_img[ii:ii+h_w, jj:jj+h_h, kk:kk+h_d]\ + image_u = np.pad(image_u, ((half_w, half_w), + (half_h, half_h), + (half_d, half_d))) + + for ii in range(out_im.shape[0]): + for jj in range(out_im.shape[1]): + for kk in range(out_im.shape[2]): + x = image_u[ii:ii+h_w, jj:jj+h_h, kk:kk+h_d]\ - image_u[ii, jj, kk] range_filter = evaluate_gaussian_distribution(x, sigma_range) - res_filter = range_filter * h_filter - # Divide the filter into two filters; - # One for the current sphere direction and - # the other for the opposite direction. - res_filter_sum = np.sum(res_filter) - center_val = res_filter[half_w, half_h, half_d] - res_filter[half_w, half_h, half_d] = 0.0 - - out_im[ii, jj, kk] = image_u[ii, jj, kk] * center_val - out_im[ii, jj, kk] += np.sum(pad_img[ii:ii+h_w, + + out_im[ii, jj, kk] += np.sum(image_u[ii:ii+h_w, jj:jj+h_h, kk:kk+h_d] * res_filter) - out_im[ii, jj, kk] /= res_filter_sum + out_im[ii, jj, kk] /= np.sum(res_filter) return out_im diff --git a/scilpy/io/utils.py b/scilpy/io/utils.py index 658cd495f..cb3e727f3 100644 --- a/scilpy/io/utils.py +++ b/scilpy/io/utils.py @@ -207,7 +207,9 @@ def add_sh_basis_args(parser, mandatory=False): def validate_nbr_processes(parser, args, default_nbr_cpu=None): - """ Check if the passed number of processes arg is valid. + """ Check if the passed number of processes arg is valid. If the number + of processes is 0, use the number of cpu. + If not valid (0 < nbr_cpu_to_use <= cpu_count), raise parser.error. Parameters diff --git a/scripts/scil_execute_asymmetric_filtering.py b/scripts/scil_execute_asymmetric_filtering.py index 1713de87c..0069e7370 100755 --- a/scripts/scil_execute_asymmetric_filtering.py +++ b/scripts/scil_execute_asymmetric_filtering.py @@ -25,19 +25,17 @@ import argparse import logging -from scilpy.reconst.utils import get_sh_order_and_fullness - import nibabel as nib import numpy as np +from dipy.reconst.shm import sph_harm_ind_list from dipy.data import SPHERE_FILES - +from scilpy.reconst.utils import get_sh_order_and_fullness from scilpy.io.utils import (add_overwrite_arg, add_verbose_arg, assert_inputs_exist, add_sh_basis_args, assert_outputs_exist) - from scilpy.denoise.asym_averaging import local_asym_filtering @@ -51,6 +49,9 @@ def _build_arg_parser(): p.add_argument('out_sh', help='File name for averaged signal.') + p.add_argument('--out_sym', default=None, + help='Name of optional symmetric output. [%(default)s]') + add_sh_basis_args(p) p.add_argument('--sphere', default='repulsion724', @@ -65,9 +66,6 @@ def _build_arg_parser(): p.add_argument('--sigma', default=1.0, type=float, help='Sigma of the gaussian to use. [%(default)s]') - p.add_argument('--out_sym', action='store_true', - help='If set, saves output in symmetric SH basis.') - add_verbose_arg(p) add_overwrite_arg(p) @@ -103,6 +101,12 @@ def main(): logging.info('Saving filtered SH to file {0}.'.format(args.out_sh)) nib.save(nib.Nifti1Image(filtered_sh, sh_img.affine), args.out_sh) + if args.out_sym: + _, orders = sph_harm_ind_list(sh_order, full_basis=True) + logging.info('Saving symmetric SH to file {0}.'.format(args.out_sym)) + nib.save(nib.Nifti1Image(filtered_sh[..., orders % 2 == 0], + sh_img.affine), args.out_sym) + if __name__ == "__main__": main() diff --git a/scripts/scil_execute_bilateral_filtering.py b/scripts/scil_execute_bilateral_filtering.py index 58f499cc7..04e067608 100755 --- a/scripts/scil_execute_bilateral_filtering.py +++ b/scripts/scil_execute_bilateral_filtering.py @@ -3,19 +3,32 @@ """ Script to compute angle-aware bilateral filtering. + +Angle-aware bilateral filtering is an extension of bilateral filtering +considering the angular distance between sphere directions for filtering +5-dimensional spatio-angular images. + +The filtering can be performed on the GPU using pyopencl by specifying +--use_gpu. Make sure you have pyopencl installed to use this option. +Otherwise, the filtering also runs entirely on the CPU, optionally using +multiple processes. + +Using default parameters, fODF filtering for a HCP subject processed with +Tractoflow takes about 12 minutes on the GPU versus 90 minutes using 16 CPU +threads. The time required scales with the sigma_spatial parameter. For +example, sigma_spatial=3.0 takes about 4.15 hours on the GPU versus 7.67 hours +on the CPU using 16 threads. """ import argparse import logging import time -from scilpy.reconst.utils import get_sh_order_and_fullness - import nibabel as nib import numpy as np from dipy.data import SPHERE_FILES from dipy.reconst.shm import sph_harm_ind_list - +from scilpy.reconst.utils import get_sh_order_and_fullness from scilpy.io.utils import (add_overwrite_arg, add_processes_arg, add_verbose_arg, @@ -23,7 +36,6 @@ add_sh_basis_args, assert_outputs_exist, validate_nbr_processes) - from scilpy.denoise.bilateral_filtering import angle_aware_bilateral_filtering @@ -40,8 +52,7 @@ def _build_arg_parser(): add_sh_basis_args(p) p.add_argument('--out_sym', default=None, - help='If set, saves additional output ' - 'in symmetric SH basis.') + help='Name of optional symmetric output. [%(default)s]') p.add_argument('--sphere', default='repulsion724', choices=sorted(SPHERE_FILES.keys()), @@ -57,7 +68,7 @@ def _build_arg_parser(): ' [%(default)s]') p.add_argument('--sigma_range', default=1.0, type=float, - help='Standard deviation for intensity range.' + help='Standard deviation for range filter.' ' [%(default)s]') p.add_argument('--use_gpu', action='store_true', @@ -92,7 +103,7 @@ def main(): sh_order, full_basis = get_sh_order_and_fullness(data.shape[-1]) t0 = time.perf_counter() - logging.info('Executing asymmetric filtering.') + logging.info('Executing angle-aware bilateral filtering.') asym_sh = angle_aware_bilateral_filtering( data, sh_order=sh_order, sh_basis=args.sh_basis, From bafaf69058122ab2ed497b7a0b832113dbaa5ed6 Mon Sep 17 00:00:00 2001 From: Charles Poirier Date: Thu, 18 Nov 2021 14:15:23 -0500 Subject: [PATCH 13/20] rename script and remove unused module --- scilpy/denoise/correlation.py | 25 ------------------- ...xecute_angle_aware_bilateral_filtering.py} | 0 2 files changed, 25 deletions(-) delete mode 100644 scilpy/denoise/correlation.py rename scripts/{scil_execute_bilateral_filtering.py => scil_execute_angle_aware_bilateral_filtering.py} (100%) diff --git a/scilpy/denoise/correlation.py b/scilpy/denoise/correlation.py deleted file mode 100644 index 7f0a0bc5a..000000000 --- a/scilpy/denoise/correlation.py +++ /dev/null @@ -1,25 +0,0 @@ -# -*- coding: utf-8 -*- - - -def correlate_sf_from_sh(in_sh, sphere, get_filter_func, sigma_range=None): - """ - Correlation operation on SF amplitudes from SH coefficients array. - - Parameters - ---------- - in_sh: ndarray (X, Y, Z, n_coeffs) - SH coefficients image to correlate. - sphere: dipy.sphere - Sphere used for SH to SF projection. - get_filter_func: Callable(direction) - Callback taking a sphere direction as input - and returning filter weights for this direction. - sigma_range: float or None, optional - If given, a range kernel is applied on-the-fly for edges preservation. - - Returns - ------- - out_sh: ndarray (X, Y, Z, n_coeffs_full) - Output SH image saved in full SH basis to preserve asymmetries. - """ - pass diff --git a/scripts/scil_execute_bilateral_filtering.py b/scripts/scil_execute_angle_aware_bilateral_filtering.py similarity index 100% rename from scripts/scil_execute_bilateral_filtering.py rename to scripts/scil_execute_angle_aware_bilateral_filtering.py From 0c5b6ccb930cceb25b1270ebc20447fcc15ab229 Mon Sep 17 00:00:00 2001 From: Charles Poirier Date: Thu, 18 Nov 2021 14:31:09 -0500 Subject: [PATCH 14/20] Fix asymmetric filtering tests --- scripts/scil_execute_asymmetric_filtering.py | 1 - scripts/tests/test_execute_asymmetric_filtering.py | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/scripts/scil_execute_asymmetric_filtering.py b/scripts/scil_execute_asymmetric_filtering.py index 0069e7370..55f14cd73 100755 --- a/scripts/scil_execute_asymmetric_filtering.py +++ b/scripts/scil_execute_asymmetric_filtering.py @@ -93,7 +93,6 @@ def main(): data, sh_order=sh_order, sh_basis=args.sh_basis, in_full_basis=full_basis, - out_full_basis=not(args.out_sym), sphere_str=args.sphere, dot_sharpness=args.sharpness, sigma=args.sigma) diff --git a/scripts/tests/test_execute_asymmetric_filtering.py b/scripts/tests/test_execute_asymmetric_filtering.py index 649b0d3a7..7aa9dc207 100644 --- a/scripts/tests/test_execute_asymmetric_filtering.py +++ b/scripts/tests/test_execute_asymmetric_filtering.py @@ -34,8 +34,8 @@ def test_sym_basis_output(script_runner): # We use a low resolution sphere to reduce execution time ret = script_runner.run('scil_execute_asymmetric_filtering.py', in_fodf, - 'out_1.nii.gz', '--out_sym', '--sphere', - 'repulsion100') + 'out_1.nii.gz', '--out_sym', 'out_sym.nii.gz', + '--sphere', 'repulsion100') assert ret.success From 981c8ccd6e78fd618cdd7bc5cc0b2cdcd9895820 Mon Sep 17 00:00:00 2001 From: Charles Poirier Date: Thu, 18 Nov 2021 17:42:09 -0500 Subject: [PATCH 15/20] Simple tests for angle-aware filtering --- ...execute_angle_aware_bilateral_filtering.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 scripts/tests/test_execute_angle_aware_bilateral_filtering.py diff --git a/scripts/tests/test_execute_angle_aware_bilateral_filtering.py b/scripts/tests/test_execute_angle_aware_bilateral_filtering.py new file mode 100644 index 000000000..e4fd32aec --- /dev/null +++ b/scripts/tests/test_execute_angle_aware_bilateral_filtering.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +import os +import tempfile + +from scilpy.io.fetcher import get_testing_files_dict, fetch_data, get_home + + +# If they already exist, this only takes 5 seconds (check md5sum) +fetch_data(get_testing_files_dict(), keys=['tracking.zip']) +tmp_dir = tempfile.TemporaryDirectory() + + +def test_help_option(script_runner): + ret = script_runner.run('scil_execute_angle_aware_bilateral_filtering.py', + '--help') + assert ret.success + + +def test_asym_basis_output(script_runner): + os.chdir(os.path.expanduser(tmp_dir.name)) + in_fodf = os.path.join(get_home(), 'tracking', 'fodf.nii.gz') + + # We use a low resolution sphere to reduce execution time + ret = script_runner.run('scil_execute_angle_aware_bilateral_filtering.py', + in_fodf, 'out_0.nii.gz', + '--sphere', 'repulsion100') + assert ret.success + + +def test_sym_basis_output(script_runner): + os.chdir(os.path.expanduser(tmp_dir.name)) + in_fodf = os.path.join(get_home(), 'tracking', 'fodf.nii.gz') + + # We use a low resolution sphere to reduce execution time + ret = script_runner.run('scil_execute_angle_aware_bilateral_filtering.py', + in_fodf, 'out_1.nii.gz', '--out_sym', + 'out_sym.nii.gz', '--sphere', 'repulsion100') + assert ret.success + + +def test_asym_input(script_runner): + os.chdir(os.path.expanduser(tmp_dir.name)) + in_fodf = os.path.join(get_home(), 'tracking', 'fodf_full.nii.gz') + + # We use a low resolution sphere to reduce execution time + ret = script_runner.run('scil_execute_angle_aware_bilateral_filtering.py', + in_fodf, 'out_2.nii.gz', + '--sphere', 'repulsion100', '-f') + assert ret.success From 841817aaaa7a07e0327b8dbb79e47b55d906b2f4 Mon Sep 17 00:00:00 2001 From: "charles.poirier" Date: Fri, 19 Nov 2021 16:42:10 -0500 Subject: [PATCH 16/20] Add comments and robustify opencl_utils --- scilpy/denoise/angle_aware_bilateral.cl | 6 +- scilpy/denoise/bilateral_filtering.py | 6 +- scilpy/denoise/opencl_utils.py | 160 +++++++++++++++++++----- 3 files changed, 137 insertions(+), 35 deletions(-) diff --git a/scilpy/denoise/angle_aware_bilateral.cl b/scilpy/denoise/angle_aware_bilateral.cl index dbdaa6c10..63f8702d4 100644 --- a/scilpy/denoise/angle_aware_bilateral.cl +++ b/scilpy/denoise/angle_aware_bilateral.cl @@ -27,8 +27,8 @@ int get_flat_index(const int x, const int y, } float sf_for_direction(const int idx, const int idy, const int idz, - const int dir_id, const float* sh_buffer, - const float* sh_to_sf_mat) + const int dir_id, global const float* sh_buffer, + global const float* sh_to_sf_mat) { float sf_coeff = 0.0f; for(int i = 0; i < IN_N_COEFFS; ++i) @@ -45,7 +45,7 @@ float sf_for_direction(const int idx, const int idy, const int idz, return sf_coeff; } -void sf_to_sh(const float* sf_coeffs, const float* sf_to_sh_mat, +void sf_to_sh(const float* sf_coeffs, global const float* sf_to_sh_mat, float* sh_coeffs) { // although vector operations are supported diff --git a/scilpy/denoise/bilateral_filtering.py b/scilpy/denoise/bilateral_filtering.py index dbd3c16e5..3a2f624cd 100644 --- a/scilpy/denoise/bilateral_filtering.py +++ b/scilpy/denoise/bilateral_filtering.py @@ -5,8 +5,7 @@ import itertools from dipy.reconst.shm import sh_to_sf_matrix from dipy.data import get_sphere -from scilpy.denoise.opencl_utils import (have_opencl, CLKernel, CLManager, - get_kernel_path) +from scilpy.denoise.opencl_utils import (have_opencl, CLKernel, CLManager) def angle_aware_bilateral_filtering(in_sh, sh_order=8, @@ -126,8 +125,7 @@ def angle_aware_bilateral_filtering_gpu(in_sh, sh_order=8, (h_half_width, h_half_width), (0, 0))) - kernel_path = get_kernel_path('denoise', 'angle_aware_bilateral.cl') - cl_kernel = CLKernel('correlate', kernel_path) + cl_kernel = CLKernel('correlate', 'denoise', 'angle_aware_bilateral.cl') cl_kernel.set_define('IM_X_DIM', volume_shape[0]) cl_kernel.set_define('IM_Y_DIM', volume_shape[1]) cl_kernel.set_define('IM_Z_DIM', volume_shape[2]) diff --git a/scilpy/denoise/opencl_utils.py b/scilpy/denoise/opencl_utils.py index 7101ac8da..f1edb1786 100644 --- a/scilpy/denoise/opencl_utils.py +++ b/scilpy/denoise/opencl_utils.py @@ -9,13 +9,22 @@ class CLManager(object): - class OutBuffer(object): - def __init__(self, buf, shape, dtype): - self.buf = buf - self.shape = shape - self.dtype = dtype + """ + Class for managing an OpenCL GPU program. + Wraps a subset of pyopencl functions to simplify its + integration with python. + + Parameters + ---------- + cl_kernel: CLKernel object + The CLKernel containing the OpenCL program to manage. + """ def __init__(self, cl_kernel): + if not have_opencl: + raise RuntimeError('pyopencl is not installed. ' + 'Cannot create CLManager instance.') + self.input_buffers = [] self.output_buffers = [] @@ -24,20 +33,91 @@ def __init__(self, cl_kernel): program = cl.Program(self.context, cl_kernel.code_string).build() self.kernel = cl.Kernel(program, cl_kernel.entry_point) + class OutBuffer(object): + """ + Structure containing output buffer information. + + Parameters + ---------- + buf: cl.Buffer + The cl.Buffer object containing the output. + shape: tuple + Shape for the output array. + dtype: dtype + Datatype for output. + """ + def __init__(self, buf, shape, dtype): + self.buf = buf + self.shape = shape + self.dtype = dtype + def add_input_buffer(self, arr, dtype=np.float32): - # convert to fortran ordered, float32 array + """ + Add an input buffer to the kernel program. Input buffers + must be added in the same order as they are declared inside + the kernel code (.cl file). + + Parameters + ---------- + arr: numpy ndarray + Data array. + dtype: dtype, optional + Optional type for array data. It is recommended to use float32 + whenever possible to avoid unexpected behaviours. + + Note + ---- + Array is reordered as fortran array and then flattened. This is + important to keep in mind when writing kernel code. + + For example, for a 3-dimensional array of shape (X, Y, Z), the flat + index for position i, j, k is idx = i + j * X + z * X * Y. + """ + # convert to fortran ordered, dtype array arr = np.asfortranarray(arr, dtype=dtype) buf = cl.Buffer(self.context, cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR, hostbuf=arr) self.input_buffers.append(buf) - def add_output_buffer(self, shape, dtype): + def add_output_buffer(self, shape, dtype=np.float32): + """ + Add an output buffer. + + Parameters + ---------- + shape: tuple + Shape of the output array. + dtype: dtype, optional + Data type for the output. It is recommended to keep + float32 to avoid unexpected behaviour. + """ buf = cl.Buffer(self.context, cl.mem_flags.WRITE_ONLY, np.prod(shape) * np.dtype(dtype).itemsize) self.output_buffers.append(self.OutBuffer(buf, shape, dtype)) def run(self, global_size, local_size=None): + """ + Execute the kernel code on the GPU. + + Parameters + ---------- + global_size: tuple + Tuple of between 1 and 3 entries representing the shape of the + grid used for GPU computing. OpenCL uses global_size to generate + a unique id for each kernel execution, which can be queried using + get_global_id(axis) with axis between 0 and 2. + local_size: tuple, optional + Dimensions of local groups. Must divide global_size exactly, + element-wise. If None, an implementation local workgroup size is + used. Memory allocated in the __local address space on the GPU is + shared between elements in a same workgroup. + + Returns + ------- + outputs: list of ndarrays + List of outputs produced by the program. + """ wait_event = self.kernel(self.queue, global_size, local_size, @@ -53,31 +133,62 @@ def run(self, global_size, local_size=None): class CLKernel(object): - class KernelConstVar(object): - def __init__(self, ctype, value): - self.ctype = ctype - self.value = value - - def __init__(self, entrypoint, path_to_kernel): + """ + Wrapper for OpenCL kernel/program code. + + Parameters + ---------- + entrypoint: string + Name of __kernel function in .cl file. + module: string + Scilpy module in which the kernel code is located. + filename: string + Name for the file containing the kernel code. + """ + def __init__(self, entrypoint, module, filename): + path_to_kernel = self._get_kernel_path(module, filename) f = open(path_to_kernel, 'r') self.code = f.readlines() self.entrypoint = entrypoint + def _get_kernel_path(self, module, filename): + """ + Get the full path for the OpenCL kernel located in scilpy + module `module` with filename `filename`. + """ + module_path = inspect.getfile(scilpy) + kernel_path = os.path.join(os.path.dirname(module_path), + module, filename) + return kernel_path + def set_define(self, def_name, value): - # warning! #define are not typed and therefore prone to compilation - # error. They are however faster than accessing a const variable on - # the GPU. + """ + Set the value for a compiler definition in the kernel code. + This method will overwrite the previous value for this definition. + + Parameters + ---------- + def_name: string + Name of definition. By convention, #define should be in upper case. + Therefore, this value will also be converted to upper case. + value: string + The value for the define. Will be replaced directly in the kernel + code. + + Note + ---- + Be careful! #define instructions are not typed and therefore prone to + compilation errors. They are however faster to access than const + variables. Moreover, they do not take additional space on the GPU. + """ def_name = def_name.upper() to_find = '#define {}'.format(def_name) - def_line = -1 + def_line = None for i, line in enumerate(self.code): if line.find(to_find) != -1: - if def_line != -1: - raise ValueError('Multiple definitions for {0}' - .format(def_name)) def_line = i break - if def_line == -1: + if def_line is None: raise ValueError('Definition {0} not found in kernel code' .format(def_name)) @@ -91,10 +202,3 @@ def entry_point(self): def code_string(self): code_str = ''.join(self.code) return code_str - - -def get_kernel_path(module, kernel_name): - module_path = inspect.getfile(scilpy) - kernel_path = os.path.join(os.path.dirname(module_path), - module, kernel_name) - return kernel_path From 800e74a65274391e087718816f7207714ae13e5b Mon Sep 17 00:00:00 2001 From: "charles.poirier" Date: Wed, 19 Jan 2022 15:11:00 -0500 Subject: [PATCH 17/20] Add missing docstrings --- scilpy/denoise/bilateral_filtering.py | 83 ++++++++++++++++++++++++--- 1 file changed, 75 insertions(+), 8 deletions(-) diff --git a/scilpy/denoise/bilateral_filtering.py b/scilpy/denoise/bilateral_filtering.py index 3a2f624cd..daa3aab7a 100644 --- a/scilpy/denoise/bilateral_filtering.py +++ b/scilpy/denoise/bilateral_filtering.py @@ -240,13 +240,44 @@ def angle_aware_bilateral_filtering_cpu(in_sh, sh_order=8, return out_sh -def evaluate_gaussian_distribution(x, sigma): +def _evaluate_gaussian_distribution(x, sigma): + """ + 1-dimensional 0-centered Gaussian distribution + with standard deviation sigma. + + Parameters + ---------- + x: ndarray or float + Points where the distribution is evaluated. + sigma: float + Standard deviation. + + Returns + ------- + out: ndarray or float + Values at x. + """ assert sigma > 0.0, "Sigma must be greater than 0." cnorm = 1.0 / sigma / np.sqrt(2.0*np.pi) return cnorm * np.exp(-x**2/2/sigma**2) def _get_window_directions(shape): + """ + Get directions from center voxel to all neighbours + for a window of given shape. + + Parameters + ---------- + shape: tuple + Dimensions of the window. + + Returns + ------- + grid: ndarray + Grid containing the direction from the center voxel to + the current position for all positions inside the window. + """ grid = np.indices(shape) grid = np.moveaxis(grid, 0, -1) grid = grid - np.asarray(shape) // 2 @@ -254,6 +285,23 @@ def _get_window_directions(shape): def _get_spatial_weights(sigma_spatial): + """ + Compute the spatial filter, which is an isotropic Gaussian filter + of standard deviation sigma_spatial forweighting by the distance + between voxel positions. The size of the filter is given by + 6 * sigma_spatial, in order to cover the range + [-3*sigma_spatial, 3*sigma_spatial]. + + Parameters + ---------- + sigma_spatial: float + Standard deviation of spatial filter. + + Returns + ------- + spatial_weights: ndarray + Spatial filter. + """ shape = int(6 * sigma_spatial) if shape % 2 == 0: shape += 1 @@ -262,7 +310,7 @@ def _get_spatial_weights(sigma_spatial): grid = _get_window_directions(shape) distances = np.linalg.norm(grid, axis=-1) - spatial_weights = evaluate_gaussian_distribution(distances, sigma_spatial) + spatial_weights = _evaluate_gaussian_distribution(distances, sigma_spatial) # normalize filter spatial_weights /= np.sum(spatial_weights) @@ -270,6 +318,25 @@ def _get_spatial_weights(sigma_spatial): def _get_angular_weights(shape, sphere, sigma_angular): + """ + Compute the angular filter, weighted by the alignment between a + sphere direction and the direction to a neighbour. The parameter + sigma_angular controls the sharpness of the kernel. + + Parameters + ---------- + shape: tuple + Shape of the angular filter. + sphere: dipy Sphere + Sphere on which the SH coefficeints are projected. + sigma_angular: float + Standard deviation of Gaussian distribution. + + Returns + ------- + angular_weights: ndarray + Angular filter for each position and for each sphere direction. + """ grid_dirs = _get_window_directions(shape).astype(np.float32) dir_norms = np.linalg.norm(grid_dirs, axis=-1) @@ -278,7 +345,7 @@ def _get_angular_weights(shape, sphere, sigma_angular): angles = np.arccos(np.dot(grid_dirs, sphere.vertices.T)) angles[np.logical_not(dir_norms > 0), :] = 0.0 - angular_weights = evaluate_gaussian_distribution(angles, sigma_angular) + angular_weights = _evaluate_gaussian_distribution(angles, sigma_angular) # normalize filter per direction angular_weights /= np.sum(angular_weights, axis=(0, 1, 2)) @@ -327,13 +394,13 @@ def _process_subset_directions(args): # Generate 1-channel images for directions u and -u current_sf = np.dot(in_sh, B[:, sph_id]) - out_sf[..., offset_i] = correlate_spatial(current_sf, - w_filter, - sigma_range) + out_sf[..., offset_i] = _correlate_spatial(current_sf, + w_filter, + sigma_range) return out_sf -def correlate_spatial(image_u, h_filter, sigma_range): +def _correlate_spatial(image_u, h_filter, sigma_range): """ Implementation of correlate function. """ @@ -349,7 +416,7 @@ def correlate_spatial(image_u, h_filter, sigma_range): for kk in range(out_im.shape[2]): x = image_u[ii:ii+h_w, jj:jj+h_h, kk:kk+h_d]\ - image_u[ii, jj, kk] - range_filter = evaluate_gaussian_distribution(x, sigma_range) + range_filter = _evaluate_gaussian_distribution(x, sigma_range) res_filter = range_filter * h_filter out_im[ii, jj, kk] += np.sum(image_u[ii:ii+h_w, From df4ce91dc3ee53e1b4e1b38cb752ce5dfe0babff Mon Sep 17 00:00:00 2001 From: "charles.poirier" Date: Wed, 19 Jan 2022 16:33:41 -0500 Subject: [PATCH 18/20] Speed up tests --- .../test_execute_angle_aware_bilateral_filtering.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/scripts/tests/test_execute_angle_aware_bilateral_filtering.py b/scripts/tests/test_execute_angle_aware_bilateral_filtering.py index e4fd32aec..ea68da1d3 100644 --- a/scripts/tests/test_execute_angle_aware_bilateral_filtering.py +++ b/scripts/tests/test_execute_angle_aware_bilateral_filtering.py @@ -8,7 +8,7 @@ # If they already exist, this only takes 5 seconds (check md5sum) -fetch_data(get_testing_files_dict(), keys=['tracking.zip']) +fetch_data(get_testing_files_dict(), keys=['processing.zip']) tmp_dir = tempfile.TemporaryDirectory() @@ -20,7 +20,8 @@ def test_help_option(script_runner): def test_asym_basis_output(script_runner): os.chdir(os.path.expanduser(tmp_dir.name)) - in_fodf = os.path.join(get_home(), 'tracking', 'fodf.nii.gz') + in_fodf = os.path.join(get_home(), 'processing', + 'fodf_descoteaux07_sub.nii.gz') # We use a low resolution sphere to reduce execution time ret = script_runner.run('scil_execute_angle_aware_bilateral_filtering.py', @@ -31,7 +32,8 @@ def test_asym_basis_output(script_runner): def test_sym_basis_output(script_runner): os.chdir(os.path.expanduser(tmp_dir.name)) - in_fodf = os.path.join(get_home(), 'tracking', 'fodf.nii.gz') + in_fodf = os.path.join(get_home(), 'processing', + 'fodf_descoteaux07_sub.nii.gz') # We use a low resolution sphere to reduce execution time ret = script_runner.run('scil_execute_angle_aware_bilateral_filtering.py', @@ -42,7 +44,8 @@ def test_sym_basis_output(script_runner): def test_asym_input(script_runner): os.chdir(os.path.expanduser(tmp_dir.name)) - in_fodf = os.path.join(get_home(), 'tracking', 'fodf_full.nii.gz') + in_fodf = os.path.join(get_home(), 'processing', + 'fodf_descoteaux07_sub_full.nii.gz') # We use a low resolution sphere to reduce execution time ret = script_runner.run('scil_execute_angle_aware_bilateral_filtering.py', From 738232696cbc4c8f218247ea3f1357d519f96d37 Mon Sep 17 00:00:00 2001 From: "charles.poirier" Date: Wed, 19 Jan 2022 16:40:26 -0500 Subject: [PATCH 19/20] Add docstring for _correlate_spatial --- scilpy/denoise/bilateral_filtering.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/scilpy/denoise/bilateral_filtering.py b/scilpy/denoise/bilateral_filtering.py index daa3aab7a..003d128b9 100644 --- a/scilpy/denoise/bilateral_filtering.py +++ b/scilpy/denoise/bilateral_filtering.py @@ -402,7 +402,22 @@ def _process_subset_directions(args): def _correlate_spatial(image_u, h_filter, sigma_range): """ - Implementation of correlate function. + Implementation of the correlation operation for anisotropic filtering. + + Parameters + ---------- + image_u: ndarray (X, Y, Z) + SF image for some sphere direction u. + h_filter: ndarray (W, H, D) + 3-dimensional filter to apply. + sigma_range: float + Standard deviation of the Gaussian distribution defining the + range kernel. + + Returns + ------- + out_im: ndarray (X, Y, Z) + Filtered SF image. """ h_w, h_h, h_d = h_filter.shape[:3] half_w, half_h, half_d = h_w // 2, h_h // 2, h_d // 2 From 1233a5d4309588e98772c19f9fe9a4c1c4e53553 Mon Sep 17 00:00:00 2001 From: Charles Poirier Date: Thu, 10 Feb 2022 09:56:18 -0500 Subject: [PATCH 20/20] Select optimal GPU based on maximum buffer size --- scilpy/denoise/opencl_utils.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/scilpy/denoise/opencl_utils.py b/scilpy/denoise/opencl_utils.py index f1edb1786..c53025469 100644 --- a/scilpy/denoise/opencl_utils.py +++ b/scilpy/denoise/opencl_utils.py @@ -28,7 +28,19 @@ def __init__(self, cl_kernel): self.input_buffers = [] self.output_buffers = [] - self.context = cl.create_some_context(interactive=False) + # Find the best device for running GPU tasks + platforms = cl.get_platforms() + best_device = None + for p in platforms: + devices = p.get_devices() + for d in devices: + if best_device is None: + best_device = d + elif (d.info.IMAGE_MAX_BUFFER_SIZE > + best_device.info.IMAGE_MAX_BUFFER_SIZE): + best_device = d + + self.context = cl.Context(devices=[best_device]) self.queue = cl.CommandQueue(self.context) program = cl.Program(self.context, cl_kernel.code_string).build() self.kernel = cl.Kernel(program, cl_kernel.entry_point)