Fixes #586

Author: MarcCote

nibabel/streamlines/tck.py

@@ -7,6 +7,7 @@
 
 import os
 import warnings
+import itertools
 
 import numpy as np
 
@@ -191,8 +192,18 @@ def save(self, fileobj):
             # Write temporary header that we will update at the end
             self._write_header(f, header)
 
+            # Make sure streamlines are in rasmm.
+            tractogram = self.tractogram.to_world(lazy=True)
+            # Assume looping over the streamlines can be done only once.
+            tractogram = iter(tractogram)
+
             try:
-                first_item = next(iter(self.tractogram))
+                # Use the first element to check
+                #  1) the tractogram is not empty;
+                #  2) quantity of information saved along each streamline.
+                first_item = next(tractogram)
+                # Put back the first element at its place.
+                tractogram = itertools.chain([first_item], tractogram)
             except StopIteration:
                 # Empty tractogram
                 header[Field.NB_STREAMLINES] = 0
@@ -216,9 +227,6 @@ def save(self, fileobj):
                        " alongside points. Dropping: {}".format(keys))
                 warnings.warn(msg, DataWarning)
 
-            # Make sure streamlines are in rasmm.
-            tractogram = self.tractogram.to_world(lazy=True)
-
             for t in tractogram:
                 data = np.r_[t.streamline, self.FIBER_DELIMITER]
                 f.write(data.astype(dtype).tostring())

nibabel/streamlines/tests/test_streamlines.py

@@ -272,3 +272,19 @@ def test_load_unknown_format(self):
 
     def test_save_unknown_format(self):
         assert_raises(ValueError, nib.streamlines.save, Tractogram(), "")
+
+    def test_save_from_generator(self):
+        tractogram = Tractogram(DATA['streamlines'],
+                                affine_to_rasmm=np.eye(4))
+
+        # Just to create a generator
+        for ext, _ in FORMATS.items():
+            filtered = (s for s in tractogram.streamlines if True)
+            lazy_tractogram = LazyTractogram(lambda: filtered,
+                                             affine_to_rasmm=np.eye(4))
+
+            with InTemporaryDirectory():
+                filename = 'streamlines' + ext
+                nib.streamlines.save(lazy_tractogram, filename)
+                tfile = nib.streamlines.load(filename, lazy_load=False)
+                assert_tractogram_equal(tfile.tractogram, tractogram)

nibabel/streamlines/trk.py

@@ -5,8 +5,9 @@
 
 import os
 import struct
-import warnings
 import string
+import warnings
+import itertools
 
 import numpy as np
 import nibabel as nib
@@ -423,8 +424,20 @@ def save(self, fileobj):
             i4_dtype = np.dtype("<i4")  # Always save in little-endian.
             f4_dtype = np.dtype("<f4")  # Always save in little-endian.
 
+            # Make sure streamlines are in rasmm then send them to voxmm.
+            tractogram = self.tractogram.to_world(lazy=True)
+            affine_to_trackvis = get_affine_rasmm_to_trackvis(header)
+            tractogram = tractogram.apply_affine(affine_to_trackvis, lazy=True)
+            # Assume looping over the streamlines can be done only once.
+            tractogram = iter(tractogram)
+
             try:
-                first_item = next(iter(self.tractogram))
+                # Use the first element to check
+                #  1) the tractogram is not empty;
+                #  2) quantity of information saved along each streamline.
+                first_item = next(tractogram)
+                # Put back the first element at its place.
+                tractogram = itertools.chain([first_item], tractogram)
             except StopIteration:
                 # Empty tractogram
                 header[Field.NB_STREAMLINES] = 0
@@ -470,11 +483,6 @@ def save(self, fileobj):
                 scalar_name[i] = encode_value_in_name(nb_values, name)
             header['scalar_name'][:] = scalar_name
 
-            # Make sure streamlines are in rasmm then send them to voxmm.
-            tractogram = self.tractogram.to_world(lazy=True)
-            affine_to_trackvis = get_affine_rasmm_to_trackvis(header)
-            tractogram = tractogram.apply_affine(affine_to_trackvis, lazy=True)
-
             for t in tractogram:
                 if any((len(d) != len(t.streamline)
                         for d in t.data_for_points.values())):