NF: Handle metadata in read/write_geometry

Author: effigies

nibabel/freesurfer/io.py

@@ -3,6 +3,7 @@
 import numpy as np
 import getpass
 import time
+from collections import OrderedDict
 
 
 from ..externals.six.moves import xrange
@@ -44,21 +45,32 @@ def _fread3_many(fobj, n):
     return (b1 << 16) + (b2 << 8) + b3
 
 
-def read_geometry(filepath):
+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
+    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 : dict-like
+        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 = None
+
     with open(filepath, "rb") as fobj:
         magic = _fread3(fobj)
         if magic == 16777215:  # Quad file
@@ -86,20 +98,46 @@ def read_geometry(filepath):
                     nface += 1
 
         elif magic == 16777214:  # Triangle file
-            fobj.readline()  # create_stamp
+            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)
+
+            extra = fobj.read() if read_metadata else b''
+            if extra:
+                if extra[:4] != b'\x00\x00\x00\x14':
+                    warnings.warn("Unknown extension code.")
+                volume_info = OrderedDict()
+                try:
+                    for line in extra[4:].split(b'\n'):
+                        key, val = map(bytes.strip, line.split(b'=', 1))
+                        key = key.decode('utf-8')
+                        if key in ('voxelsize', 'xras', 'yras', 'zras'):
+                            val = np.fromstring(val, sep=' ')
+                        elif key == 'volume':
+                            val = np.fromstring(val, sep=' ', dtype=np.uint)
+                        volume_info[key] = val
+                except ValueError:
+                    raise ValueError("Error parsing volume info")
+
         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:
+        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,13 +150,27 @@ 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
     """
     magic_bytes = np.array([255, 255, 254], dtype=np.uint8)
 
     if create_stamp is None:
         create_stamp = "created by %s on %s" % (getpass.getuser(),
                                                 time.ctime())
 
+    postlude = b''
+    if volume_info is not None:
+        postlude = [b'\x00\x00\x00\x14']
+        for key, val in volume_info.items():
+            if key in ('voxelsize', 'xras', 'yras', 'zras'):
+                val = '{:.3f} {:.3f} {:.3f}'.format(*val)
+            elif key == 'volume':
+                val = '{:d} {:d} {:d}'.format(*val)
+            key = key.ljust(6)
+            postlude.append('{} = {}'.format(key, val).encode('utf-8'))
+        postlude = b'\n'.join(postlude)
+
     with open(filepath, 'wb') as fobj:
         magic_bytes.tofile(fobj)
         fobj.write(("%s\n\n" % create_stamp).encode('utf-8'))
@@ -129,6 +181,9 @@ 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
+        fobj.write(postlude)
+
 
 def read_morph_data(filepath):
     """Read a Freesurfer morphometry data file.