Merge pull request #592 from pauldmccarthy/fix_freesurfer_annot

MRG: Fix freesurfer `read_annot`, and enhance `write_annot`

This PR comprises two sections, related to the nibabel.freesurfer.io.read_annot and write_annot functions. The first involves changes to read_annot, to fix some bugs, and the second involves enhancements to write_annot, to make saving .annot files a little easier.

Author: matthew-brett

nibabel/freesurfer/io.py

@@ -12,6 +12,14 @@
 from ..openers import Opener
 
 
+_ANNOT_DT = ">i4"
+"""Data type for Freesurfer `.annot` files.
+
+Used by :func:`read_annot` and :func:`write_annot`.  All data (apart from
+strings) in an `.annot` file is stored as big-endian int32.
+"""
+
+
 def _fread3(fobj):
     """Read a 3-byte int from an open binary file object
 
@@ -73,6 +81,26 @@ def _read_volume_info(fobj):
     return volume_info
 
 
+def _pack_rgba(rgba):
+    """Pack an RGBA sequence into a single integer.
+
+    Used by :func:`read_annot` and :func:`write_annot` to generate
+    "annotation values" for a Freesurfer ``.annot`` file.
+
+    Parameters
+    ----------
+    rgba : ndarray, shape (n, 4)
+        RGBA colors
+
+    Returns
+    -------
+    out : ndarray, shape (n, 1)
+        Annotation values for each color.
+    """
+    bitshifts = 2 ** np.array([[0], [8], [16], [24]], dtype=rgba.dtype)
+    return rgba.dot(bitshifts)
+
+
 def read_geometry(filepath, read_metadata=False, read_stamp=False):
     """Read a triangular format Freesurfer surface mesh.
 
@@ -296,7 +324,18 @@ def write_morph_data(file_like, values, fnum=0):
 
 
 def read_annot(filepath, orig_ids=False):
-    """Read in a Freesurfer annotation from a .annot file.
+    """Read in a Freesurfer annotation from a ``.annot`` file.
+
+    An ``.annot`` file contains a sequence of vertices with a label (also known
+    as an "annotation value") associated with each vertex, and then a sequence
+    of colors corresponding to each label.
+
+    Annotation file format versions 1 and 2 are supported, corresponding to
+    the "old-style" and "new-style" color table layout.
+
+    See:
+     * https://surfer.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles#Annotation
+     * https://github.com/freesurfer/freesurfer/blob/dev/matlab/read_annotation.m
 
     Parameters
     ----------
@@ -314,55 +353,38 @@ def read_annot(filepath, orig_ids=False):
         to any label and orig_ids=False, its id will be set to -1.
     ctab : ndarray, shape (n_labels, 5)
         RGBA + label id colortable array.
