Merge remote-tracking branch 'upstream/master' into enh/cifti2

satra · satra · commit 9caabb97927c · 2016-07-20T10:41:23.000-04:00
* upstream/master: (23 commits)
  Better test.
  Fix.
  Test for new quad file.
  Fix to surface reading with new quad.
  TST: add freesurfer quad file for testing
  FIX: typo
  FIX: Copy affine_to_rasmm when slicing a Tractogram object
  Added a test.
  Test coverage.
  General format.
  Fix.
  String to bytes.
  _serialize_volume_info
  Address comments.
  Fixes.
  Changed from strings to np.arrays. Allow head == 20.
  Fix
  Reading of extra info.
  don't count warnings
  address comments
  ...
diff --git a/nibabel-data/nitest-freesurfer b/nibabel-data/nitest-freesurfer
@@ -1 +1 @@
-Subproject commit df9ccc64adb8f2e1e39e99c1ad22cbb89d3d3e6f
+Subproject commit 0d307865704df71c3b2248139714806aad47139d
diff --git a/nibabel/freesurfer/io.py b/nibabel/freesurfer/io.py
@@ -1,10 +1,11 @@
 from __future__ import division, print_function, absolute_import
 
+import warnings
 import numpy as np
 import getpass
 import time
 
-
+from ..externals import OrderedDict
 from ..externals.six.moves import xrange
 from ..openers import Opener
 
@@ -44,28 +45,79 @@ def _fread3_many(fobj, n):
     return (b1 << 16) + (b2 << 8) + b3
 
 
-def read_geometry(filepath):
+def _read_volume_info(fobj):
+    """Helper for reading the footer from a surface file."""
+    volume_info = OrderedDict()
+    head = np.fromfile(fobj, '>i4', 1)
+    if not np.array_equal(head, [20]):  # Read two bytes more
+        head = np.concatenate([head, np.fromfile(fobj, '>i4', 2)])
+        if not np.array_equal(head, [2, 0, 20]):
+            warnings.warn("Unknown extension code.")
+            return volume_info
+
+    volume_info['head'] = head
+    for key in ['valid', 'filename', 'volume', 'voxelsize', 'xras', 'yras',
+                'zras', 'cras']:
+        pair = fobj.readline().decode('utf-8').split('=')
+        if pair[0].strip() != key or len(pair) != 2:
+            raise IOError('Error parsing volume info.')
+        if key in ('valid', 'filename'):
+            volume_info[key] = pair[1].strip()
+        elif key == 'volume':
+            volume_info[key] = np.array(pair[1].split()).astype(int)
+        else:
+            volume_info[key] = np.array(pair[1].split()).astype(float)
+    # Ignore the rest
+    return volume_info
+
+
+def read_geometry(filepath, read_metadata=False, read_stamp=False):
     """Read a triangular format Freesurfer surface mesh.
 
     Parameters
     ----------
     filepath : str
         Path to surface file
+    read_metadata : bool
+        Read metadata as key-value pairs.
+        Valid keys:
+            * 'head' : array of int
+            * 'valid' : str
+            * 'filename' : str
+            * 'volume' : array of int, shape (3,)
+            * 'voxelsize' : array of float, shape (3,)
+            * 'xras' : array of float, shape (3,)
+            * 'yras' : array of float, shape (3,)
+            * 'zras' : array of float, shape (3,)
+            * 'cras' : array of float, shape (3,)
+    read_stamp : bool
+        Return the comment from the file
 
     Returns
     -------
     coords : numpy array
         nvtx x 3 array of vertex (x, y, z) coordinates
     faces : numpy array
         nfaces x 3 array of defining mesh triangles
+    volume_info : OrderedDict
+        If read_metadata is true, key-value pairs found in the geometry file
+    create_stamp : str
+        If read_stamp is true, the comment added by the program that saved
+        the file
     """
