Merge pull request #13 from matthew-brett/refactor-create-empty-header

MRG: refactoring for tractogram headers

Author: MarcCote

nibabel/streamlines/tck.py

@@ -72,9 +72,6 @@ def __init__(self, tractogram, header=None):
         This is in contrast with TRK's internal convention where it would
         have referred to a corner.
         """
-        if header is None:
-            header = self.create_empty_header()
-
         super(TckFile, self).__init__(tractogram, header)
 
     @classmethod
@@ -103,7 +100,7 @@ def is_correct_format(cls, fileobj):
 
     @classmethod
     def create_empty_header(cls):
-        """ Return an empty compliant TCK header. """
+        """ Return an empty compliant TCK header as dict """
         header = {}
 
         # Default values

nibabel/streamlines/tests/test_tractogram_file.py

@@ -1,8 +1,11 @@
-from nose.tools import assert_raises
+""" Test tractogramFile base class
+"""
 
 from ..tractogram import Tractogram
 from ..tractogram_file import TractogramFile
 
+from nose.tools import assert_raises, assert_equal
+
 
 def test_subclassing_tractogram_file():
 
@@ -37,7 +40,7 @@ def create_empty_header(cls):
 
     assert_raises(TypeError, DummyTractogramFile, Tractogram())
 
-    # Missing 'create_empty_header' method
+    # Now we have everything required.
     class DummyTractogramFile(TractogramFile):
         @classmethod
         def is_correct_format(cls, fileobj):
@@ -50,13 +53,16 @@ def load(cls, fileobj, lazy_load=True):
         def save(self, fileobj):
             pass
 
-    assert_raises(TypeError, DummyTractogramFile, Tractogram())
+    # No error
+    dtf = DummyTractogramFile(Tractogram())
+
+    # Default create_empty_header is empty dict
+    assert_equal(dtf.header, {})
 
 
 def test_tractogram_file():
     assert_raises(NotImplementedError, TractogramFile.is_correct_format, "")
     assert_raises(NotImplementedError, TractogramFile.load, "")
-    assert_raises(NotImplementedError, TractogramFile.create_empty_header)
 
     # Testing calling the 'save' method of `TractogramFile` object.
     class DummyTractogramFile(TractogramFile):
@@ -72,10 +78,6 @@ def load(cls, fileobj, lazy_load=True):
         def save(self, fileobj):
             pass
 
-        @classmethod
-        def create_empty_header(cls):
-            return None
-
     assert_raises(NotImplementedError,
                   super(DummyTractogramFile,
                         DummyTractogramFile(Tractogram)).save, "")

nibabel/streamlines/tractogram_file.py

@@ -39,7 +39,7 @@ class TractogramFile(with_metaclass(ABCMeta)):
 
     def __init__(self, tractogram, header=None):
         self._tractogram = tractogram
-        self._header = {} if header is None else header
+        self._header = self.create_empty_header() if header is None else header
 
     @property
     def tractogram(self):
@@ -77,10 +77,10 @@ def is_correct_format(cls, fileobj):
         """
         raise NotImplementedError()
 
-    @abstractclassmethod
+    @classmethod
     def create_empty_header(cls):
         """ Returns an empty header for this streamlines file format. """
-        raise NotImplementedError()
+        return {}
 
     @abstractclassmethod
     def load(cls, fileobj, lazy_load=True):

nibabel/streamlines/trk.py

@@ -72,8 +72,8 @@ def get_affine_trackvis_to_rasmm(header):
 
     Parameters
     ----------
-    header : dict
-        Dict containing trackvis header.
+    header : dict or ndarray
+        Dict or numpy structured array containing trackvis header.
 
     Returns
     -------
@@ -101,9 +101,12 @@ def get_affine_trackvis_to_rasmm(header):
     # If the voxel order implied by the affine does not match the voxel
     # order in the TRK header, change the orientation.
     # voxel (header) -> voxel (affine)
-    header_ornt = asstr(header[Field.VOXEL_ORDER])
+    vox_order = header[Field.VOXEL_ORDER]
+    # Input header can be dict or structured array
+    if hasattr(vox_order, 'item'):  # structured array
+        vox_order = header[Field.VOXEL_ORDER].item()
     affine_ornt = "".join(aff2axcodes(header[Field.VOXEL_TO_RASMM]))
-    header_ornt = axcodes2ornt(header_ornt)
+    header_ornt = axcodes2ornt(vox_order.decode('latin1'))
     affine_ornt = axcodes2ornt(affine_ornt)
     ornt = nib.orientations.ornt_transform(header_ornt, affine_ornt)
     M = nib.orientations.inv_ornt_aff(ornt, header[Field.DIMENSIONS])
@@ -235,10 +238,6 @@ def __init__(self, tractogram, header=None):
         and *mm* space where coordinate (0,0,0) refers to the center
         of the voxel.
         """
-        if header is None:
-            header_rec = self.create_empty_header()
-            header = dict(zip(header_rec.dtype.names, header_rec[0]))
-
         super(TrkFile, self).__init__(tractogram, header)
 
     @classmethod
@@ -266,20 +265,28 @@ def is_correct_format(cls, fileobj):
             return magic_number == cls.MAGIC_NUMBER
 
     @classmethod
-    def create_empty_header(cls):
-        """ Return an empty compliant TRK header. """
-        header = np.zeros(1, dtype=header_2_dtype)
+    def _default_structarr(cls):
+        """ Return an empty compliant TRK header as numpy structured array
+        """
+        st_arr = np.zeros((), dtype=header_2_dtype)
 
         # Default values
-        header[Field.MAGIC_NUMBER] = cls.MAGIC_NUMBER
-        header[Field.VOXEL_SIZES] = np.array((1, 1, 1), dtype="f4")
-        header[Field.DIMENSIONS] = np.array((1, 1, 1), dtype="h")
-        header[Field.VOXEL_TO_RASMM] = np.eye(4, dtype="f4")
-        header[Field.VOXEL_ORDER] = b"RAS"
-        header['version'] = 2
-        header['hdr_size'] = cls.HEADER_SIZE
+        st_arr[Field.MAGIC_NUMBER] = cls.MAGIC_NUMBER
+        st_arr[Field.VOXEL_SIZES] = np.array((1, 1, 1), dtype="f4")
+        st_arr[Field.DIMENSIONS] = np.array((1, 1, 1), dtype="h")
+        st_arr[Field.VOXEL_TO_RASMM] = np.eye(4, dtype="f4")
+        st_arr[Field.VOXEL_ORDER] = b"RAS"
+        st_arr['version'] = 2
+        st_arr['hdr_size'] = cls.HEADER_SIZE
 
-        return header
+        return st_arr
+
+    @classmethod
+    def create_empty_header(cls):
+        """ Return an empty compliant TRK header as dict
+        """
+        st_arr = cls._default_structarr()
+        return dict(zip(st_arr.dtype.names, st_arr.tolist()))
 
     @classmethod
     def load(cls, fileobj, lazy_load=False):
@@ -388,7 +395,7 @@ def save(self, fileobj):
             of the TRK header data).
         """
         # Enforce little-endian byte order for header
-        header = self.create_empty_header().newbyteorder('<')
+        header = self._default_structarr().newbyteorder('<')
 
         # Override hdr's fields by those contained in `header`.
         for k, v in self.header.items():
@@ -406,7 +413,6 @@ def save(self, fileobj):
         nb_scalars = 0
         nb_properties = 0
 
-        header = header[0]
         with Opener(fileobj, mode="wb") as f:
             # Keep track of the beginning of the header.
             beginning = f.tell()