1""" Read / write FreeSurfer geometry, morphometry, label, annotation formats
2"""
3
4import warnings
5import numpy as np
6import getpass
7import time
8
9from collections import OrderedDict
10from ..openers import Opener
11
12
13_ANNOT_DT = ">i4"
14"""Data type for Freesurfer `.annot` files.
15
16Used by :func:`read_annot` and :func:`write_annot`.  All data (apart from
17strings) in an `.annot` file is stored as big-endian int32.
18"""
19
20
21def _fread3(fobj):
22    """Read a 3-byte int from an open binary file object
23
24    Parameters
25    ----------
26    fobj : file
27        File descriptor
28
29    Returns
30    -------
31    n : int
32        A 3 byte int
33    """
34    b1, b2, b3 = np.fromfile(fobj, ">u1", 3)
35    return (b1 << 16) + (b2 << 8) + b3
36
37
38def _fread3_many(fobj, n):
39    """Read 3-byte ints from an open binary file object.
40
41    Parameters
42    ----------
43    fobj : file
44        File descriptor
45
46    Returns
47    -------
48    out : 1D array
49        An array of 3 byte int
50    """
51    b1, b2, b3 = np.fromfile(fobj, ">u1", 3 * n).reshape(-1,
52                                                         3).astype(int).T
53    return (b1 << 16) + (b2 << 8) + b3
54
55
56def _read_volume_info(fobj):
57    """Helper for reading the footer from a surface file."""
58    volume_info = OrderedDict()
59    head = np.fromfile(fobj, '>i4', 1)
60    if not np.array_equal(head, [20]):  # Read two bytes more
61        head = np.concatenate([head, np.fromfile(fobj, '>i4', 2)])
62        if not np.array_equal(head, [2, 0, 20]):
63            warnings.warn("Unknown extension code.")
64            return volume_info
65
66    volume_info['head'] = head
67    for key in ['valid', 'filename', 'volume', 'voxelsize', 'xras', 'yras',
68                'zras', 'cras']:
69        pair = fobj.readline().decode('utf-8').split('=')
70        if pair[0].strip() != key or len(pair) != 2:
71            raise IOError('Error parsing volume info.')
72        if key in ('valid', 'filename'):
73            volume_info[key] = pair[1].strip()
74        elif key == 'volume':
75            volume_info[key] = np.array(pair[1].split()).astype(int)
76        else:
77            volume_info[key] = np.array(pair[1].split()).astype(float)
78    # Ignore the rest
79    return volume_info
80
81
82def _pack_rgb(rgb):
83    """Pack an RGB sequence into a single integer.
84
85    Used by :func:`read_annot` and :func:`write_annot` to generate
86    "annotation values" for a Freesurfer ``.annot`` file.
87
88    Parameters
89    ----------
90    rgb : ndarray, shape (n, 3)
91        RGB colors
92
93    Returns
94    -------
95    out : ndarray, shape (n, 1)
96        Annotation values for each color.
97    """
98    bitshifts = 2 ** np.array([[0], [8], [16]], dtype=rgb.dtype)
99    return rgb.dot(bitshifts)
100
101
102def read_geometry(filepath, read_metadata=False, read_stamp=False):
103    """Read a triangular format Freesurfer surface mesh.
104
105    Parameters
106    ----------
107    filepath : str
108        Path to surface file.
109    read_metadata : bool, optional
110        If True, read and return metadata as key-value pairs.
111
112        Valid keys:
113
114        * 'head' : array of int
115        * 'valid' : str
116        * 'filename' : str
117        * 'volume' : array of int, shape (3,)
118        * 'voxelsize' : array of float, shape (3,)
119        * 'xras' : array of float, shape (3,)
120        * 'yras' : array of float, shape (3,)
121        * 'zras' : array of float, shape (3,)
122        * 'cras' : array of float, shape (3,)
123
124    read_stamp : bool, optional
125        Return the comment from the file
126
127    Returns
128    -------
129    coords : numpy array
130        nvtx x 3 array of vertex (x, y, z) coordinates.
131    faces : numpy array
132        nfaces x 3 array of defining mesh triangles.
133    volume_info : OrderedDict
134        Returned only if `read_metadata` is True.  Key-value pairs found in the
135        geometry file.
136    create_stamp : str
137        Returned only if `read_stamp` is True.  The comment added by the
138        program that saved the file.
139    """
140    volume_info = OrderedDict()
141
142    TRIANGLE_MAGIC = 16777214
143    QUAD_MAGIC = 16777215
144    NEW_QUAD_MAGIC = 16777213
145    with open(filepath, "rb") as fobj:
146        magic = _fread3(fobj)
147        if magic in (QUAD_MAGIC, NEW_QUAD_MAGIC):  # Quad file
148            nvert = _fread3(fobj)
149            nquad = _fread3(fobj)
150            (fmt, div) = (">i2", 100.) if magic == QUAD_MAGIC else (">f4", 1.)
151            coords = np.fromfile(fobj, fmt, nvert * 3).astype(np.float64) / div
152            coords = coords.reshape(-1, 3)
153            quads = _fread3_many(fobj, nquad * 4)
154            quads = quads.reshape(nquad, 4)
155            #
156            #   Face splitting follows
157            #
158            faces = np.zeros((2 * nquad, 3), dtype=int)
159            nface = 0
160            for quad in quads:
161                if (quad[0] % 2) == 0:
162                    faces[nface] = quad[0], quad[1], quad[3]
163                    nface += 1
164                    faces[nface] = quad[2], quad[3], quad[1]
165                    nface += 1
166                else:
167                    faces[nface] = quad[0], quad[1], quad[2]
168                    nface += 1
169                    faces[nface] = quad[0], quad[2], quad[3]
170                    nface += 1
171
172        elif magic == TRIANGLE_MAGIC:  # Triangle file
173            create_stamp = fobj.readline().rstrip(b'\n').decode('utf-8')
174            fobj.readline()
175            vnum = np.fromfile(fobj, ">i4", 1)[0]
176            fnum = np.fromfile(fobj, ">i4", 1)[0]
177            coords = np.fromfile(fobj, ">f4", vnum * 3).reshape(vnum, 3)
178            faces = np.fromfile(fobj, ">i4", fnum * 3).reshape(fnum, 3)
179
180            if read_metadata:
181                volume_info = _read_volume_info(fobj)
182        else:
183            raise ValueError("File does not appear to be a Freesurfer surface")
184
185    coords = coords.astype(np.float64)  # XXX: due to mayavi bug on mac 32bits
186
187    ret = (coords, faces)
188    if read_metadata:
189        if len(volume_info) == 0:
190            warnings.warn('No volume information contained in the file')
191        ret += (volume_info,)
192    if read_stamp:
193        ret += (create_stamp,)
194
195    return ret
196
197
198def write_geometry(filepath, coords, faces, create_stamp=None,
199                   volume_info=None):
200    """Write a triangular format Freesurfer surface mesh.
201
202    Parameters
203    ----------
204    filepath : str
205        Path to surface file.
206    coords : numpy array
207        nvtx x 3 array of vertex (x, y, z) coordinates.
208    faces : numpy array
209        nfaces x 3 array of defining mesh triangles.
210    create_stamp : str, optional
211        User/time stamp (default: "created by <user> on <ctime>")
212    volume_info : dict-like or None, optional
213        Key-value pairs to encode at the end of the file.
214
215        Valid keys:
216
217        * 'head' : array of int
218        * 'valid' : str
219        * 'filename' : str
220        * 'volume' : array of int, shape (3,)
221        * 'voxelsize' : array of float, shape (3,)
222        * 'xras' : array of float, shape (3,)
223        * 'yras' : array of float, shape (3,)
224        * 'zras' : array of float, shape (3,)
225        * 'cras' : array of float, shape (3,)
226
227    """
228    magic_bytes = np.array([255, 255, 254], dtype=np.uint8)
229
230    if create_stamp is None:
231        create_stamp = f"created by {getpass.getuser()} on {time.ctime()}"
232
233    with open(filepath, 'wb') as fobj:
234        magic_bytes.tofile(fobj)
235        fobj.write((f"{create_stamp}\n\n").encode('utf-8'))
236
237        np.array([coords.shape[0], faces.shape[0]], dtype='>i4').tofile(fobj)
238
239        # Coerce types, just to be safe
240        coords.astype('>f4').reshape(-1).tofile(fobj)
241        faces.astype('>i4').reshape(-1).tofile(fobj)
242
243        # Add volume info, if given
244        if volume_info is not None and len(volume_info) > 0:
245            fobj.write(_serialize_volume_info(volume_info))
246
247
248def read_morph_data(filepath):
249    """Read a Freesurfer morphometry data file.
250
251    This function reads in what Freesurfer internally calls "curv" file types,
252    (e.g. ?h. curv, ?h.thickness), but as that has the potential to cause
253    confusion where "curv" also refers to the surface curvature values,
254    we refer to these files as "morphometry" files with PySurfer.
255
256    Parameters
257    ----------
258    filepath : str
259        Path to morphometry file
260
261    Returns
262    -------
263    curv : numpy array
264        Vector representation of surface morpometry values
265    """
266    with open(filepath, "rb") as fobj:
267        magic = _fread3(fobj)
268        if magic == 16777215:
269            vnum = np.fromfile(fobj, ">i4", 3)[0]
270            curv = np.fromfile(fobj, ">f4", vnum)
271        else:
272            vnum = magic
273            _fread3(fobj)
274            curv = np.fromfile(fobj, ">i2", vnum) / 100
275    return curv
276
277
278def write_morph_data(file_like, values, fnum=0):
279    """Write Freesurfer morphometry data `values` to file-like `file_like`
280
281    Equivalent to FreeSurfer's `write_curv.m`_
282
283    See also:
284    http://www.grahamwideman.com/gw/brain/fs/surfacefileformats.htm#CurvNew
285
286    .. _write_curv.m: \
287    https://github.com/neurodebian/freesurfer/blob/debian-sloppy/matlab/write_curv.m
288
289    Parameters
290    ----------
291    file_like : file-like
292        String containing path of file to be written, or file-like object, open
293        in binary write (`'wb'` mode, implementing the `write` method)
294    values : array-like
295        Surface morphometry values.  Shape must be (N,), (N, 1), (1, N) or (N,
296        1, 1)
297    fnum : int, optional
298        Number of faces in the associated surface.
299    """
300    magic_bytes = np.array([255, 255, 255], dtype=np.uint8)
301
302    vector = np.asarray(values)
303    vnum = np.prod(vector.shape)
304    if vector.shape not in ((vnum,), (vnum, 1), (1, vnum), (vnum, 1, 1)):
305        raise ValueError("Invalid shape: argument values must be a vector")
306
307    i4info = np.iinfo('i4')
308    if vnum > i4info.max:
309        raise ValueError("Too many values for morphometry file")
310    if not i4info.min <= fnum <= i4info.max:
311        raise ValueError(f"Argument fnum must be between {i4info.min} and {i4info.max}")
312
313    with Opener(file_like, 'wb') as fobj:
314        fobj.write(magic_bytes)
315
316        # vertex count, face count (unused), vals per vertex (only 1 supported)
317        fobj.write(np.array([vnum, fnum, 1], dtype='>i4'))
318
319        fobj.write(vector.astype('>f4'))
320
321
322def read_annot(filepath, orig_ids=False):
323    """Read in a Freesurfer annotation from a ``.annot`` file.
324
325    An ``.annot`` file contains a sequence of vertices with a label (also known
326    as an "annotation value") associated with each vertex, and then a sequence
327    of colors corresponding to each label.
328
329    Annotation file format versions 1 and 2 are supported, corresponding to
330    the "old-style" and "new-style" color table layout.
331
332    Note that the output color table ``ctab`` is in RGBT form, where T
333    (transparency) is 255 - alpha.
334
335    See:
336     * https://surfer.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles#Annotation
337     * https://github.com/freesurfer/freesurfer/blob/dev/matlab/read_annotation.m
338     * https://github.com/freesurfer/freesurfer/blob/8b88b34/utils/colortab.c
339
340    Parameters
341    ----------
342    filepath : str
343        Path to annotation file.
344    orig_ids : bool
345        Whether to return the vertex ids as stored in the annotation
346        file or the positional colortable ids. With orig_ids=False
347        vertices with no id have an id set to -1.
348
349    Returns
350    -------
351    labels : ndarray, shape (n_vertices,)
352        Annotation id at each vertex. If a vertex does not belong
353        to any label and orig_ids=False, its id will be set to -1.
354    ctab : ndarray, shape (n_labels, 5)
355        RGBT + label id colortable array.
356    names : list of bytes
357        The names of the labels. The length of the list is n_labels.
358    """
359    with open(filepath, "rb") as fobj:
360        dt = _ANNOT_DT
361
362        # number of vertices
363        vnum = np.fromfile(fobj, dt, 1)[0]
364
365        # vertex ids + annotation values
366        data = np.fromfile(fobj, dt, vnum * 2).reshape(vnum, 2)
367        labels = data[:, 1]
368
369        # is there a color table?
370        ctab_exists = np.fromfile(fobj, dt, 1)[0]
371        if not ctab_exists:
372            raise Exception('Color table not found in annotation file')
373
374        # in old-format files, the next field will contain the number of
375        # entries in the color table. In new-format files, this must be
376        # equal to -2
377        n_entries = np.fromfile(fobj, dt, 1)[0]
378
379        # We've got an old-format .annot file.
380        if n_entries > 0:
381            ctab, names = _read_annot_ctab_old_format(fobj, n_entries)
382        # We've got a new-format .annot file
383        else:
384            ctab, names = _read_annot_ctab_new_format(fobj, -n_entries)
385
386    # generate annotation values for each LUT entry
387    ctab[:, [4]] = _pack_rgb(ctab[:, :3])
388
389    if not orig_ids:
390        ord = np.argsort(ctab[:, -1])
391        mask = labels != 0
392        labels[~mask] = -1
393        labels[mask] = ord[np.searchsorted(ctab[ord, -1], labels[mask])]
394    return labels, ctab, names
395
396
397def _read_annot_ctab_old_format(fobj, n_entries):
398    """Read in an old-style Freesurfer color table from `fobj`.
399
400    Note that the output color table ``ctab`` is in RGBT form, where T
401    (transparency) is 255 - alpha.
402
403    This function is used by :func:`read_annot`.
404
405    Parameters
406    ----------
407
408    fobj : file-like
409        Open file handle to a Freesurfer `.annot` file, with seek point
410        at the beginning of the color table data.
411    n_entries : int
412        Number of entries in the color table.
413
414    Returns
415    -------
416
417    ctab : ndarray, shape (n_entries, 5)
418        RGBT colortable array - the last column contains all zeros.
419    names : list of str
420        The names of the labels. The length of the list is n_entries.
421    """
422    assert hasattr(fobj, 'read')
423
424    dt = _ANNOT_DT
425    # orig_tab string length + string
426    length = np.fromfile(fobj, dt, 1)[0]
427    orig_tab = np.fromfile(fobj, '>c', length)
428    orig_tab = orig_tab[:-1]
429    names = list()
430    ctab = np.zeros((n_entries, 5), dt)
431    for i in range(n_entries):
432        # structure name length + string
433        name_length = np.fromfile(fobj, dt, 1)[0]
434        name = np.fromfile(fobj, "|S%d" % name_length, 1)[0]
435        names.append(name)
436        # read RGBT for this entry
437        ctab[i, :4] = np.fromfile(fobj, dt, 4)
438
439    return ctab, names
440
441
442def _read_annot_ctab_new_format(fobj, ctab_version):
443    """Read in a new-style Freesurfer color table from `fobj`.
444
445    Note that the output color table ``ctab`` is in RGBT form, where T
446    (transparency) is 255 - alpha.
447
448    This function is used by :func:`read_annot`.
449
450    Parameters
451    ----------
452
453    fobj : file-like
454        Open file handle to a Freesurfer `.annot` file, with seek point
455        at the beginning of the color table data.
456    ctab_version : int
457        Color table format version - must be equal to 2
458
459    Returns
460    -------
461
462    ctab : ndarray, shape (n_labels, 5)
463        RGBT colortable array - the last column contains all zeros.
464    names : list of str
465        The names of the labels. The length of the list is n_labels.
466    """
467    assert hasattr(fobj, 'read')
468
469    dt = _ANNOT_DT
470    # This code works with a file version == 2, nothing else
471    if ctab_version != 2:
472        raise Exception('Unrecognised .annot file version (%i)', ctab_version)
473    # maximum LUT index present in the file
474    max_index = np.fromfile(fobj, dt, 1)[0]
475    ctab = np.zeros((max_index, 5), dt)
476    # orig_tab string length + string
477    length = np.fromfile(fobj, dt, 1)[0]
478    np.fromfile(fobj, "|S%d" % length, 1)[0]  # Orig table path
479    # number of LUT entries present in the file
480    entries_to_read = np.fromfile(fobj, dt, 1)[0]
481    names = list()
482    for _ in range(entries_to_read):
483        # index of this entry
484        idx = np.fromfile(fobj, dt, 1)[0]
485        # structure name length + string
486        name_length = np.fromfile(fobj, dt, 1)[0]
487        name = np.fromfile(fobj, "|S%d" % name_length, 1)[0]
488        names.append(name)
489        # RGBT
490        ctab[idx, :4] = np.fromfile(fobj, dt, 4)
491
492    return ctab, names
493
494
495def write_annot(filepath, labels, ctab, names, fill_ctab=True):
496    """Write out a "new-style" Freesurfer annotation file.
497
498    Note that the color table ``ctab`` is in RGBT form, where T (transparency)
499    is 255 - alpha.
500
501    See:
502     * https://surfer.nmr.mgh.harvard.edu/fswiki/LabelsClutsAnnotationFiles#Annotation
503     * https://github.com/freesurfer/freesurfer/blob/dev/matlab/write_annotation.m
504     * https://github.com/freesurfer/freesurfer/blob/8b88b34/utils/colortab.c
505
506    Parameters
507    ----------
508    filepath : str
509        Path to annotation file to be written
510    labels : ndarray, shape (n_vertices,)
511        Annotation id at each vertex.
512    ctab : ndarray, shape (n_labels, 5)
513        RGBT + label id colortable array.
514    names : list of str
515        The names of the labels. The length of the list is n_labels.
516    fill_ctab : {True, False} optional
517        If True, the annotation values for each vertex  are automatically
518        generated. In this case, the provided `ctab` may have shape
519        (n_labels, 4) or (n_labels, 5) - if the latter, the final column is
520        ignored.
521    """
522    with open(filepath, "wb") as fobj:
523        dt = _ANNOT_DT
524        vnum = len(labels)
525
526        def write(num, dtype=dt):
527            np.array([num]).astype(dtype).tofile(fobj)
528
529        def write_string(s):
530            s = (s if isinstance(s, bytes) else s.encode()) + b'\x00'
531            write(len(s))
532            write(s, dtype='|S%d' % len(s))
533
534        # Generate annotation values for each ctab entry
535        if fill_ctab:
536            ctab = np.hstack((ctab[:, :4], _pack_rgb(ctab[:, :3])))
537        elif not np.array_equal(ctab[:, [4]], _pack_rgb(ctab[:, :3])):
538            warnings.warn(f'Annotation values in {filepath} will be incorrect')
539
540        # vtxct
541        write(vnum)
542
543        # convert labels into coded CLUT values
544        clut_labels = ctab[:, -1][labels]
545        clut_labels[np.where(labels == -1)] = 0
546
547        # vno, label
548        data = np.vstack((np.array(range(vnum)),
549                          clut_labels)).T.astype(dt)
550        data.tofile(fobj)
551
552        # tag
553        write(1)
554
555        # ctabversion
556        write(-2)
557
558        # maxstruc
559        write(max(np.max(labels) + 1, ctab.shape[0]))
560
561        # File of LUT is unknown.
562        write_string('NOFILE')
563
564        # num_entries
565        write(ctab.shape[0])
566
567        for ind, (clu, name) in enumerate(zip(ctab, names)):
568            write(ind)
569            write_string(name)
570            for val in clu[:-1]:
571                write(val)
572
573
574def read_label(filepath, read_scalars=False):
575    """Load in a Freesurfer .label file.
576
577    Parameters
578    ----------
579    filepath : str
580        Path to label file.
581    read_scalars : bool, optional
582        If True, read and return scalars associated with each vertex.
583
584    Returns
585    -------
586    label_array : numpy array
587        Array with indices of vertices included in label.
588    scalar_array : numpy array (floats)
589        Only returned if `read_scalars` is True.  Array of scalar data for each
590        vertex.
591    """
592    label_array = np.loadtxt(filepath, dtype=int, skiprows=2, usecols=[0])
593    if read_scalars:
594        scalar_array = np.loadtxt(filepath, skiprows=2, usecols=[-1])
595        return label_array, scalar_array
596    return label_array
597
598
599def _serialize_volume_info(volume_info):
600    """Helper for serializing the volume info."""
601    keys = ['head', 'valid', 'filename', 'volume', 'voxelsize', 'xras', 'yras',
602            'zras', 'cras']
603    diff = set(volume_info.keys()).difference(keys)
604    if len(diff) > 0:
605        raise ValueError(f'Invalid volume info: {diff.pop()}.')
606
607    strings = list()
608    for key in keys:
609        if key == 'head':
610            if not (np.array_equal(volume_info[key], [20]) or np.array_equal(
611                    volume_info[key], [2, 0, 20])):
612                warnings.warn("Unknown extension code.")
613            strings.append(np.array(volume_info[key], dtype='>i4').tobytes())
614        elif key in ('valid', 'filename'):
615            val = volume_info[key]
616            strings.append(f'{key} = {val}\n'.encode('utf-8'))
617        elif key == 'volume':
618            val = volume_info[key]
619            strings.append(f'{key} = {val[0]} {val[1]} {val[2]}\n'.encode('utf-8'))
620        else:
621            val = volume_info[key]
622            strings.append(
623                f'{key:6s} = {val[0]:.10g} {val[1]:.10g} {val[2]:.10g}\n'.encode('utf-8'))
624    return b''.join(strings)
625