removes 2d median/dezingering using cucim

Author: dkazanc

httomolibgpu/misc/corr.py

@@ -23,13 +23,15 @@
 """
 
 import numpy as np
-from httomolibgpu import cupywrapper
 from typing import Union
 
+from httomolibgpu import cupywrapper
+
 cp = cupywrapper.cp
 nvtx = cupywrapper.nvtx
 
 from numpy import float32
+from httomolibgpu.cuda_kernels import load_cuda_module
 
 __all__ = [
     "median_filter",
@@ -40,20 +42,17 @@
 def median_filter(
     data: cp.ndarray,
     kernel_size: int = 3,
-    axis: Union[int, None] = 0,
     dif: float = 0.0,
 ) -> cp.ndarray:
     """
-    Apply 2D or 3D median filter to a 3D CuPy array. For more detailed information, see :ref:`method_median_filter`.
+    Apply 3D median filter to a 3D CuPy array. For more detailed information, see :ref:`method_median_filter`.
 
     Parameters
     ----------
     data : cp.ndarray
         Input CuPy 3D array either float32 or uint16 data type.
     kernel_size : int, optional
         The size of the filter's kernel (a diameter).
-    axis: int or None, optional:
-        Axis along which the 2D filter kernel should be applied. If set to None, then the kernel is 3D.
     dif : float, optional
         Expected difference value between outlier value and the
         median value of the array, leave equal to 0 for classical median.
