Fixed a bug that made StaticVolume crash on cpu; updated test_devices to validate differences between devices

Author: the-lay

tests/test_devices.py

@@ -2,9 +2,12 @@
 import numpy as np
 import cupy as cp
 import matplotlib.pyplot as plt
+plt.ioff()
 
 size = (50, 50, 50)
 data = np.random.random(size).astype(np.float32)
+rotation = np.random.uniform(0, 180, 3)
+scale = np.random.uniform(0.7, 1.3, 3)
 
 rows = [
     [
@@ -24,33 +27,52 @@
     [
         {
             'interpolation': 'linear',
-            'device': 'gpu:0'
+            'device': 'gpu'
         },
         {
             'interpolation': 'bspline',
-            'device': 'gpu:0'
+            'device': 'gpu'
         },
         {
             'interpolation': 'filt_bspline',
-            'device': 'gpu:0'
+            'device': 'gpu'
         },
     ]
 ]
 
+### testing transforms methods
 fig, ax = plt.subplots(len(rows), len(rows[0]), sharex=True, sharey=True)
 for i, r in enumerate(rows):
     for j, case in enumerate(r):
         print(f'Test case: {case["interpolation"]} / {case["device"]}')
-        tf = vt.transform(data, rotation=(0, 30, 0), interpolation=case['interpolation'],
+        tf = vt.transform(data, rotation=rotation, scale=scale, interpolation=case['interpolation'],
                           profile=True, device=case['device'])
 
         ax[i][j].set_title(f'{case["interpolation"]} / {case["device"]}')
-
-        if isinstance(tf, cp.ndarray):
-            ax[i][j].imshow(tf[size[0] // 2].get())
-        else:
-            ax[i][j].imshow(tf[size[0] // 2])
+        ax[i][j].imshow(tf[size[0] // 2])
 
 plt.show()
 
 
+### testing static volume methods
+print('\n\n\n')
+st_volumes = [
+    vt.StaticVolume(data, interpolation='linear', device='cpu'),
+    vt.StaticVolume(data, interpolation='bspline', device='cpu'),
+    vt.StaticVolume(data, interpolation='filt_bspline', device='cpu'),
+    vt.StaticVolume(data, interpolation='linear', device='gpu'),
+    vt.StaticVolume(data, interpolation='bspline', device='gpu'),
+    vt.StaticVolume(data, interpolation='filt_bspline', device='gpu')
+]
+
+fig, ax = plt.subplots(2, 3, sharex=True, sharey=True)
+
+for n, v in enumerate(st_volumes):
+    print(f'Test case: {v.interpolation} / {v.device}')
+    tf = v.transform(scale=scale, rotation=rotation, profile=True)
+    i, j = int(n / 3), n % 3
+    ax[i][j].set_title(f'{v.interpolation} / {v.device}')
+    ax[i][j].imshow(tf[size[0] // 2])
+
+plt.show()
+

voltools/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.3.2'
+__version__ = '0.3.3'
 
 from .transforms import AVAILABLE_INTERPOLATIONS, AVAILABLE_DEVICES, scale, shear, rotate, translate, transform, affine
 from .volume import StaticVolume

voltools/volume.py

@@ -52,6 +52,7 @@ def __init__(self, data: np.ndarray, interpolation: str = 'linear', device: str
             del data
 
         elif device == 'cpu':
+            self.shape = data.shape
             self.data = data
 
     def affine(self, transform_m: np.ndarray, profile: bool = False, output: cp.ndarray = None) -> Union[np.ndarray, None]: