Merge pull request #121 from DiamondLightSource/hidecupy

* fully working version

conda/environment.yml

@@ -24,3 +24,4 @@ dependencies:
   - astra-toolbox::astra-toolbox
   - httomo::tomobar
   - rapidsai::cucim
+  - anaconda::scikit-image

httomolibgpu/__init__.py

@@ -1,8 +1,14 @@
-from httomolibgpu.misc.corr import *
-from httomolibgpu.misc.morph import *
-from httomolibgpu.prep.alignment import *
-from httomolibgpu.prep.normalize import *
-from httomolibgpu.prep.phase import *
-from httomolibgpu.prep.stripe import *
-from httomolibgpu.recon.algorithm import *
-from httomolibgpu.recon.rotation import *
+from httomolibgpu.misc.corr import median_filter, remove_outlier
+from httomolibgpu.misc.morph import sino_360_to_180, data_resampler
+from httomolibgpu.misc.rescale import rescale_to_int
+from httomolibgpu.prep.alignment import distortion_correction_proj_discorpy
+from httomolibgpu.prep.normalize import normalize
+from httomolibgpu.prep.phase import paganin_filter_savu, paganin_filter_tomopy
+from httomolibgpu.prep.stripe import (
+    remove_stripe_based_sorting,
+    remove_stripe_ti,
+    remove_all_stripe,
+)
+
+from httomolibgpu.recon.algorithm import FBP, SIRT, CGLS
+from httomolibgpu.recon.rotation import find_center_vo, find_center_360, find_center_pc

httomolibgpu/cuda_kernels/__init__.py

@@ -1,10 +1,13 @@
 import os
 from typing import List, Optional, Tuple
-import cupy as cp
 
+from httomolibgpu import cupywrapper
+cp = cupywrapper.cp
 
 def load_cuda_module(
-    file: str, name_expressions: Optional[List[str]] = None, options: Tuple[str, ...] = tuple()
+    file: str,
+    name_expressions: Optional[List[str]] = None,
+    options: Tuple[str, ...] = tuple(),
 ) -> cp.RawModule:
     """Load a CUDA module file, i.e. a .cu file, from the file system,
     compile it, and return is as a CuPy RawModule for further

httomolibgpu/cupywrapper.py

@@ -0,0 +1,19 @@
+cupy_run = False
+try:
+    import cupy as cp
+    import nvtx
+
+    try:
+        cp.cuda.Device(0).compute_capability
+        cupy_run = True
+    except cp.cuda.runtime.CUDARuntimeError:
+        print("CuPy library is a major dependency for HTTomolibgpu, please install")
+        import numpy as cp
+except ImportError as e:
+    print(
+        f"Failed to import module in {__file__} with error: {e}; defaulting to CPU-only mode"
+    )
+    from unittest.mock import Mock
+    import numpy as cp
+
+    nvtx = Mock()

httomolibgpu/misc/corr.py

@@ -20,24 +20,20 @@
 # ---------------------------------------------------------------------------
 """ Module for data correction """
 
-try:
-    import cupy as cp
-except ImportError:
-    print("Cupy library is a required dependency for HTTomolibgpu, please install")
-from typing import Tuple
 import numpy as np
-import nvtx
-from cupy import float32
+from httomolibgpu import cupywrapper
 
-from httomolibgpu.cuda_kernels import load_cuda_module
+cp = cupywrapper.cp
+
+nvtx = cupywrapper.nvtx
+from numpy import float32
 
 __all__ = [
     "median_filter",
     "remove_outlier",
 ]
 
 
-@nvtx.annotate()
 def median_filter(
     data: cp.ndarray,
     kernel_size: int = 3,
@@ -69,8 +65,29 @@ def median_filter(
     ValueError
         If the input array is not three dimensional.
     """
-    from cucim.skimage.filters import median
-    from cucim.skimage.morphology import disk
+    if cupywrapper.cupy_run:
+        return __median_filter(data, kernel_size, axis, dif)
+    else:
+        print("median_filter won't be executed because CuPy is not installed")
+        return data
+
+
+@nvtx.annotate()
+def __median_filter(
+    data: cp.ndarray,
+    kernel_size: int = 3,
+    axis: int = 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
 
@@ -175,4 +192,8 @@ def remove_outlier(
     if dif <= 0.0:
         raise ValueError("Threshold value (dif) must be positive and nonzero.")
 
-    return median_filter(data=data, kernel_size=kernel_size, axis=axis, dif=dif)
+    if cupywrapper.cupy_run:
+        return __median_filter(data, kernel_size, axis, dif)
+    else:
+        print("remove_outlier won't be executed because CuPy is not installed")
+        return data

httomolibgpu/misc/morph.py

@@ -20,21 +20,20 @@
 # ---------------------------------------------------------------------------
 """Module for data type morphing functions"""
 
-import cupy as cp
 import numpy as np
-import nvtx
-from typing import Literal, Tuple
+from httomolibgpu import cupywrapper
 
-from cupyx.scipy.interpolate import interpn
+cp = cupywrapper.cp
 
+nvtx = cupywrapper.nvtx
+from typing import Literal
 
 __all__ = [
     "sino_360_to_180",
     "data_resampler",
 ]
 
 
-@nvtx.annotate()
 def sino_360_to_180(
     data: cp.ndarray, overlap: int = 0, rotation: Literal["left", "right"] = "left"
 ) -> cp.ndarray:
@@ -57,6 +56,18 @@ def sino_360_to_180(
     cp.ndarray
         Output 3D data.
     """
+    if cupywrapper.cupy_run:
+        return __sino_360_to_180(data, overlap, rotation)
+    else:
+        print("sino_360_to_180 won't be executed because CuPy is not installed")
+        return data
+
+
+@nvtx.annotate()
+def __sino_360_to_180(
+    data: cp.ndarray, overlap: int = 0, rotation: Literal["left", "right"] = "left"
+) -> cp.ndarray:
+
     if data.ndim != 3:
         raise ValueError("only 3D data is supported")
 
@@ -94,7 +105,6 @@ def sino_360_to_180(
     return out
 
 
-@nvtx.annotate()
 def data_resampler(
     data: cp.ndarray, newshape: list, axis: int = 1, interpolation: str = "linear"
 ) -> cp.ndarray:
@@ -115,6 +125,19 @@ def data_resampler(
     Returns:
         cp.ndarray: Up/Down-scaled 3D cupy array
     """
+    if cupywrapper.cupy_run:
+        return __data_resampler(data, newshape, axis, interpolation)
+    else:
+        print("data_resampler won't be executed because CuPy is not installed")
+        return data
+
+
+@nvtx.annotate()
+def __data_resampler(
+    data: cp.ndarray, newshape: list, axis: int = 1, interpolation: str = "linear"
+) -> cp.ndarray:
+
+    from cupyx.scipy.interpolate import interpn
 
     if data.ndim != 3:
         raise ValueError("only 3D data is supported")