RF: Accept generator input, test with generator and with affine.

dipy · Dec 30, 2015 · a956f71 · a956f71
1 parent deb04ac
commit a956f71
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 6 deletions.
diff --git a/dipy/tracking/tests/test_utils.py b/dipy/tracking/tests/test_utils.py
@@ -596,8 +596,30 @@ def test_vals_from_img():
                              [3.9, 0, 0],
                              [4.1, 0, 0]])]
 
+    vv = vals_from_img(img, streamlines)
+    npt.assert_almost_equal(vv, [[100., 150., 200., 250.],
+                                 [200., 310., 390., 410.]])
+
     vv = vals_from_img(img, np.array(streamlines))
 
-    # Interpolation is to floating point precision:
+    npt.assert_almost_equal(vv, [[100., 150., 200., 250.],
+                                 [200., 310., 390., 410.]])
+
+    affine = np.eye(4)
+    affine[:, 3] = [-100, 10, 1, 1]
+    x_streamlines = move_streamlines(streamlines, affine)
+
+    vv = vals_from_img(img, x_streamlines, affine=affine)
+    npt.assert_almost_equal(vv, [[100., 150., 200., 250.],
+                                 [200., 310., 390., 410.]])
+
+    # The generator has already been consumed so needs to be regenerated:
+    x_streamlines = list(move_streamlines(streamlines, affine))
+    vv = vals_from_img(img, x_streamlines, affine=affine)
+    npt.assert_almost_equal(vv, [[100., 150., 200., 250.],
+                                 [200., 310., 390., 410.]])
+
+    vv = vals_from_img(img, np.array(x_streamlines), affine=affine)
+
     npt.assert_almost_equal(vv, [[100., 150., 200., 250.],
                                  [200., 310., 390., 410.]])
diff --git a/dipy/tracking/utils.py b/dipy/tracking/utils.py
@@ -50,6 +50,7 @@
 
 from functools import wraps
 from warnings import warn
+import types
 
 from nibabel.affines import apply_affine
 from scipy.spatial.distance import cdist
@@ -965,13 +966,15 @@ def vals_from_img(img, streamlines, affine=None):
 
     """
     data = img.get_data().astype(np.float)
-    if isinstance(streamlines, list):
-        if affine is not None:
-            streamlines = dtu.move_streamlines(streamlines,
-                                               np.linalg.inv(affine))
+    if affine is not None:
+        streamlines = move_streamlines(streamlines,
+                                       np.linalg.inv(affine))
+
+    if (isinstance(streamlines, list) or
+            isinstance(streamlines, types.GeneratorType)):
         vals = []
         for sl in streamlines:
-            vals.append(vfu.interpolate_scalar_3d(data, sl))
+            vals.append(vfu.interpolate_scalar_3d(data, sl)[0])
 
     elif isinstance(streamlines, np.ndarray):
         sl_shape = streamlines.shape