-    names : list of str
+    names : list of str (python 2), list of bytes (python 3)
         The names of the labels. The length of the list is n_labels.
     """
     with open(filepath, "rb") as fobj:
-        dt = ">i4"
+        dt = _ANNOT_DT
+
+        # number of vertices
         vnum = np.fromfile(fobj, dt, 1)[0]
+
+        # vertex ids + annotation values
         data = np.fromfile(fobj, dt, vnum * 2).reshape(vnum, 2)
         labels = data[:, 1]
 
+        # is there a color table?
         ctab_exists = np.fromfile(fobj, dt, 1)[0]
         if not ctab_exists:
             raise Exception('Color table not found in annotation file')
+
+        # in old-format files, the next field will contain the number of
+        # entries in the color table. In new-format files, this must be
+        # equal to -2
         n_entries = np.fromfile(fobj, dt, 1)[0]
+
+        # We've got an old-format .annot file.
         if n_entries > 0:
-            length = np.fromfile(fobj, dt, 1)[0]
-            orig_tab = np.fromfile(fobj, '>c', length)
-            orig_tab = orig_tab[:-1]
-
-            names = list()
-            ctab = np.zeros((n_entries, 5), np.int)
-            for i in xrange(n_entries):
-                name_length = np.fromfile(fobj, dt, 1)[0]
-                name = np.fromfile(fobj, "|S%d" % name_length, 1)[0]
-                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, 3] * (2 ** 24))
+            ctab, names = _read_annot_ctab_old_format(fobj, n_entries)
+        # We've got a new-format .annot file
         else:
-            ctab_version = -n_entries
-            if ctab_version != 2:
-                raise Exception('Color table version not supported')
-            n_entries = np.fromfile(fobj, dt, 1)[0]
-            ctab = np.zeros((n_entries, 5), np.int)
-            length = np.fromfile(fobj, dt, 1)[0]
-            np.fromfile(fobj, "|S%d" % length, 1)[0]  # Orig table path
-            entries_to_read = np.fromfile(fobj, dt, 1)[0]
-            names = list()
-            for i in xrange(entries_to_read):
-                np.fromfile(fobj, dt, 1)[0]  # Structure
-                name_length = np.fromfile(fobj, dt, 1)[0]
-                name = np.fromfile(fobj, "|S%d" % name_length, 1)[0]
-                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[:, 3] = 255
-
-    labels = labels.astype(np.int)
+            ctab, names = _read_annot_ctab_new_format(fobj, -n_entries)
+
+    # generate annotation values for each LUT entry
+    ctab[:, [4]] = _pack_rgba(ctab[:, :4])
 
     if not orig_ids:
         ord = np.argsort(ctab[:, -1])
@@ -372,11 +394,104 @@ def read_annot(filepath, orig_ids=False):
     return labels, ctab, names
 
 
-def write_annot(filepath, labels, ctab, names):
-    """Write out a Freesurfer annotation file.
+def _read_annot_ctab_old_format(fobj, n_entries):
+    """Read in an old-style Freesurfer color table from `fobj`.
+
+    This function is used by :func:`read_annot`.
+
+    Parameters
+    ----------
+
+    fobj : file-like
+        Open file handle to a Freesurfer `.annot` file, with seek point
+        at the beginning of the color table data.
+    n_entries : int
+        Number of entries in the color table.
+
+    Returns
+    -------
+
+    ctab : ndarray, shape (n_entries, 5)
+        RGBA colortable array - the last column contains all zeros.
+    names : list of str
+        The names of the labels. The length of the list is n_entries.
+    """
+    assert hasattr(fobj, 'read')
+
+    dt = _ANNOT_DT
+    # orig_tab string length + string
+    length = np.fromfile(fobj, dt, 1)[0]
+    orig_tab = np.fromfile(fobj, '>c', length)
+    orig_tab = orig_tab[:-1]
+    names = list()
+    ctab = np.zeros((n_entries, 5), dt)
+    for i in xrange(n_entries):
+        # structure name length + string
+        name_length = np.fromfile(fobj, dt, 1)[0]
+        name = np.fromfile(fobj, "|S%d" % name_length, 1)[0]
+        names.append(name)
+        # read RGBA for this entry
+        ctab[i, :4] = np.fromfile(fobj, dt, 4)
+
+    return ctab, names
+
+
+def _read_annot_ctab_new_format(fobj, ctab_version):
+    """Read in a new-style Freesurfer color table from `fobj`.
+
+    This function is used by :func:`read_annot`.
+
+    Parameters
+    ----------
+
+    fobj : file-like
+        Open file handle to a Freesurfer `.annot` file, with seek point
+        at the beginning of the color table data.
+    ctab_version : int
+        Color table format version - must be equal to 2
+
+    Returns
+    -------
+
+    ctab : ndarray, shape (n_labels, 5)
+        RGBA colortable array - the last column contains all zeros.
+    names : list of str
+        The names of the labels. The length of the list is n_labels.
+    """
+    assert hasattr(fobj, 'read')
+
+    dt = _ANNOT_DT
+    # This code works with a file version == 2, nothing else
+    if ctab_version != 2:
+        raise Exception('Unrecognised .annot file version (%i)', ctab_version)
+    # maximum LUT index present in the file
+    max_index = np.fromfile(fobj, dt, 1)[0]
+    ctab = np.zeros((max_index, 5), dt)
+    # orig_tab string length + string
+    length = np.fromfile(fobj, dt, 1)[0]
+    np.fromfile(fobj, "|S%d" % length, 1)[0]  # Orig table path
+    # number of LUT entries present in the file
+    entries_to_read = np.fromfile(fobj, dt, 1)[0]
+    names = list()
+    for _ in xrange(entries_to_read):
+        # index of this entry
+        idx = np.fromfile(fobj, dt, 1)[0]
+        # structure name length + string
+        name_length = np.fromfile(fobj, dt, 1)[0]
+        name = np.fromfile(fobj, "|S%d" % name_length, 1)[0]
+        names.append(name)
+        # RGBA
+        ctab[idx, :4] = np.fromfile(fobj, dt, 4)
+
+    return ctab, names
+
+
+def write_annot(filepath, labels, ctab, names, fill_ctab=True):
+    """Write out a "new-style" Freesurfer annotation file.
 
     See:
-    https://surfer.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles#Annotation
+     * https://surfer.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles#Annotation
+     * https://github.com/freesurfer/freesurfer/blob/dev/matlab/write_annotation.m
 
     Parameters
     ----------
@@ -388,18 +503,31 @@ def write_annot(filepath, labels, ctab, names):
         RGBA + label id colortable array.
     names : list of str
         The names of the labels. The length of the list is n_labels.
+    fill_ctab : {True, False} optional
+        If True, the annotation values for each vertex  are automatically
+        generated. In this case, the provided `ctab` may have shape
+        (n_labels, 4) or (n_labels, 5) - if the latter, the final column is
+        ignored.
     """
     with open(filepath, "wb") as fobj:
-        dt = ">i4"
+        dt = _ANNOT_DT
         vnum = len(labels)
 
         def write(num, dtype=dt):
             np.array([num]).astype(dtype).tofile(fobj)
 
         def write_string(s):
+            s = (s if isinstance(s, bytes) else s.encode()) + b'\x00'
             write(len(s))
             write(s, dtype='|S%d' % len(s))
 
+        # Generate annotation values for each ctab entry
+        if fill_ctab:
+            ctab = np.hstack((ctab[:, :4], _pack_rgba(ctab[:, :4])))
+        elif not np.array_equal(ctab[:, [4]], _pack_rgba(ctab[:, :4])):
+            warnings.warn('Annotation values in {} will be incorrect'.format(
+                filepath))
+
         # vtxct
         write(vnum)