Fixed import of OrderedDict and added support to apply transformation on streamlines.

Author: MarcCote

nibabel/streamlines/base_format.py

@@ -1,3 +1,4 @@
+import numpy as np
 from warnings import warn
 
 from nibabel.externals.six.moves import zip_longest
@@ -111,14 +112,56 @@ def __getitem__(self, idx):
     def __len__(self):
         return len(self.points)
 
-    def to_world_space(self, as_generator=False):
-        affine = self.header.voxel_to_world
-        new_points = (apply_affine(affine, pts) for pts in self.points)
+    def copy(self):
+        """ Returns a copy of this `Streamlines` object. """
+        streamlines = Streamlines(self.points, self.scalars, self.properties)
+        streamlines._header = self.header.copy()
+        return streamlines
 
-        if not as_generator:
-            return list(new_points)
+    def transform(self, affine, lazy=False):
+        """ Applies an affine transformation on the points of each streamline.
 
-        return new_points
+        Parameters
+        ----------
+        affine : 2D array (4,4)
+            Transformation that will be applied on each streamline.
+        lazy : bool (optional)
+            If true output will be a generator of arrays instead of a list.
+
+        Returns
+        -------
+        streamlines
+            If `lazy` is true, a `LazyStreamlines` object is returned,
+            otherwise a `Streamlines` object is returned. In both case,
+            streamlines are in a space defined by `affine`.
+        """
+        points = lambda: (apply_affine(affine, pts) for pts in self.points)
+        if not lazy:
+            points = list(points())
+
+        streamlines = self.copy()
+        streamlines.points = points
+        streamlines.header.to_world_space = np.dot(streamlines.header.to_world_space,
+                                                   np.linalg.inv(affine))
+
+        return streamlines
+
+    def to_world_space(self, lazy=False):
+        """ Sends the streamlines back into world space.
+
+        Parameters
+        ----------
+        lazy : bool (optional)
+            If true output will be a generator of arrays instead of a list.
+
+        Returns
+        -------
+        streamlines
+            If `lazy` is true, a `LazyStreamlines` object is returned,
+            otherwise a `Streamlines` object is returned. In both case,
+            streamlines are in world space.
+        """
+        return self.transform(self.header.to_world_space, lazy)
 
 
 class LazyStreamlines(Streamlines):
@@ -255,8 +298,36 @@ def __len__(self):
 
         return self.header.nb_streamlines
 
+    def copy(self):
+        """ Returns a copy of this `LazyStreamlines` object. """
+        streamlines = LazyStreamlines(self._points, self._scalars, self._properties)
+        streamlines._header = self.header.copy()
+        return streamlines
+
+    def transform(self, affine):
+        """ Applies an affine transformation on the points of each streamline.
+
+        Parameters
+        ----------
+        affine : 2D array (4,4)
+            Transformation that will be applied on each streamline.
+
+        Returns
+        -------
+        streamlines : `LazyStreamlines` object
+            Streamlines living in a space defined by `affine`.
+        """
+        return super(LazyStreamlines, self).transform(affine, lazy=True)
+
     def to_world_space(self):
-        return super(LazyStreamlines, self).to_world_space(as_generator=True)
+        """ Sends the streamlines back into world space.
+
+        Returns
+        -------
+        streamlines : `LazyStreamlines` object
+            Streamlines living in world space.
+        """
+        return super(LazyStreamlines, self).to_world_space(lazy=True)
 
 
 class StreamlinesFile:

nibabel/streamlines/header.py

@@ -1,6 +1,7 @@
+import copy
 import numpy as np
 from nibabel.orientations import aff2axcodes
-from collections import OrderedDict
+from nibabel.externals import OrderedDict
 
 
 class Field:
@@ -18,7 +19,7 @@ class Field:
     DIMENSIONS = "dimensions"
     MAGIC_NUMBER = "magic_number"
     ORIGIN = "origin"
-    VOXEL_TO_WORLD = "voxel_to_world"
+    to_world_space = "to_world_space"
     VOXEL_ORDER = "voxel_order"
     ENDIAN = "endian"
 
