fix using cuda

examples/cpd_affine3d_cuda.py

@@ -7,9 +7,7 @@
 else:
     cp = np
     to_cpu = lambda x: x
-import open3d as o3
 from probreg import cpd
-from probreg import callbacks
 import utils
 import time
 

examples/cpd_nonrigid3d_cuda.py

@@ -9,7 +9,6 @@
     to_cpu = lambda x: x
 import open3d as o3
 from probreg import cpd
-from probreg import callbacks
 import utils
 import time
 

examples/cpd_rigid_cuda.py

@@ -7,10 +7,8 @@
 else:
     cp = np
     to_cpu = lambda x: x
-import open3d as o3
-import transforms3d as trans
+import transforms3d as t3d
 from probreg import cpd
-from probreg import callbacks
 import utils
 import time
 

probreg/cpd.py

@@ -25,6 +25,14 @@
 """
 
 
+class DistModule:
+    def __init__(self, xp):
+        self.xp = xp
+
+    def cdist(self, x1, x2, metric):
+        return self.xp.stack([self.xp.sum(self.xp.square(x2 - ts), axis=1) for ts in x1])
+
+
 @six.add_metaclass(abc.ABCMeta)
 class CoherentPointDrift:
     """Coherent Point Drift algorithm.
@@ -50,12 +58,11 @@ def __init__(self, source: Optional[np.ndarray] = None, use_color: bool = False,
         self._use_color = use_color
         if use_cuda:
             import cupy as cp
-            from cupyx.scipy.spatial import distance as cupy_distance
 
             from . import cupy_utils
 
             self.xp = cp
-            self.distance_module = cupy_distance
+            self.distance_module = DistModule(cp)
             self.cupy_utils = cupy_utils
             self._squared_kernel_sum = cupy_utils.squared_kernel_sum
         else: