RF: Reducing read_annot function size. .annot data type is now

module-level. Other small doc adjustments.

Author: pauldmccarthy

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
 
@@ -74,8 +82,10 @@ def _read_volume_info(fobj):
 
 
 def _pack_rgba(rgba):
-    """Used by :meth:`read_annot` and :meth:`write_annot` to pack an RGBA
-    sequence into a single integer.
+    """Pack an RGBA sequence into a single integer.
+
+    Used by :func:`read_annot` and :func:`write_annot` to generate
+    "annotation values" for a Freesuerfer `.annot` file.
 
     Parameters
     ----------
@@ -86,7 +96,7 @@ def _pack_rgba(rgba):
     Returns
     -------
 
-    out : ndarray, shape (n, )
+    out : ndarray, shape (n, 1)
         Annotation values for each colour.
     """
     bitshifts = 2 ** np.array([[0], [8], [16], [24]], dtype=rgba.dtype)
@@ -316,7 +326,20 @@ 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 colours corresponding to each label.
+
+    The colour table itself may be stored in either an "old-style" format, or
+    a "new-style" format - the :func:`_read_annot_ctab_old_format` and
+    :func:`_read_annot_ctab_new_format` functions are respectively used to
+    read in the colour table.
+
+    See:
+    https://surfer.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles#Annotation
+    https://github.com/freesurfer/freesurfer/blob/dev/matlab/read_annotation.m
 
     Parameters
     ----------
@@ -336,11 +359,10 @@ def read_annot(filepath, orig_ids=False):
         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, "rb") as fobj:
-        # all data (apart from strings) in an .annot file is stored as
-        # big-endian int32
-        dt = ">i4"
+        dt = _ANNOT_DT
 
         # number of vertices
         vnum = np.fromfile(fobj, dt, 1)[0]
@@ -355,50 +377,16 @@ def read_annot(filepath, orig_ids=False):
             raise Exception('Color table not found in annotation file')
 
         # in old-format files, the next field will contain the number of
-        # entries in the colour table. In new-format files, this will be
+        # entries in the colour 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:
-
-            # 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)
+            ctab, names = _read_annot_ctab_old_format(fobj, n_entries)
         # We've got a new-format .annot file
         else:
-            # file version number
-            ctab_version = -n_entries
-            if ctab_version != 2:
-                raise Exception('Color table version not supported')
-            # 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)
+            ctab, names = _read_annot_ctab_new_format(fobj, -n_entries)
 
     # generate annotation values for each LUT entry
     ctab[:, [4]] = _pack_rgba(ctab[:, :4])
@@ -414,11 +402,101 @@ def read_annot(filepath, orig_ids=False):
     return labels, ctab, names
 
 
+def _read_annot_ctab_old_format(fobj, n_entries):
+    """Read in an old-style Freesurfer colour 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 colour table data.
+    n_entries : int
+        Number of entries in the colour 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.
+    """
+    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 colour 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 colour table data.
+    ctab_version : int
+        Colour 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.
+    """
+    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 Freesurfer annotation file.
 
     See:
     https://surfer.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles#Annotation
+    https://github.com/freesurfer/freesurfer/blob/dev/matlab/write_annotation.m
 
     Parameters
     ----------
@@ -430,14 +508,14 @@ def write_annot(filepath, labels, ctab, names, fill_ctab=True):
         RGBA + label id colortable array.
     names : list of str
         The names of the labels. The length of the list is n_labels.
-    fill_ctab : bool
+    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):