Add function to save FreeSurfer annotation files.

nibabel/freesurfer/__init__.py

@@ -2,5 +2,5 @@
 """
 
 from .io import read_geometry, read_morph_data, \
-               read_annot, read_label, write_geometry
+    read_annot, read_label, write_geometry, write_annot
 from .mghformat import load, save, MGHImage

nibabel/freesurfer/io.py

@@ -36,7 +36,7 @@ def _fread3_many(fobj, n):
         An array of 3 byte int
     """
     b1, b2, b3 = np.fromfile(fobj, ">u1", 3 * n).reshape(-1,
-                                                    3).astype(np.int).T
+                                                         3).astype(np.int).T
     return (b1 << 16) + (b2 << 8) + b3
 
 
@@ -113,7 +113,7 @@ def write_geometry(filepath, coords, faces, create_stamp=None):
 
     if create_stamp is None:
         create_stamp = "created by %s on %s" % (getpass.getuser(),
-            time.ctime())
+                                                time.ctime())
 
     with open(filepath, 'wb') as fobj:
         magic_bytes.tofile(fobj)
@@ -184,6 +184,7 @@ def read_annot(filepath, orig_ids=False):
         vnum = np.fromfile(fobj, dt, 1)[0]
         data = np.fromfile(fobj, dt, vnum * 2).reshape(vnum, 2)
         labels = data[:, 1]
+
         ctab_exists = np.fromfile(fobj, dt, 1)[0]
         if not ctab_exists:
             raise Exception('Color table not found in annotation file')
@@ -220,7 +221,7 @@ def read_annot(filepath, orig_ids=False):
                 names.append(name)
                 ctab[i, :4] = np.fromfile(fobj, dt, 4)
                 ctab[i, 4] = (ctab[i, 0] + ctab[i, 1] * (2 ** 8) +
-                                ctab[i, 2] * (2 ** 16))
+                              ctab[i, 2] * (2 ** 16))
         ctab[:, 3] = 255
     if not orig_ids:
         ord = np.argsort(ctab[:, -1])
@@ -230,6 +231,68 @@ def read_annot(filepath, orig_ids=False):
     return labels, ctab, names
 
 
+def write_annot(filepath, labels, ctab, names):
+    """Write out a Freesurfer annotation file.
+
+    See:
+    http://ftp.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles#Annotation
+
+    Parameters
+    ----------
+    filepath : str
+        Path to annotation file to be written
+    labels : ndarray, shape (n_vertices,)
+        Annotation id at each vertex.
+    ctab : ndarray, shape (n_labels, 5)
+        RGBA + label id colortable array.
+    names : list of str
+        The names of the labels. The length of the list is n_labels.
+    """
+    with open(filepath, "wb") as fobj:
+        dt = ">i4"
+        vnum = len(labels)
+
+        def write(num, dtype=dt):
+            np.array([num]).astype(dtype).tofile(fobj)
+
+        def write_string(s):
+            write(len(s))
+            write(s, dtype='|S%d' % len(s))
+
+        # vtxct
+        write(vnum)
+
+        # convert labels into coded CLUT values
+        clut_labels = ctab[:, -1][labels]
+        clut_labels[np.where(labels == -1)] = 0
+
+        # vno, label
+        data = np.vstack((np.array(range(vnum)).astype(dt),
+                          clut_labels.astype(dt))).T
+        data.byteswap().tofile(fobj)
+
+        # tag
+        write(1)
+
+        # ctabversion
+        write(-2)
+
+        # maxstruc
+        write(np.max(labels) + 1)
+
+        # File of LUT is unknown.
+        write_string('NOFILE')
+
+        # num_entries
+        write(ctab.shape[0])
+
+        for ind, (clu, name) in enumerate(zip(ctab, names)):
+            write(ind)
+            write_string(name)
+            for val in clu[:-1]:
+                write(val)
+
+
 def read_label(filepath, read_scalars=False):
     """Load in a Freesurfer .label file.
 

nibabel/freesurfer/tests/test_io.py

@@ -11,7 +11,7 @@
 from numpy.testing import assert_equal
 
 from .. import read_geometry, read_morph_data, read_annot, read_label, \
-                write_geometry
+    write_geometry, write_annot
 
 
 have_freesurfer = True
@@ -47,7 +47,7 @@ def test_geometry():
     with InTemporaryDirectory():
         surf_path = 'test'
         create_stamp = "created by %s on %s" % (getpass.getuser(),
-            time.ctime())
+                                                time.ctime())
         write_geometry(surf_path, coords, faces, create_stamp)
 
         coords2, faces2 = read_geometry(surf_path)
@@ -87,11 +87,28 @@ def test_annot():
         assert_true(labels.shape == (163842, ))
         assert_true(ctab.shape == (len(names), 5))
 
+        labels_orig = None
         if a == 'aparc':
             labels_orig, _, _ = read_annot(annot_path, orig_ids=True)
             np.testing.assert_array_equal(labels == -1, labels_orig == 0)
             assert_true(np.sum(labels_orig == 0) > 0)
 
+        # Test equivalence of freesurfer- and nibabel-generated annot files
+        # with respect to read_annot()
+        with InTemporaryDirectory():
+            annot_path = 'test'
+            write_annot(annot_path, labels, ctab, names)
+
+            labels2, ctab2, names2 = read_annot(annot_path)
+            if labels_orig is not None:
+                labels_orig_2, _, _ = read_annot(annot_path, orig_ids=True)
+
+        np.testing.assert_array_equal(labels, labels2)
+        if labels_orig is not None:
+            np.testing.assert_array_equal(labels_orig, labels_orig_2)
+        np.testing.assert_array_equal(ctab, ctab2)
+        assert_equal(names, names2)
+
 
 @freesurfer_test
 def test_label():