Affine align with GPUtools

jni · Sep 29, 2019 · 0c6fb69 · 0c6fb69
1 parent 1341a70
commit 0c6fb69
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 9 deletions.
diff --git a/doc/examples/transform/plot_register_affine.py b/doc/examples/transform/plot_register_affine.py
@@ -18,6 +18,7 @@
         2017. isbn: 1491922877, 9781491922873.
 
 """
+import time
 
 import numpy as np
 from scipy import ndimage as ndi
@@ -41,7 +42,7 @@
                              [s, c, 50],
                              [0, 0,  1]])
 
-image = astronaut()[..., 1]  # Just green channel
+image = astronaut()[..., 1].copy()  # Just green channel
 target = ndi.affine_transform(image, matrix_transform)
 
 ###############################################################################
@@ -57,8 +58,16 @@
 
 level_alignments = []
 
-register_matrix = register_affine(image, target,
-                                  level_callback=level_alignments.append)
+t0 = time.time()
+register_matrix = register_affine(image, target, use_gputools=True)
+t1 = time.time()
+_ = register_affine(image, target, use_gputools=False,
+                    level_callback=level_alignments.append)
+t2 = time.time()
+
+print('gpu time: ', t1 - t0)
+print('cpu time: ', t2 - t1)
+
 
 ###############################################################################
 # Once we have the matrix, it's easy to transform the target image to match

diff --git a/skimage/transform/registration.py b/skimage/transform/registration.py
@@ -1,4 +1,6 @@
 import numpy as np
+from gputools.transforms import affine
+from gputools import OCLArray, OCLImage
 from scipy import ndimage as ndi
 from scipy.optimize import minimize
 
@@ -81,7 +83,9 @@ def cost_nmi(image0, image1, *, bins=100):
 
 def register_affine(reference, target, *, cost=cost_nmi, minimum_size=8,
                     multichannel=False, inverse=True,
-                    level_callback=None, method='Powell', **kwargs):
+                    level_callback=None, method='Powell',
+                    use_gputools=False,
+                    **kwargs):
     """Find a transformation matrix to register a target image to a reference.
 
     Parameters
@@ -155,23 +159,47 @@ def register_affine(reference, target, *, cost=cost_nmi, minimum_size=8,
     parameter_vector = _matrix_to_parameter_vector(np.identity(ndim + 1))
 
     for ref, tgt in image_pairs:
+        if use_gputools:
+            if ndim == 2:
+                ref = ref[np.newaxis, ...]
+                tgt = tgt[np.newaxis, ...]
+            tgt_gpu = OCLImage.from_array(tgt.astype(np.float32))
+            if not multichannel:
+                result_gpu = OCLArray.empty(tgt.shape, np.float32)
+            else:
+                result_gpu = OCLArray.empty(tgt.shape[:-1], np.float32)
         def _cost(param):
             transformation = _parameter_vector_to_matrix(param, ndim)
+            if use_gputools and ndim == 2:
+                transformation2 = np.identity((ndim + 2))
+                transformation2[1:, 1:] = transformation[:, :]
             if not multichannel:
-                transformed = ndi.affine_transform(tgt, transformation,
-                                                   order=1)
+                if use_gputools:
+                    transformed = affine(tgt_gpu, transformation2,
+                                         interpolation='linear',
+                                         res_g=result_gpu).get()
+                else:
+                    transformed = ndi.affine_transform(tgt, transformation,
+                                                       order=1)
             else:
                 transformed = np.zeros_like(tgt)
                 for ch in range(tgt.shape[-1]):
-                    ndi.affine_transform(tgt[..., ch], transformation,
-                                         order=1, output=transformed[..., ch])
+                    if use_gputools:
+                        transformed[..., ch] = affine(tgt[..., ch],
+                                                      transformation,
+                                                      interpolation='linear',
+                                                      res_g=result_gpu).get()
+                    else:
+                        ndi.affine_transform(tgt[..., ch], transformation,
+                                             order=1,
+                                             output=transformed[..., ch])
             return cost(ref, transformed)
 
         result = minimize(_cost, x0=parameter_vector, method=method, **kwargs)
         parameter_vector = result.x
         if level_callback is not None:
             level_callback(
-                (tgt,
+                (np.squeeze(tgt),
                  _parameter_vector_to_matrix(parameter_vector, ndim),
                  result.fun)
             )