Merge pull request #719 from h2see/windows-cuda-path-bugfix

Ensure the windows CUDA_PATH is included in DLL search paths.

Author: AnderBiguri

Python/tigre/__init__.py

@@ -6,6 +6,20 @@
 # https://github.com/CERN/TIGRE/issues/349
 import os
 
+_CUDA_PATH_WIN_DIR_HANDLE = None
+def _add_cuda_path_on_windows():
+    global _CUDA_PATH_WIN_DIR_HANDLE
+    import sys
+    if sys.platform == "win32" and hasattr(os, "add_dll_directory"):
+        _cuda_path = os.environ.get("CUDA_PATH")
+        if _cuda_path:
+            _bin = os.path.join(_cuda_path, "bin")
+            if os.path.isdir(_bin):
+                # a reference to the handle returned by os.add_dll_directory is necessary, 
+                # otherwise it will be garbage collected and the directory will be removed 
+                # from the search path again.
+                _CUDA_PATH_WIN_DIR_HANDLE = os.add_dll_directory(_bin)
+
 # if hasattr(os, "add_dll_directory"):
 #     # Add all the DLL directories manually
 #     # see:
@@ -21,7 +35,18 @@
 from .utilities.geometry import geometry
 from .utilities.geometry_default import ConeGeometryDefault as geometry_default
 from .utilities.geometry_default import FanGeometryDefault as fan_geometry_default
-from .utilities.Ax import Ax
+
+# always import Ax and _Ax_ext before all other ctypes extensions (checks for DLL import errors).
+from .utilities.Ax import Ax, _Ax_ext, _ensure_Ax_ext_import
+# _Ax_ext will be None if the import failed, 
+# in that case we try to add the CUDA path on windows and import again, 
+# this will raise the original import error with the full message if it still fails.
+if _Ax_ext is None:
+    from .utilities.Ax import _try_import_Ax_ext
+    _add_cuda_path_on_windows()
+    _try_import_Ax_ext()
+_ensure_Ax_ext_import() # check if the import was successful, if not this will raise the original import error.
+
 from .utilities.Atb import Atb
 from .utilities.visualization.plotproj import plotproj, plotProj, plotSinogram
 from .utilities.visualization.plotimg import plotimg, plotImg

Python/tigre/utilities/Ax.py

@@ -1,12 +1,33 @@
 import copy
 
 import numpy as np
-from _Ax import _Ax_ext
 
-from .gpu import GpuIds
+_Ax_ext = None
+GpuIds = None
+def _try_import_Ax_ext():
+    global _Ax_ext, GpuIds
+    if _Ax_ext is None:
+        try:
+            from _Ax import _Ax_ext as imported
+            _Ax_ext = imported
+        except ImportError:
+            _Ax_ext = None
+    if _Ax_ext is not None and GpuIds is None:
+        # GpuIds depends on ctypes as well, so we must import it 
+        # after sucessful completion of the _Ax_ext import.
+        from .gpu import GpuIds as imported_ids
+        GpuIds = imported_ids
+
+def _ensure_Ax_ext_import():
+    global _Ax_ext
+    if _Ax_ext is None:
+        from _Ax import _Ax_ext
+
+_try_import_Ax_ext()
 
 
 def Ax(img, geo, angles, projection_type="Siddon", **kwargs):
+    _ensure_Ax_ext_import() # check if the import was successful, if not this will raise the original import error with the full message.
     if img.dtype != np.float32:
         raise TypeError("Input data should be float32, not " + str(img.dtype))
     if not np.isreal(img).all():