Merge pull request #512 from MarcCote/bf_support_trk_v1

MRG: streamlines support reading TRK (version 1)

Add compatibility with trk file version 1; assume identity affine for v1 files.

Author: matthew-brett

nibabel/streamlines/tests/test_trk.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import copy
 import unittest
 import numpy as np
@@ -99,68 +100,94 @@ def test_load_complex_file(self):
             trk = TrkFile.load(DATA['complex_trk_fname'], lazy_load=lazy_load)
             assert_tractogram_equal(trk.tractogram, DATA['complex_tractogram'])
 
+    def trk_with_bytes(self, trk_key='simple_trk_fname', endian='<'):
+        """ Return example trk file bytes and struct view onto bytes """
+        with open(DATA[trk_key], 'rb') as fobj:
+            trk_bytes = fobj.read()
+        dt = trk_module.header_2_dtype.newbyteorder(endian)
+        trk_struct = np.ndarray((1,), dt, buffer=trk_bytes)
+        trk_struct.flags.writeable = True
+        return trk_struct, trk_bytes
+
     def test_load_file_with_wrong_information(self):
         trk_file = open(DATA['simple_trk_fname'], 'rb').read()
 
         # Simulate a TRK file where `count` was not provided.
-        count = np.array(0, dtype="int32").tostring()
-        new_trk_file = trk_file[:1000-12] + count + trk_file[1000-8:]
-        trk = TrkFile.load(BytesIO(new_trk_file), lazy_load=False)
+        trk_struct, trk_bytes = self.trk_with_bytes()
+        trk_struct[Field.NB_STREAMLINES] = 0
+        trk = TrkFile.load(BytesIO(trk_bytes), lazy_load=False)
         assert_tractogram_equal(trk.tractogram, DATA['simple_tractogram'])
 
         # Simulate a TRK where `vox_to_ras` is not recorded (i.e. all zeros).
-        vox_to_ras = np.zeros((4, 4), dtype=np.float32).tostring()
-        new_trk_file = trk_file[:440] + vox_to_ras + trk_file[440+64:]
+        trk_struct, trk_bytes = self.trk_with_bytes()
+        trk_struct[Field.VOXEL_TO_RASMM] = np.zeros((4, 4))
         with clear_and_catch_warnings(record=True, modules=[trk_module]) as w:
-            trk = TrkFile.load(BytesIO(new_trk_file))
+            trk = TrkFile.load(BytesIO(trk_bytes))
             assert_equal(len(w), 1)
             assert_true(issubclass(w[0].category, HeaderWarning))
             assert_true("identity" in str(w[0].message))
             assert_array_equal(trk.affine, np.eye(4))
 
         # Simulate a TRK where `vox_to_ras` is invalid.
-        vox_to_ras = np.zeros((4, 4), dtype=np.float32)
-        vox_to_ras[3, 3] = 1
-        vox_to_ras = vox_to_ras.tostring()
-        new_trk_file = trk_file[:440] + vox_to_ras + trk_file[440+64:]
+        trk_struct, trk_bytes = self.trk_with_bytes()
+        trk_struct[Field.VOXEL_TO_RASMM] = np.diag([0, 0, 0, 1])
         with clear_and_catch_warnings(record=True, modules=[trk_module]) as w:
-            assert_raises(HeaderError, TrkFile.load, BytesIO(new_trk_file))
+            assert_raises(HeaderError, TrkFile.load, BytesIO(trk_bytes))
 
         # Simulate a TRK file where `voxel_order` was not provided.
-        voxel_order = np.zeros(1, dtype="|S3").tostring()
-        new_trk_file = trk_file[:948] + voxel_order + trk_file[948+3:]
+        trk_struct, trk_bytes = self.trk_with_bytes()
+        trk_struct[Field.VOXEL_ORDER] = b''
         with clear_and_catch_warnings(record=True, modules=[trk_module]) as w:
-            TrkFile.load(BytesIO(new_trk_file))
+            TrkFile.load(BytesIO(trk_bytes))
             assert_equal(len(w), 1)
             assert_true(issubclass(w[0].category, HeaderWarning))
             assert_true("LPS" in str(w[0].message))
 
         # Simulate a TRK file with an unsupported version.
-        version = np.int32(123).tostring()
-        new_trk_file = trk_file[:992] + version + trk_file[992+4:]
-        assert_raises(HeaderError, TrkFile.load, BytesIO(new_trk_file))
+        trk_struct, trk_bytes = self.trk_with_bytes()
+        trk_struct['version'] = 123
+        assert_raises(HeaderError, TrkFile.load, BytesIO(trk_bytes))
 
         # Simulate a TRK file with a wrong hdr_size.