+    volume_info = OrderedDict()
+
+    TRIANGLE_MAGIC = 16777214
+    QUAD_MAGIC = 16777215
+    NEW_QUAD_MAGIC = 16777213
     with open(filepath, "rb") as fobj:
         magic = _fread3(fobj)
-        if magic == 16777215:  # Quad file
+        if magic in (QUAD_MAGIC, NEW_QUAD_MAGIC):  # Quad file
             nvert = _fread3(fobj)
             nquad = _fread3(fobj)
-            coords = np.fromfile(fobj, ">i2", nvert * 3).astype(np.float)
-            coords = coords.reshape(-1, 3) / 100.0
+            (fmt, div) = (">i2", 100.) if magic == QUAD_MAGIC else (">f4", 1.)
+            coords = np.fromfile(fobj, fmt, nvert * 3).astype(np.float) / div
+            coords = coords.reshape(-1, 3)
             quads = _fread3_many(fobj, nquad * 4)
             quads = quads.reshape(nquad, 4)
             #
@@ -85,21 +137,34 @@ def read_geometry(filepath):
                     faces[nface] = quad[0], quad[2], quad[3]
                     nface += 1
 
-        elif magic == 16777214:  # Triangle file
-            fobj.readline()  # create_stamp
+        elif magic == TRIANGLE_MAGIC:  # Triangle file
+            create_stamp = fobj.readline().rstrip(b'\n').decode('utf-8')
             fobj.readline()
             vnum = np.fromfile(fobj, ">i4", 1)[0]
             fnum = np.fromfile(fobj, ">i4", 1)[0]
             coords = np.fromfile(fobj, ">f4", vnum * 3).reshape(vnum, 3)
             faces = np.fromfile(fobj, ">i4", fnum * 3).reshape(fnum, 3)
+
+            if read_metadata:
+                volume_info = _read_volume_info(fobj)
         else:
             raise ValueError("File does not appear to be a Freesurfer surface")
 
     coords = coords.astype(np.float)  # XXX: due to mayavi bug on mac 32bits
-    return coords, faces
 
+    ret = (coords, faces)
+    if read_metadata:
+        if len(volume_info) == 0:
+            warnings.warn('No volume information contained in the file')
+        ret += (volume_info,)
+    if read_stamp:
+        ret += (create_stamp,)
 
-def write_geometry(filepath, coords, faces, create_stamp=None):
+    return ret
+
+
+def write_geometry(filepath, coords, faces, create_stamp=None,
+                   volume_info=None):
     """Write a triangular format Freesurfer surface mesh.
 
     Parameters
@@ -112,6 +177,19 @@ def write_geometry(filepath, coords, faces, create_stamp=None):
         nfaces x 3 array of defining mesh triangles
     create_stamp : str
         User/time stamp (default: "created by <user> on <ctime>")
+    volume_info : dict-like or None
+        Key-value pairs to encode at the end of the file.
+        Valid keys:
+            * 'head' : array of int
+            * 'valid' : str
+            * 'filename' : str
+            * 'volume' : array of int, shape (3,)
+            * 'voxelsize' : array of float, shape (3,)
+            * 'xras' : array of float, shape (3,)
+            * 'yras' : array of float, shape (3,)
+            * 'zras' : array of float, shape (3,)
+            * 'cras' : array of float, shape (3,)
+
     """
     magic_bytes = np.array([255, 255, 254], dtype=np.uint8)
 
@@ -129,6 +207,10 @@ def write_geometry(filepath, coords, faces, create_stamp=None):
         coords.astype('>f4').reshape(-1).tofile(fobj)
         faces.astype('>i4').reshape(-1).tofile(fobj)
 