@@ -69,7 +68,7 @@ def median_filter(
         If the input array is not three dimensional.
     """
     if cupywrapper.cupy_run:
-        return __median_filter(data, kernel_size, axis, dif)
+        return __median_filter(data, kernel_size, dif)
     else:
         print("median_filter won't be executed because CuPy is not installed")
         return data
@@ -79,18 +78,8 @@ def median_filter(
 def __median_filter(
     data: cp.ndarray,
     kernel_size: int = 3,
-    axis: Union[int, None] = 0,
     dif: float = 0.0,
 ) -> cp.ndarray:
-    try:
-        from cucim.skimage.filters import median
-        from cucim.skimage.morphology import disk
-    except ImportError:
-        print(
-            "Cucim library of Rapidsai is a required dependency for median_filter and remove_outlier modules, please install"
-        )
-    from httomolibgpu.cuda_kernels import load_cuda_module
-
     input_type = data.dtype
 
     if input_type not in ["float32", "uint16"]:
@@ -105,65 +94,32 @@ def __median_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")
 
-    if axis not in [0, 1, 2, None]:
-        raise ValueError("The axis should be 0,1,2 or None for full 3d processing")
-
     dz, dy, dx = data.shape
     output = cp.copy(data, order="C")
 
-    if axis == 0:
-        for j in range(dz):
-            median(data[j, :, :], footprint=disk(kernel_size // 2), out=output[j, :, :])
-    elif axis == 1:
-        for j in range(dy):
-            median(data[:, j, :], footprint=disk(kernel_size // 2), out=output[:, j, :])
-    elif axis == 2:
-        for j in range(dx):
-            median(data[:, :, j], footprint=disk(kernel_size // 2), out=output[:, :, j])
-    else:
-        # 3d median or dezinger
-        kernel_args = "median_general_kernel3d<{0}, {1}>".format(
-            "float" if input_type == "float32" else "unsigned short", kernel_size
-        )
-        block_x = 128
-        # setting grid/block parameters
-        block_dims = (block_x, 1, 1)
-        grid_x = (dx + block_x - 1) // block_x
-        grid_y = dy
-        grid_z = dz
-        grid_dims = (grid_x, grid_y, grid_z)
-        params = (data, output, cp.float32(dif), dz, dy, dx)
-
-        median_module = load_cuda_module(
-            "median_kernel", name_expressions=[kernel_args]
-        )
-        median_filt = median_module.get_function(kernel_args)
-
-        median_filt(grid_dims, block_dims, params)
-
-    if axis is not None and dif > 0:
-        # 2d dezingering enabled
-        kernel_name = "thresholding"
-        kernel = r"""
-            float dif_curr = abs(float(data) - float(output));
-            if (dif_curr > dif) {
-                output = data;
-            }
-            """
-        thresholding_kernel = cp.ElementwiseKernel(
-            "T data, raw float32 dif",
-            "T output",
-            kernel,
-            kernel_name,
-            options=("-std=c++11",),
-            no_return=True,
-        )
-        thresholding_kernel(data, float32(dif), output)
+    # 3d median or dezinger
+    kernel_args = "median_general_kernel3d<{0}, {1}>".format(
+        "float" if input_type == "float32" else "unsigned short", kernel_size
+    )
+    block_x = 128
+    # setting grid/block parameters
+    block_dims = (block_x, 1, 1)
+    grid_x = (dx + block_x - 1) // block_x
+    grid_y = dy
+    grid_z = dz
+    grid_dims = (grid_x, grid_y, grid_z)
+    params = (data, output, cp.float32(dif), dz, dy, dx)
+
+    median_module = load_cuda_module("median_kernel", name_expressions=[kernel_args])
+    median_filt = median_module.get_function(kernel_args)
+
+    median_filt(grid_dims, block_dims, params)
+
     return output
 
 
 def remove_outlier(
-    data: cp.ndarray, kernel_size: int = 3, axis: Union[int, None] = 0, dif: float = 0.1
+    data: cp.ndarray, kernel_size: int = 3, dif: float = 0.1
 ) -> cp.ndarray:
     """Selectively applies 3D median filter to a 3D CuPy array to remove outliers. Also called a dezinger.
     For more detailed information, see :ref:`method_outlier_removal`.
@@ -174,8 +130,6 @@ def remove_outlier(
         Input CuPy 3D array either float32 or uint16 data type.
     kernel_size : int, optional
         The size of the filter's kernel (a diameter).
-    axis: int or None, optional:
-        Axis along which the 2D filter kernel should be applied. If set to None, then the kernel is 3D.
     dif : float, optional
         Expected difference value between outlier value and the
         median value of the array.
@@ -195,7 +149,7 @@ def remove_outlier(
         raise ValueError("Threshold value (dif) must be positive and nonzero.")
 
     if cupywrapper.cupy_run:
-        return __median_filter(data, kernel_size, axis, dif)
+        return __median_filter(data, kernel_size, dif)
     else:
         print("remove_outlier won't be executed because CuPy is not installed")
         return data

tests/test_misc/test_corr.py

@@ -19,18 +19,14 @@ def test_median_filter3d_vs_scipy_on_arange(ensure_clean_memory):
     mat = np.arange(4 * 5 * 6).reshape(4, 5, 6)
     assert_equal(
         scipy.ndimage.median_filter(np.float32(mat), size=3),
-        median_filter(
-            cp.asarray(mat, dtype=cp.float32), kernel_size=3, axis=None
-        ).get(),
+        median_filter(cp.asarray(mat, dtype=cp.float32), kernel_size=3).get(),
     )
 
 
 def test_median_filter3d_vs_scipy(host_data, ensure_clean_memory):
     assert_equal(
         scipy.ndimage.median_filter(np.float32(host_data), size=3),
-        median_filter(
-            cp.asarray(host_data, dtype=cp.float32), kernel_size=3, axis=None
-        ).get(),
+        median_filter(cp.asarray(host_data, dtype=cp.float32), kernel_size=3).get(),
     )
 
 
@@ -43,7 +39,6 @@ def test_median_filter3d_benchmark(
         median_filter,
         cp.asarray(host_data, dtype=cp.float32),
         kernel_size=kernel_size,
-        axis=None,
     )
 
 
@@ -76,7 +71,7 @@ def test_median_filter3d_wrong_dtype(data):
 
 
 def test_median_filter3d(data):
-    filtered_data = median_filter(data, kernel_size=3, axis=None).get()
+    filtered_data = median_filter(data, kernel_size=3).get()
 
     assert filtered_data.ndim == 3
     assert_allclose(np.mean(filtered_data), 808.753494, rtol=eps)
@@ -88,49 +83,25 @@ def test_median_filter3d(data):
     assert filtered_data.flags.c_contiguous
 
     assert (
-        median_filter(data.astype(cp.float32), kernel_size=5, axis=None, dif=1.5)
-        .get()
-        .dtype
+        median_filter(data.astype(cp.float32), kernel_size=5, dif=1.5).get().dtype
         == np.float32
     )
 
 
-@pytest.mark.parametrize("axis", [0, 1, 2])
-def test_median_filter2d_axes(data, axis):
-    filtered_data = median_filter(data, kernel_size=3, axis=axis).get()
-
-    assert filtered_data.ndim == 3
-    assert filtered_data.dtype == np.uint16
-    assert filtered_data.flags.c_contiguous
-
-
-def test_median_filter2d(data):
-    filtered_data = median_filter(data, kernel_size=3, axis=0).get()
-
-    assert filtered_data.ndim == 3
-    assert_allclose(np.mean(filtered_data), 808.86170708, rtol=eps)
-    assert_allclose(np.mean(filtered_data, axis=(1, 2)).sum(), 145595.1072753)
-    assert_allclose(np.max(filtered_data), 1080)
-    assert_allclose(np.min(filtered_data), 80)
-
-    assert filtered_data.dtype == np.uint16
-    assert filtered_data.flags.c_contiguous
-
-
 @pytest.mark.perf
 def test_median_filter3d_performance(ensure_clean_memory):
     dev = cp.cuda.Device()
     data_host = np.random.random_sample(size=(450, 2160, 2560)).astype(np.float32) * 2.0
     data = cp.asarray(data_host, dtype=cp.float32)
 
     # warm up
-    median_filter(data, kernel_size=3, axis=None)
+    median_filter(data, kernel_size=3)
     dev.synchronize()
 
     start = time.perf_counter_ns()
     nvtx.RangePush("Core")
     for _ in range(10):
-        median_filter(data, kernel_size=3, axis=None)
+        median_filter(data, kernel_size=3)
     nvtx.RangePop()
     dev.synchronize()
     duration_ms = float(time.perf_counter_ns() - start) * 1e-6 / 10
@@ -139,7 +110,7 @@ def test_median_filter3d_performance(ensure_clean_memory):
 
 
 def test_remove_outlier3d(data):
-    filtered_data = remove_outlier(data, kernel_size=3, axis=None, dif=1.5).get()
+    filtered_data = remove_outlier(data, kernel_size=3, dif=1.5).get()
 
     assert filtered_data.ndim == 3
     assert_allclose(np.mean(filtered_data), 808.753494, rtol=eps)
@@ -154,16 +125,3 @@ def test_remove_outlier3d(data):
         == np.float32
     )
     assert filtered_data.flags.c_contiguous
-
-
-def test_remove_outlier2d(data):
-    filtered_data = remove_outlier(data, kernel_size=3, axis=0, dif=1.5).get()
-
-    assert filtered_data.ndim == 3
-    assert_allclose(np.mean(filtered_data), 809.049971, rtol=eps)
-    assert_allclose(np.mean(filtered_data, axis=(1, 2)).sum(), 145628.994824)
-    assert_allclose(np.max(filtered_data), 1136)
-    assert_allclose(np.min(filtered_data), 62)
-
-    assert filtered_data.dtype == np.uint16
-    assert filtered_data.flags.c_contiguous