diff --git a/scilpy/denoise/angle_aware_bilateral.cl b/scilpy/denoise/angle_aware_bilateral.cl new file mode 100644 index 000000000..63f8702d4 --- /dev/null +++ b/scilpy/denoise/angle_aware_bilateral.cl @@ -0,0 +1,131 @@ +/* +OpenCL kernel code for computing angle-aware bilateral filtering. +*/ + +// 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, + 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, 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) + { + 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, global 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/asym_enhancement.py b/scilpy/denoise/asym_averaging.py similarity index 91% rename from scilpy/denoise/asym_enhancement.py rename to scilpy/denoise/asym_averaging.py index 6b061cf7b..27e313161 100644 --- a/scilpy/denoise/asym_enhancement.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 new file mode 100644 index 000000000..003d128b9 --- /dev/null +++ b/scilpy/denoise/bilateral_filtering.py @@ -0,0 +1,443 @@ +# -*- 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 scilpy.denoise.opencl_utils import (have_opencl, CLKernel, 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, + 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('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, + 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): + """ + 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(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))) + + 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]) + + 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) + + 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 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): + """ + Angle-aware bilateral filtering on the CPU + (optionally using multiple threads). + + 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. + 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. + """ + # Load the sphere used for projection of SH + sphere = get_sphere(sphere_str) + + # Normalized filter for each sf direction + 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, + 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(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, 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) + # By default, return only asymmetric SH + return out_sh + + +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 + return grid + + +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 + shape = (shape, shape, shape) + + grid = _get_window_directions(shape) + + distances = np.linalg.norm(grid, axis=-1) + spatial_weights = _evaluate_gaussian_distribution(distances, sigma_spatial) + + # normalize filter + spatial_weights /= np.sum(spatial_weights) + return spatial_weights + + +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) + + # 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_distribution(angles, sigma_angular) + + # normalize filter per direction + angular_weights /= np.sum(angular_weights, axis=(0, 1, 2)) + return angular_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] + sigma_range = args[5] + + 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]) + out_sf[..., offset_i] = _correlate_spatial(current_sf, + w_filter, + sigma_range) + return out_sf + + +def _correlate_spatial(image_u, h_filter, sigma_range): + """ + 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 + out_im = np.zeros_like(image_u) + 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 + + 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] /= np.sum(res_filter) + + return out_im diff --git a/scilpy/denoise/opencl_utils.py b/scilpy/denoise/opencl_utils.py new file mode 100644 index 000000000..c53025469 --- /dev/null +++ b/scilpy/denoise/opencl_utils.py @@ -0,0 +1,216 @@ +# -*- coding: utf-8 -*- +import numpy as np +import inspect +import os +import scilpy + +from dipy.utils.optpkg import optional_package +cl, have_opencl, _ = optional_package('pyopencl') + + +class CLManager(object): + """ + 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 = [] + + # 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) + + 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): + """ + 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=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, + *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 + + +class CLKernel(object): + """ + 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): + """ + 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 = None + for i, line in enumerate(self.code): + if line.find(to_find) != -1: + def_line = i + break + if def_line is None: + 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 diff --git a/scripts/scil_execute_angle_aware_bilateral_filtering.py b/scripts/scil_execute_angle_aware_bilateral_filtering.py new file mode 100755 index 000000000..04e067608 --- /dev/null +++ b/scripts/scil_execute_angle_aware_bilateral_filtering.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +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 +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, + assert_inputs_exist, + add_sh_basis_args, + assert_outputs_exist, + validate_nbr_processes) +from scilpy.denoise.bilateral_filtering import angle_aware_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='Name of optional symmetric output. [%(default)s]') + + 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 range filter.' + ' [%(default)s]') + + p.add_argument('--use_gpu', action='store_true', + help='Use GPU for computation.') + + 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) + + nbr_processes = 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]) + + t0 = time.perf_counter() + logging.info('Executing angle-aware 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, + sigma_spatial=args.sigma_spatial, + sigma_angular=args.sigma_angular, + sigma_range=args.sigma_range, + 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) + + 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(asym_sh[..., orders % 2 == 0], sh_img.affine), + args.out_sym) + + +if __name__ == "__main__": + main() diff --git a/scripts/scil_execute_asymmetric_filtering.py b/scripts/scil_execute_asymmetric_filtering.py index 000eec3c5..55f14cd73 100755 --- a/scripts/scil_execute_asymmetric_filtering.py +++ b/scripts/scil_execute_asymmetric_filtering.py @@ -25,20 +25,18 @@ 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_enhancement import local_asym_filtering +from scilpy.denoise.asym_averaging import local_asym_filtering def _build_arg_parser(): @@ -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) @@ -95,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) @@ -103,6 +100,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/tests/test_execute_angle_aware_bilateral_filtering.py b/scripts/tests/test_execute_angle_aware_bilateral_filtering.py new file mode 100644 index 000000000..ea68da1d3 --- /dev/null +++ b/scripts/tests/test_execute_angle_aware_bilateral_filtering.py @@ -0,0 +1,54 @@ +#!/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=['processing.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(), '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', + 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(), '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', + 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(), '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', + in_fodf, 'out_2.nii.gz', + '--sphere', 'repulsion100', '-f') + assert ret.success diff --git a/scripts/tests/test_execute_asymmetric_filtering.py b/scripts/tests/test_execute_asymmetric_filtering.py index 5f0e10df8..0f7a28f43 100644 --- a/scripts/tests/test_execute_asymmetric_filtering.py +++ b/scripts/tests/test_execute_asymmetric_filtering.py @@ -36,8 +36,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