+        # Add volume info, if given
+        if volume_info is not None and len(volume_info) > 0:
+            fobj.write(_serialize_volume_info(volume_info))
+
 
 def read_morph_data(filepath):
     """Read a Freesurfer morphometry data file.
@@ -369,3 +451,32 @@ def read_label(filepath, read_scalars=False):
         scalar_array = np.loadtxt(filepath, skiprows=2, usecols=[-1])
         return label_array, scalar_array
     return label_array
+
+
+def _serialize_volume_info(volume_info):
+    """Helper for serializing the volume info."""
+    keys = ['head', 'valid', 'filename', 'volume', 'voxelsize', 'xras', 'yras',
+            'zras', 'cras']
+    diff = set(volume_info.keys()).difference(keys)
+    if len(diff) > 0:
+        raise ValueError('Invalid volume info: %s.' % diff.pop())
+
+    strings = list()
+    for key in keys:
+        if key == 'head':
+            if not (np.array_equal(volume_info[key], [20]) or np.array_equal(
+                    volume_info[key], [2, 0, 20])):
+                warnings.warn("Unknown extension code.")
+            strings.append(np.array(volume_info[key], dtype='>i4').tostring())
+        elif key in ('valid', 'filename'):
+            val = volume_info[key]
+            strings.append('{0} = {1}\n'.format(key, val).encode('utf-8'))
+        elif key == 'volume':
+            val = volume_info[key]
+            strings.append('{0} = {1} {2} {3}\n'.format(
+                key, val[0], val[1], val[2]).encode('utf-8'))
+        else:
+            val = volume_info[key]
+            strings.append('{0} = {1:0.10g} {2:0.10g} {3:0.10g}\n'.format(
+                key.ljust(6), val[0], val[1], val[2]).encode('utf-8'))
+    return b''.join(strings)
diff --git a/nibabel/freesurfer/tests/test_io.py b/nibabel/freesurfer/tests/test_io.py
@@ -4,20 +4,21 @@
 import getpass
 import time
 import hashlib
+import warnings
 
 
 from ...tmpdirs import InTemporaryDirectory
 
 from nose.tools import assert_true
 import numpy as np
-from numpy.testing import assert_equal, assert_raises, dec
+from numpy.testing import assert_equal, assert_raises, dec, assert_allclose
 
 from .. import (read_geometry, read_morph_data, read_annot, read_label,
                 write_geometry, write_morph_data, write_annot)
 
-from ...tests.nibabel_data import get_nibabel_data
+from ...tests.nibabel_data import get_nibabel_data, needs_nibabel_data
 from ...fileslice import strided_scalar
-
+from ...testing import clear_and_catch_warnings
 
 DATA_SDIR = 'fsaverage'
 
@@ -54,26 +55,53 @@ def test_geometry():
     assert_equal(0, faces.min())
     assert_equal(coords.shape[0], faces.max() + 1)
 
-    # Test quad with sphere
     surf_path = pjoin(data_path, "surf", "%s.%s" % ("lh", "sphere"))
-    coords, faces = read_geometry(surf_path)
+    coords, faces, volume_info, create_stamp = read_geometry(
+        surf_path, read_metadata=True, read_stamp=True)
+
     assert_equal(0, faces.min())
     assert_equal(coords.shape[0], faces.max() + 1)
+    assert_equal(9, len(volume_info))
+    assert_equal([2, 0, 20], volume_info['head'])
+    assert_equal(u'created by greve on Thu Jun  8 19:17:51 2006',
+                 create_stamp)
 
     # Test equivalence of freesurfer- and nibabel-generated triangular files
     # with respect to read_geometry()
     with InTemporaryDirectory():
         surf_path = 'test'
         create_stamp = "created by %s on %s" % (getpass.getuser(),
                                                 time.ctime())
-        write_geometry(surf_path, coords, faces, create_stamp)
+        volume_info['cras'] = [1., 2., 3.]
+        write_geometry(surf_path, coords, faces, create_stamp, volume_info)
 
-        coords2, faces2 = read_geometry(surf_path)
+        coords2, faces2, volume_info2 = \
+            read_geometry(surf_path, read_metadata=True)
 
+        for key in ('xras', 'yras', 'zras', 'cras'):
+            assert_allclose(volume_info2[key], volume_info[key],
+                            rtol=1e-7, atol=1e-30)
+        assert_equal(volume_info2['cras'], volume_info['cras'])
         with open(surf_path, 'rb') as fobj:
             np.fromfile(fobj, ">u1", 3)
             read_create_stamp = fobj.readline().decode().rstrip('\n')
 
+        # now write an incomplete file
+        write_geometry(surf_path, coords, faces)
+        with clear_and_catch_warnings() as w:
+            warnings.filterwarnings('always', category=DeprecationWarning)
+            read_geometry(surf_path, read_metadata=True)
+        assert_true(any('volume information contained' in str(ww.message)
+                        for ww in w))
+        assert_true(any('extension code' in str(ww.message) for ww in w))
+        volume_info['head'] = [1, 2]
+        with clear_and_catch_warnings() as w:
+            write_geometry(surf_path, coords, faces, create_stamp, volume_info)
+        assert_true(any('Unknown extension' in str(ww.message) for ww in w))
+        volume_info['a'] = 0
+        assert_raises(ValueError, write_geometry, surf_path, coords,
+                      faces, create_stamp, volume_info)
+
     assert_equal(create_stamp, read_create_stamp)
 
     np.testing.assert_array_equal(coords, coords2)
@@ -86,6 +114,23 @@ def test_geometry():
     np.testing.assert_array_equal(faces_swapped, faces)
 
 
+@freesurfer_test
+@needs_nibabel_data('nitest-freesurfer')
+def test_quad_geometry():
+    """Test IO of freesurfer quad files."""
+    new_quad = pjoin(get_nibabel_data(), 'nitest-freesurfer', 'subjects',
+                     'bert', 'surf', 'lh.inflated.nofix')
+    coords, faces = read_geometry(new_quad)
+    assert_equal(0, faces.min())
+    assert_equal(coords.shape[0], faces.max() + 1)
+    with InTemporaryDirectory():
+        new_path = 'test'
+        write_geometry(new_path, coords, faces)
+        coords2, faces2 = read_geometry(new_path)
+        assert_equal(coords, coords2)
+        assert_equal(faces, faces2)
+
+
 @freesurfer_test
 def test_morph_data():
     """Test IO of morphometry data file (eg. curvature)."""
diff --git a/nibabel/streamlines/tests/test_tractogram.py b/nibabel/streamlines/tests/test_tractogram.py
@@ -18,6 +18,7 @@
 
 def setup():
     global DATA
+    DATA['rng'] = np.random.RandomState(1234)
     DATA['streamlines'] = [np.arange(1*3, dtype="f4").reshape((1, 3)),
                            np.arange(2*3, dtype="f4").reshape((2, 3)),
                            np.arange(5*3, dtype="f4").reshape((5, 3))]
@@ -373,6 +374,13 @@ def test_tractogram_getitem(self):
                          DATA['tractogram'].data_per_streamline[::-1],
                          DATA['tractogram'].data_per_point[::-1])
 
+        # Make sure slicing conserves the affine_to_rasmm property.
+        tractogram = DATA['tractogram'].copy()
+        tractogram.affine_to_rasmm = DATA['rng'].rand(4, 4)
+        tractogram_view = tractogram[::2]
+        assert_array_equal(tractogram_view.affine_to_rasmm,
+                           tractogram.affine_to_rasmm)
+
     def test_tractogram_add_new_data(self):
         # Tractogram with only streamlines
         t = DATA['simple_tractogram'].copy()
diff --git a/nibabel/streamlines/tractogram.py b/nibabel/streamlines/tractogram.py
@@ -335,7 +335,8 @@ def __getitem__(self, idx):
         if isinstance(idx, (numbers.Integral, np.integer)):
             return TractogramItem(pts, data_per_streamline, data_per_point)
 
-        return Tractogram(pts, data_per_streamline, data_per_point)
+        return Tractogram(pts, data_per_streamline, data_per_point,
+                          affine_to_rasmm=self.affine_to_rasmm)
 
     def __len__(self):
         return len(self.streamlines)
