Update raven filter

Author: neon60

httomolibgpu/cuda_kernels/raven_filter.cu

@@ -1,70 +1,13 @@
-#define IDX2R(i,j,N) (((i)*(N))+(j))
-
-template <typename Type>
-__global__ void fftshift_2D(Type *out, int N1, int N2)
-{
-    int i = threadIdx.y + blockDim.y * blockIdx.y;
-    int j = threadIdx.x + blockDim.x * blockIdx.x;
-
-    if (i < N1 && j < N2) 
-    {
-        double a = 1-2*((i+j)&1);
-        data[IDX2R(i,j,N2)].x *= a;
-        data[IDX2R(i,j,N2)].y *= a;
-    }
-}
 
 template <typename Type, int diameter>
-__global__ void raven__filter_kernel3d(const Type *in, Type *out, float dif,
-                                      int Z, int M, int N) {
-  constexpr int radius = diameter / 2;
-  constexpr int d3 = diameter * diameter * diameter;
-  constexpr int midpoint = d3 / 2;
+__global__ void raven__filter_kernel3d(Type *in, int Z, int M, int N) {
 
-  Type ValVec[d3];
   const long i = blockDim.x * blockIdx.x + threadIdx.x;
   const long j = blockDim.y * blockIdx.y + threadIdx.y;
   const long k = blockDim.z * blockIdx.z + threadIdx.z;
 
   if (i >= N || j >= M || k >= Z)
     return;
 
-  long long index = static_cast<long long>(i) + N * static_cast<long long>(j) + N * M * static_cast<long long>(k);
-
-  int counter = 0;
-  for (int i_m = -radius; i_m <= radius; i_m++) {
-    long long i1 = i + i_m;   // using long long to avoid integer overflows
-    if ((i1 < 0) || (i1 >= N))
-      i1 = i;
-    for (int j_m = -radius; j_m <= radius; j_m++) {
-      long long j1 = j + j_m;
-      if ((j1 < 0) || (j1 >= M))
-        j1 = j;
-      for (int k_m = -radius; k_m <= radius; k_m++) {
-        long long k1 = k + k_m;
-        if ((k1 < 0) || (k1 >= Z))
-          k1 = k;
-        ValVec[counter] = in[i1 + N * j1 + N * M * k1];
-        counter++;
-      }
-    }
-  }  
-
-  /* do bubble sort here */
-  for (int x = 0; x < d3 - 1; x++) {
-    for (int y = 0; y < d3 - x - 1; y++) {
-      if (ValVec[y] > ValVec[y + 1]) {
-        Type temp = ValVec[y];
-        ValVec[y] = ValVec[y + 1];
-        ValVec[y + 1] = temp;
-      }
-    }
-  }
-
-  if (dif > 0.0f) {
-    /* perform dezingering */
-    out[index] =
-        fabsf(in[index] - ValVec[midpoint]) >= dif ? ValVec[midpoint] : in[index];
-  }
-  else out[index] = ValVec[midpoint]; /* median filtering */
+  
 }

httomolibgpu/misc/raven_filter.py

@@ -35,19 +35,26 @@
 
 if cupy_run:
     from httomolibgpu.cuda_kernels import load_cuda_module
+    from cupyx.scipy.fft import fft2, ifft2, fftshift, ifftshift
 else:
     load_cuda_module = Mock()
+    fft2 = Mock()
+    ifft2 = Mock()
+    fftshift = Mock()
+    ifftshift = Mock()
 
 
 __all__ = [
     "raven_filter",
 ]
 
 
-def raven_filter(
+def raven_filter_savu(
     data: cp.ndarray,
     kernel_size: int = 3,
-    dif: float = 0.0,
+    pad_y: int = 100,
+    pad_x: int = 100,
+    pad_method: str = "edge",
 ) -> cp.ndarray:
     """
     Applies raven filter to a 3D CuPy array. For more detailed information, see :ref:`method_raven_filter`.
@@ -58,14 +65,20 @@ def raven_filter(
         Input CuPy 3D array either float32 or uint16 data type.
     kernel_size : int, optional
         The size of the filter's kernel (a diameter).
-    dif : float, optional
-        Expected difference value between outlier value and the
-        median value of the array, leave equal to 0 for classical median.
+
+    pad_y : int, optional
+        Pad the top and bottom of projections.
+
+    pad_x : int, optional
+        Pad the left and right of projections.
+
+    pad_method : str, optional
+        Numpy pad method to use.
 
     Returns
     -------
     ndarray
-        Median filtered 3D CuPy array either float32 or uint16 data type.
+        Raven filtered 3D CuPy array either float32 or uint16 data type.
 
     Raises
     ------
@@ -86,11 +99,25 @@ def raven_filter(
     if kernel_size not in [3, 5, 7, 9, 11, 13]:
         raise ValueError("Please select a correct kernel size: 3, 5, 7, 9, 11, 13")
 
-    dz, dy, dx = data.shape
-    output = cp.copy(data, order="C")
 
-    # 3d median or dezinger
-    kernel_args = "median_general_kernel3d<{0}, {1}>".format(
+    dz_orig, dy_orig, dx_orig = data.shape
+
+    padded_data, pad_tup = cp.pad(data, fftpad, "edge")
+    dz, dy, dx = padded_data.shape
+
+    # 3D FFT of data
+    padded_data = cp.pad(data, ((0, 0), (pad_y, pad_y), (pad_x, pad_x)), mode=pad_method)
+    fft_data = fft2(padded_data, axes=(-2, -1), overwrite_x=True)
+    fft_data_shifted = fftshift(fft_data)
+
+    # Setup various values for the filter
+    _, height, width = data.shape
+
+    height1 = height + 2 * pad_y
+    width1 = width + 2 * pad_x
+
+    # raven
+    kernel_args = "raven_general_kernel3d<{0}, {1}>".format(
         "float" if input_type == "float32" else "unsigned short", kernel_size
     )
     block_x = 128
@@ -100,11 +127,15 @@ def raven_filter(
     grid_y = dy
     grid_z = dz
     grid_dims = (grid_x, grid_y, grid_z)
-    params = (data, output, cp.float32(dif), dz, dy, dx)
+    params = (fft_data_shifted, dz, dy, dx)
+
+    raven_module = load_cuda_module("raven_kernel", name_expressions=[kernel_args])
+    raven_filt = raven_module.get_function(kernel_args)
+
+    raven_filt(grid_dims, block_dims, params)
 
-    median_module = load_cuda_module("raven_filter", name_expressions=[kernel_args])
-    median_filt = median_module.get_function(kernel_args)
+    fft_data = fftshift(fft_data_shifted)
 
-    median_filt(grid_dims, block_dims, params)
+    data = ifft2(fft_data, axes=(-2, -1), overwrite_x=True, norm="forward")
 
-    return output
+    return data