@@ -28,33 +29,33 @@ def __init__(self):
         self._nb_streamlines = None
         self._nb_scalars_per_point = None
         self._nb_properties_per_streamline = None
-        self._voxel_to_world = np.eye(4)
+        self._to_world_space = np.eye(4)
         self.extra = OrderedDict()
 
     @property
-    def voxel_to_world(self):
-        return self._voxel_to_world
+    def to_world_space(self):
+        return self._to_world_space
 
-    @voxel_to_world.setter
-    def voxel_to_world(self, value):
-        self._voxel_to_world = np.array(value, dtype=np.float32)
+    @to_world_space.setter
+    def to_world_space(self, value):
+        self._to_world_space = np.array(value, dtype=np.float32)
 
     @property
     def voxel_sizes(self):
-        """ Get voxel sizes from voxel_to_world. """
-        return np.sqrt(np.sum(self.voxel_to_world[:3, :3]**2, axis=0))
+        """ Get voxel sizes from to_world_space. """
+        return np.sqrt(np.sum(self.to_world_space[:3, :3]**2, axis=0))
 
     @voxel_sizes.setter
     def voxel_sizes(self, value):
         scaling = np.r_[np.array(value), [1]]
         old_scaling = np.r_[np.array(self.voxel_sizes), [1]]
         # Remove old scaling and apply new one
-        self.voxel_to_world = np.dot(np.diag(scaling/old_scaling), self.voxel_to_world)
+        self.to_world_space = np.dot(np.diag(scaling/old_scaling), self.to_world_space)
 
     @property
     def voxel_order(self):
-        """ Get voxel order from voxel_to_world. """
-        return "".join(aff2axcodes(self.voxel_to_world))
+        """ Get voxel order from to_world_space. """
+        return "".join(aff2axcodes(self.to_world_space))
 
     @property
     def nb_streamlines(self):
@@ -88,8 +89,17 @@ def extra(self):
     def extra(self, value):
         self._extra = OrderedDict(value)
 
+    def copy(self):
+        header = StreamlinesHeader()
+        header._nb_streamlines = self.nb_streamlines
+        header.nb_scalars_per_point = self.nb_scalars_per_point
+        header.nb_properties_per_streamline = self.nb_properties_per_streamline
+        header.to_world_space = self.to_world_space.copy()
+        header.extra = copy.deepcopy(self.extra)
+        return header
+
     def __eq__(self, other):