-        hdr_size = np.int32(1234).tostring()
-        new_trk_file = trk_file[:996] + hdr_size + trk_file[996+4:]
-        assert_raises(HeaderError, TrkFile.load, BytesIO(new_trk_file))
+        trk_struct, trk_bytes = self.trk_with_bytes()
+        trk_struct['hdr_size'] = 1234
+        assert_raises(HeaderError, TrkFile.load, BytesIO(trk_bytes))
 
         # Simulate a TRK file with a wrong scalar_name.
-        trk_file = open(DATA['complex_trk_fname'], 'rb').read()
-        noise = np.int32(42).tostring()
-        new_trk_file = trk_file[:47] + noise + trk_file[47+4:]
-        assert_raises(HeaderError, TrkFile.load, BytesIO(new_trk_file))
+        trk_struct, trk_bytes = self.trk_with_bytes('complex_trk_fname')
+        trk_struct['scalar_name'][0, 0] = b'colors\x003\x004'
+        assert_raises(HeaderError, TrkFile.load, BytesIO(trk_bytes))
 
         # Simulate a TRK file with a wrong property_name.
-        noise = np.int32(42).tostring()
-        new_trk_file = trk_file[:254] + noise + trk_file[254+4:]
-        assert_raises(HeaderError, TrkFile.load, BytesIO(new_trk_file))
+        trk_struct, trk_bytes = self.trk_with_bytes('complex_trk_fname')
+        trk_struct['property_name'][0, 0] = b'colors\x003\x004'
+        assert_raises(HeaderError, TrkFile.load, BytesIO(trk_bytes))
+
+    def test_load_trk_version_1(self):
+        # Simulate and test a TRK (version 1).
+        # First check that setting the RAS affine works in version 2.
+        trk_struct, trk_bytes = self.trk_with_bytes()
+        trk_struct[Field.VOXEL_TO_RASMM] = np.diag([2, 3, 4, 1])
+        trk = TrkFile.load(BytesIO(trk_bytes))
+        assert_array_equal(trk.affine, np.diag([2, 3, 4, 1]))
+        # Next check that affine assumed identity if version 1.
+        trk_struct['version'] = 1
+        with clear_and_catch_warnings(record=True, modules=[trk_module]) as w:
+            trk = TrkFile.load(BytesIO(trk_bytes))
+            assert_equal(len(w), 1)
+            assert_true(issubclass(w[0].category, HeaderWarning))
+            assert_true("identity" in str(w[0].message))
+            assert_array_equal(trk.affine, np.eye(4))
+            assert_array_equal(trk.header['version'], 1)
 
     def test_load_complex_file_in_big_endian(self):
-        trk_file = open(DATA['complex_trk_big_endian_fname'], 'rb').read()
+        trk_struct, trk_bytes = self.trk_with_bytes(
+            'complex_trk_big_endian_fname', endian='>')
         # We use hdr_size as an indicator of little vs big endian.
