Changed from strings to np.arrays. Allow head == 20.

Author: jaeilepp

nibabel/freesurfer/io.py

@@ -47,18 +47,26 @@ def _fread3_many(fobj, n):
 
 def _read_volume_info(fobj):
     volume_info = OrderedDict()
-    head = np.fromfile(fobj, '>i4', 3)
-    if any(head != [2, 0, 20]):
-        warnings.warn("Unknown extension code.")
-    else:
-        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.')
-            volume_info[pair[0]] = pair[1]
-        # Ignore the rest
+    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
 
 
@@ -122,14 +130,17 @@ def read_geometry(filepath, read_metadata=False, read_stamp=False):
             coords = np.fromfile(fobj, ">f4", vnum * 3).reshape(vnum, 3)
             faces = np.fromfile(fobj, ">i4", fnum * 3).reshape(fnum, 3)
 
-            volume_info = _read_volume_info(fobj)
+            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
 
     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,)
@@ -174,9 +185,15 @@ def write_geometry(filepath, coords, faces, create_stamp=None,
         if volume_info is not None and len(volume_info) > 0:
             for key, val in volume_info.items():
                 if key == 'head':
-                    val.tofile(fobj)
-                    continue
-                fobj.write('{0}={1}'.format(key, val).encode('utf-8'))
+                    np.array(val, dtype='>i4').tofile(fobj)
+                elif key in ('valid', 'filename'):
+                    fobj.write('{0} = {1}\n'.format(key, val).encode('utf-8'))
+                elif key == 'volume':
+                    fobj.write('{0} = {1} {2} {3}\n'.format(
+                        key, val[0], val[1], val[2]).encode('utf-8'))
+                else:
+                    fobj.write('{0} = {1:.4f} {2:.4f} {3:.4f}\n'.format(
+                        key.ljust(6), val[0], val[1], val[2]).encode('utf-8'))
 
 
 def read_morph_data(filepath):

nibabel/freesurfer/tests/test_io.py

@@ -55,36 +55,43 @@ def test_geometry():
     assert_equal(0, faces.min())
     assert_equal(coords.shape[0], faces.max() + 1)
 
-    # Test quad with sphere
-    with clear_and_catch_warnings():
-        warnings.filterwarnings('always', category=DeprecationWarning)
-        surf_path = pjoin(data_path, "surf", "%s.%s" % ("lh", "sphere"))
-        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)
+    surf_path = pjoin(data_path, "surf", "%s.%s" % ("lh", "sphere"))
+    warnings.filterwarnings('always', category=DeprecationWarning)
+    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())
-        volume_info['cras   '] = '1. 2. 3.'
+        volume_info['cras'] = [1., 2., 3.]
         write_geometry(surf_path, coords, faces, create_stamp, volume_info)
 
         coords2, faces2, volume_info2 = \
             read_geometry(surf_path, read_metadata=True)
 
-        assert_equal(volume_info2['cras   '], volume_info['cras   '])
+        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_equal(create_stamp, read_create_stamp)
 
     np.testing.assert_array_equal(coords, coords2)