-        return (np.allclose(self.voxel_to_world, other.voxel_to_world) and
+        return (np.allclose(self.to_world_space, other.to_world_space) and
                 self.nb_streamlines == other.nb_streamlines and
                 self.nb_scalars_per_point == other.nb_scalars_per_point and
                 self.nb_properties_per_streamline == other.nb_properties_per_streamline and
@@ -100,7 +110,7 @@ def __repr__(self):
         txt += "nb_streamlines: " + repr(self.nb_streamlines) + '\n'
         txt += "nb_scalars_per_point: " + repr(self.nb_scalars_per_point) + '\n'
         txt += "nb_properties_per_streamline: " + repr(self.nb_properties_per_streamline) + '\n'
-        txt += "voxel_to_world: " + repr(self.voxel_to_world) + '\n'
+        txt += "to_world_space: " + repr(self.to_world_space) + '\n'
         txt += "voxel_sizes: " + repr(self.voxel_sizes) + '\n'
 
         txt += "Extra fields: {\n"

nibabel/streamlines/tests/test_base_format.py

@@ -117,8 +117,8 @@ def test_to_world_space(self):
         # World space is (RAS+) with voxel size of 2x3x4mm.
         streamlines.header.voxel_sizes = (2, 3, 4)
 
-        new_points = streamlines.to_world_space()
-        for new_pts, pts in zip(new_points, self.points):
+        new_streamlines = streamlines.to_world_space()
+        for new_pts, pts in zip(new_streamlines.points, self.points):
             for dim, size in enumerate(streamlines.header.voxel_sizes):
                 assert_array_almost_equal(new_pts[:, dim], size*pts[:, dim])
 
@@ -129,7 +129,7 @@ def test_header(self):
         assert_equal(streamlines.header.nb_scalars_per_point, 0)
         assert_equal(streamlines.header.nb_properties_per_streamline, 0)
         assert_array_equal(streamlines.header.voxel_sizes, (1, 1, 1))
-        assert_array_equal(streamlines.header.voxel_to_world, np.eye(4))
+        assert_array_equal(streamlines.header.to_world_space, np.eye(4))
         assert_equal(streamlines.header.extra, {})
 
         streamlines = Streamlines(self.points, self.colors, self.mean_curvature_torsion)
@@ -138,15 +138,15 @@ def test_header(self):
         assert_equal(streamlines.header.nb_scalars_per_point, self.colors[0].shape[1])
         assert_equal(streamlines.header.nb_properties_per_streamline, self.mean_curvature_torsion[0].shape[0])
 
-        # Modifying voxel_sizes should be reflected in voxel_to_world
+        # Modifying voxel_sizes should be reflected in to_world_space
         streamlines.header.voxel_sizes = (2, 3, 4)
         assert_array_equal(streamlines.header.voxel_sizes, (2, 3, 4))
-        assert_array_equal(np.diag(streamlines.header.voxel_to_world), (2, 3, 4, 1))
+        assert_array_equal(np.diag(streamlines.header.to_world_space), (2, 3, 4, 1))
 
-        # Modifying scaling of voxel_to_world should be reflected in voxel_sizes
-        streamlines.header.voxel_to_world = np.diag([4, 3, 2, 1])
+        # Modifying scaling of to_world_space should be reflected in voxel_sizes
+        streamlines.header.to_world_space = np.diag([4, 3, 2, 1])
         assert_array_equal(streamlines.header.voxel_sizes, (4, 3, 2))
-        assert_array_equal(streamlines.header.voxel_to_world, np.diag([4, 3, 2, 1]))
+        assert_array_equal(streamlines.header.to_world_space, np.diag([4, 3, 2, 1]))
 
         # Test that we can run __repr__ without error.
         repr(streamlines.header)
@@ -313,7 +313,7 @@ def test_lazy_streamlines_header(self):
         assert_equal(streamlines.header.nb_scalars_per_point, 0)
         assert_equal(streamlines.header.nb_properties_per_streamline, 0)
         assert_array_equal(streamlines.header.voxel_sizes, (1, 1, 1))
-        assert_array_equal(streamlines.header.voxel_to_world, np.eye(4))
+        assert_array_equal(streamlines.header.to_world_space, np.eye(4))
         assert_equal(streamlines.header.extra, {})
 
         points = lambda: (x for x in self.points)

nibabel/streamlines/tests/test_header.py

@@ -12,7 +12,7 @@ def test_streamlines_header():
     assert_true(header.nb_scalars_per_point is None)
     assert_true(header.nb_properties_per_streamline is None)
     assert_array_equal(header.voxel_sizes, (1, 1, 1))
-    assert_array_equal(header.voxel_to_world, np.eye(4))
+    assert_array_equal(header.to_world_space, np.eye(4))
     assert_equal(header.extra, {})
 
     # Modify simple attributes
@@ -23,15 +23,15 @@ def test_streamlines_header():
     assert_equal(header.nb_scalars_per_point, 2)
     assert_equal(header.nb_properties_per_streamline, 3)
 
-    # Modifying voxel_sizes should be reflected in voxel_to_world
+    # Modifying voxel_sizes should be reflected in to_world_space
     header.voxel_sizes = (2, 3, 4)
     assert_array_equal(header.voxel_sizes, (2, 3, 4))
-    assert_array_equal(np.diag(header.voxel_to_world), (2, 3, 4, 1))
+    assert_array_equal(np.diag(header.to_world_space), (2, 3, 4, 1))
 
-    # Modifying scaling of voxel_to_world should be reflected in voxel_sizes
-    header.voxel_to_world = np.diag([4, 3, 2, 1])
+    # Modifying scaling of to_world_space should be reflected in voxel_sizes
+    header.to_world_space = np.diag([4, 3, 2, 1])
     assert_array_equal(header.voxel_sizes, (4, 3, 2))
-    assert_array_equal(header.voxel_to_world, np.diag([4, 3, 2, 1]))
+    assert_array_equal(header.to_world_space, np.diag([4, 3, 2, 1]))
 
     # Test that we can run __repr__ without error.
     repr(header)

nibabel/streamlines/tests/test_streamlines.py

@@ -140,28 +140,28 @@ def setUp(self):
         self.nb_streamlines = len(self.points)
         self.nb_scalars_per_point = self.colors[0].shape[1]
         self.nb_properties_per_streamline = len(self.mean_curvature_torsion[0])
-        self.voxel_to_world = np.eye(4)
+        self.to_world_space = np.eye(4)
 
     def test_load_empty_file(self):
         for empty_filename in self.empty_filenames:
             streamlines = nib.streamlines.load(empty_filename,
-                                               ref=self.voxel_to_world,
+                                               ref=self.to_world_space,
                                                lazy_load=False)
             assert_true(type(streamlines), Streamlines)
             check_streamlines(streamlines, 0, [], [], [])
 
     def test_load_simple_file(self):
         for simple_filename in self.simple_filenames:
             streamlines = nib.streamlines.load(simple_filename,
-                                               ref=self.voxel_to_world,
+                                               ref=self.to_world_space,
                                                lazy_load=False)
             assert_true(type(streamlines), Streamlines)
             check_streamlines(streamlines, self.nb_streamlines,
                               self.points, [], [])
 
             # Test lazy_load
             streamlines = nib.streamlines.load(simple_filename,
-                                               ref=self.voxel_to_world,
+                                               ref=self.to_world_space,
                                                lazy_load=True)
             assert_true(type(streamlines), LazyStreamlines)
             check_streamlines(streamlines, self.nb_streamlines,
@@ -180,15 +180,15 @@ def test_load_complex_file(self):
                 properties = self.mean_curvature_torsion
 
             streamlines = nib.streamlines.load(complex_filename,
-                                               ref=self.voxel_to_world,
+                                               ref=self.to_world_space,
                                                lazy_load=False)
             assert_true(type(streamlines), Streamlines)
             check_streamlines(streamlines, self.nb_streamlines,
                               self.points, scalars, properties)
 
             # Test lazy_load
             streamlines = nib.streamlines.load(complex_filename,
-                                               ref=self.voxel_to_world,
+                                               ref=self.to_world_space,
                                                lazy_load=True)
             assert_true(type(streamlines), LazyStreamlines)
             check_streamlines(streamlines, self.nb_streamlines,
@@ -199,7 +199,7 @@ def test_save_simple_file(self):
         for ext in nib.streamlines.FORMATS.keys():
             with tempfile.NamedTemporaryFile(mode="w+b", suffix=ext) as f:
                 nib.streamlines.save(streamlines, f.name)
-                loaded_streamlines = nib.streamlines.load(f, ref=self.voxel_to_world, lazy_load=False)
+                loaded_streamlines = nib.streamlines.load(f, ref=self.to_world_space, lazy_load=False)
                 check_streamlines(loaded_streamlines, self.nb_streamlines,
                                   self.points, [], [])
 
@@ -224,6 +224,6 @@ def test_save_complex_file(self):
                 if cls.can_save_properties():
                     properties = self.mean_curvature_torsion
 
-                loaded_streamlines = nib.streamlines.load(f, ref=self.voxel_to_world, lazy_load=False)
+                loaded_streamlines = nib.streamlines.load(f, ref=self.to_world_space, lazy_load=False)
                 check_streamlines(loaded_streamlines, self.nb_streamlines,
                                   self.points, scalars, properties)