-        hdr_size_big_endian = np.array(1000, dtype=">i4").tostring()
-        assert_equal(trk_file[996:996+4], hdr_size_big_endian)
+        good_orders = '>' if sys.byteorder == 'little' else '>='
+        hdr_size = trk_struct['hdr_size']
+        assert_true(hdr_size.dtype.byteorder in good_orders)
+        assert_equal(hdr_size, 1000)
 
         for lazy_load in [False, True]:
             trk = TrkFile.load(DATA['complex_trk_big_endian_fname'],

nibabel/streamlines/trk.py

@@ -26,36 +26,11 @@
 MAX_NB_NAMED_SCALARS_PER_POINT = 10
 MAX_NB_NAMED_PROPERTIES_PER_STREAMLINE = 10
 
-# See http://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html
-header_1_dtd = [(Field.MAGIC_NUMBER, 'S6'),
-                (Field.DIMENSIONS, 'h', 3),
-                (Field.VOXEL_SIZES, 'f4', 3),
-                (Field.ORIGIN, 'f4', 3),
-                (Field.NB_SCALARS_PER_POINT, 'h'),
-                ('scalar_name', 'S20', MAX_NB_NAMED_SCALARS_PER_POINT),
-                (Field.NB_PROPERTIES_PER_STREAMLINE, 'h'),
-                ('property_name', 'S20',
-                 MAX_NB_NAMED_PROPERTIES_PER_STREAMLINE),
-                ('reserved', 'S508'),
-                (Field.VOXEL_ORDER, 'S4'),
-                ('pad2', 'S4'),
-                ('image_orientation_patient', 'f4', 6),
-                ('pad1', 'S2'),
-                ('invert_x', 'S1'),
-                ('invert_y', 'S1'),
-                ('invert_z', 'S1'),
-                ('swap_xy', 'S1'),
-                ('swap_yz', 'S1'),
-                ('swap_zx', 'S1'),
-                (Field.NB_STREAMLINES, 'i4'),
-                ('version', 'i4'),
-                ('hdr_size', 'i4'),
-                ]
-
-# Version 2 adds a 4x4 matrix giving the affine transformtation going
+# Version 2 adds a 4x4 matrix giving the affine transformation going
 # from voxel coordinates in the referenced 3D voxel matrix, to xyz
 # coordinates (axes L->R, P->A, I->S). If (0 based) value [3, 3] from
 # this matrix is 0, this means the matrix is not recorded.
+# See http://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html
 header_2_dtd = [(Field.MAGIC_NUMBER, 'S6'),
                 (Field.DIMENSIONS, 'h', 3),
                 (Field.VOXEL_SIZES, 'f4', 3),
@@ -83,7 +58,6 @@
                 ]
 
 # Full header numpy dtypes
-header_1_dtype = np.dtype(header_1_dtd)
 header_2_dtype = np.dtype(header_2_dtd)
 
 
@@ -341,8 +315,8 @@ def load(cls, fileobj, lazy_load=False):
         data_per_point_slice = {}
         if hdr[Field.NB_SCALARS_PER_POINT] > 0:
             cpt = 0
-            for scalar_name in hdr['scalar_name']:
-                scalar_name, nb_scalars = decode_value_from_name(scalar_name)
+            for scalar_field in hdr['scalar_name']:
+                scalar_name, nb_scalars = decode_value_from_name(scalar_field)
 
                 if nb_scalars == 0:
                     continue
@@ -358,8 +332,8 @@ def load(cls, fileobj, lazy_load=False):
         data_per_streamline_slice = {}
         if hdr[Field.NB_PROPERTIES_PER_STREAMLINE] > 0:
             cpt = 0
-            for property_name in hdr['property_name']:
-                results = decode_value_from_name(property_name)
+            for property_field in hdr['property_name']:
+                results = decode_value_from_name(property_field)
                 property_name, nb_properties = results
 
                 if nb_properties == 0:
@@ -584,8 +558,8 @@ def _read_header(fileobj):
                                                  TrkFile.HEADER_SIZE))
 
             if header_rec['version'] == 1:
-                header_rec = np.fromstring(string=header_str,
-                                           dtype=header_1_dtype)
+                # There is no 4x4 matrix for voxel to RAS transformation.
+                header_rec[Field.VOXEL_TO_RASMM] = np.zeros((4, 4))
             elif header_rec['version'] == 2:
                 pass  # Nothing more to do.
             else:
@@ -724,12 +698,17 @@ def __str__(self):
                 hdr_field = getattr(Field, attr)
                 if hdr_field in vars:
                     vars[attr] = vars[hdr_field]
-        vars['scalar_names'] = '\n  '.join([asstr(s)
-                                            for s in vars['scalar_name']
-                                            if len(s) > 0])
-        vars['property_names'] = "\n  ".join([asstr(s)
-                                              for s in vars['property_name']
-                                              if len(s) > 0])
+
+        nb_scalars = self.header[Field.NB_SCALARS_PER_POINT]
+        scalar_names = [asstr(s)
+                        for s in vars['scalar_name'][:nb_scalars]
+                        if len(s) > 0]
+        vars['scalar_names'] = '\n  '.join(scalar_names)
+        nb_properties = self.header[Field.NB_PROPERTIES_PER_STREAMLINE]
+        property_names = [asstr(s)
+                          for s in vars['property_name'][:nb_properties]
+                          if len(s) > 0]
+        vars['property_names'] = "\n  ".join(property_names)
         # Make all byte strings into strings
         # Fixes recursion error on Python 3.3
         vars = dict((k, asstr(v) if hasattr(v, 'decode') else v)
@@ -739,11 +718,11 @@ def __str__(self):
 v.{version}
 dim: {DIMENSIONS}
 voxel_sizes: {VOXEL_SIZES}
-orgin: {ORIGIN}
+origin: {ORIGIN}
 nb_scalars: {NB_SCALARS_PER_POINT}
-scalar_name:\n  {scalar_names}
+scalar_names:\n  {scalar_names}
 nb_properties: {NB_PROPERTIES_PER_STREAMLINE}
-property_name:\n  {property_names}
+property_names:\n  {property_names}
 vox_to_world:\n{VOXEL_TO_RASMM}
 voxel_order: {VOXEL_ORDER}
 image_orientation_patient: {image_orientation_patient}