Merge pull request #17 from mgxd/enh/lta

ENH: FreeSurfer LTA file support

Author: oesteban

nitransforms/io.py

@@ -0,0 +1,265 @@
+import numpy as np
+from nibabel.volumeutils import Recoder
+from nibabel.affines import voxel_sizes
+
+from .patched import LabeledWrapStruct
+
+transform_codes = Recoder((
+    (0, 'LINEAR_VOX_TO_VOX'),
+    (1, 'LINEAR_RAS_TO_RAS'),
+    (2, 'LINEAR_PHYSVOX_TO_PHYSVOX'),
+    (14, 'REGISTER_DAT'),
+    (21, 'LINEAR_COR_TO_COR')),
+    fields=('code', 'label'))
+
+
+class StringBasedStruct(LabeledWrapStruct):
+    def __init__(self,
+                 binaryblock=None,
+                 endianness=None,
+                 check=True):
+        if binaryblock is not None and getattr(binaryblock, 'dtype',
+                                               None) == self.dtype:
+            self._structarr = binaryblock.copy()
+            return
+        super(StringBasedStruct, self).__init__(binaryblock, endianness, check)
+
+    def __array__(self):
+        return self._structarr
+
+
+class VolumeGeometry(StringBasedStruct):
+    template_dtype = np.dtype([
+        ('valid', 'i4'),              # Valid values: 0, 1
+        ('volume', 'i4', (3, 1)),     # width, height, depth
+        ('voxelsize', 'f4', (3, 1)),  # xsize, ysize, zsize
+        ('xras', 'f4', (3, 1)),       # x_r, x_a, x_s
+        ('yras', 'f4', (3, 1)),       # y_r, y_a, y_s
+        ('zras', 'f4', (3, 1)),       # z_r, z_a, z_s
+        ('cras', 'f4', (3, 1)),       # c_r, c_a, c_s
+        ('filename', 'U1024')])       # Not conformant (may be >1024 bytes)
+    dtype = template_dtype
+
+    def as_affine(self):
+        affine = np.eye(4)
+        sa = self.structarr
+        A = np.hstack((sa['xras'], sa['yras'], sa['zras'])) * sa['voxelsize']
+        b = sa['cras'] - A.dot(sa['volume']) / 2
+        affine[:3, :3] = A
+        affine[:3, [3]] = b
+        return affine
+
+    def to_string(self):
+        sa = self.structarr
+        lines = [
+            'valid = {}  # volume info {:s}valid'.format(
+                sa['valid'], '' if sa['valid'] else 'in'),
+            'filename = {}'.format(sa['filename']),
+            'volume = {:d} {:d} {:d}'.format(*sa['volume'].flatten()),
+            'voxelsize = {:.15e} {:.15e} {:.15e}'.format(
+                *sa['voxelsize'].flatten()),
+            'xras   = {:.15e} {:.15e} {:.15e}'.format(*sa['xras'].flatten()),
+            'yras   = {:.15e} {:.15e} {:.15e}'.format(*sa['yras'].flatten()),
+            'zras   = {:.15e} {:.15e} {:.15e}'.format(*sa['zras'].flatten()),
+            'cras   = {:.15e} {:.15e} {:.15e}'.format(*sa['cras'].flatten()),
+        ]
+        return '\n'.join(lines)
+
+    @classmethod
+    def from_image(klass, img):
+        volgeom = klass()
+        sa = volgeom.structarr
+        sa['valid'] = 1
+        sa['volume'][:, 0] = img.shape[:3]    # Assumes xyzt-ordered image
+        sa['voxelsize'][:, 0] = voxel_sizes(img.affine)[:3]
+        A = img.affine[:3, :3]
+        b = img.affine[:3, [3]]
+        cols = A / sa['voxelsize']
+        sa['xras'] = cols[:, [0]]
+        sa['yras'] = cols[:, [1]]
+        sa['zras'] = cols[:, [2]]
+        sa['cras'] = b + A.dot(sa['volume']) / 2
+        try:
+            sa['filename'] = img.file_map['image'].filename
+        except Exception:
+            pass
+
+        return volgeom
+
+    @classmethod
+    def from_string(klass, string):
+        volgeom = klass()
+        sa = volgeom.structarr
+        lines = string.splitlines()
+        for key in ('valid', 'filename', 'volume', 'voxelsize',
+                    'xras', 'yras', 'zras', 'cras'):
+            label, valstring = lines.pop(0).split(' = ')
+            assert label.strip() == key
+
+            val = np.genfromtxt([valstring.encode()],
+                                dtype=klass.dtype[key])
+            sa[key] = val.reshape(sa[key].shape) if val.size else ''
+
+        return volgeom
+
+
+class LinearTransform(StringBasedStruct):
+    template_dtype = np.dtype([
+        ('mean', 'f4', (3, 1)),       # x0, y0, z0
+        ('sigma', 'f4'),
+        ('m_L', 'f4', (4, 4)),
+        ('m_dL', 'f4', (4, 4)),
+        ('m_last_dL', 'f4', (4, 4)),
+        ('src', VolumeGeometry),
+        ('dst', VolumeGeometry),
+        ('label', 'i4')])
+    dtype = template_dtype
+
+    def __getitem__(self, idx):
+        val = super(LinearTransform, self).__getitem__(idx)
+        if idx in ('src', 'dst'):
+            val = VolumeGeometry(val)
+        return val
+
+    def to_string(self):
+        sa = self.structarr
+        lines = [
+            'mean      = {:6.4f} {:6.4f} {:6.4f}'.format(
+                *sa['mean'].flatten()),
+            'sigma     = {:6.4f}'.format(float(sa['sigma'])),
+            '1 4 4',
+            ('{:18.15e} ' * 4).format(*sa['m_L'][0]),
+            ('{:18.15e} ' * 4).format(*sa['m_L'][1]),
+            ('{:18.15e} ' * 4).format(*sa['m_L'][2]),
+            ('{:18.15e} ' * 4).format(*sa['m_L'][3]),
+            'src volume info',
+            self['src'].to_string(),
+            'dst volume info',
+            self['dst'].to_string(),
+        ]
+        return '\n'.join(lines)
+
+    @classmethod
+    def from_string(klass, string):
+        lt = klass()
+        sa = lt.structarr
+        lines = string.splitlines()
+        for key in ('mean', 'sigma'):
+            label, valstring = lines.pop(0).split(' = ')
+            assert label.strip() == key
+
+            val = np.genfromtxt([valstring.encode()],
+                                dtype=klass.dtype[key])
+            sa[key] = val.reshape(sa[key].shape)
+        assert lines.pop(0) == '1 4 4'  # xforms, shape + 1, shape + 1
+        val = np.genfromtxt([valstring.encode() for valstring in lines[:4]],
+                            dtype='f4')
+        sa['m_L'] = val
+        lines = lines[4:]
+        assert lines.pop(0) == 'src volume info'
+        sa['src'] = np.asanyarray(VolumeGeometry.from_string('\n'.join(lines[:8])))
+        lines = lines[8:]
+        assert lines.pop(0) == 'dst volume info'
+        sa['dst'] = np.asanyarray(VolumeGeometry.from_string('\n'.join(lines)))
+        return lt
+
+
+class LinearTransformArray(StringBasedStruct):
+    template_dtype = np.dtype([
+        ('type', 'i4'),
+        ('nxforms', 'i4'),
+        ('subject', 'U1024'),
+        ('fscale', 'f4')])
+    dtype = template_dtype
+    _xforms = None
+
+    def __init__(self,
+                 binaryblock=None,
+                 endianness=None,
+                 check=True):
+        super(LinearTransformArray, self).__init__(binaryblock, endianness, check)
+        self._xforms = [LinearTransform()
+                        for _ in range(self.structarr['nxforms'])]
+
+    def __getitem__(self, idx):
+        if idx == 'xforms':
+            return self._xforms
+        if idx == 'nxforms':
+            return len(self._xforms)
+        return super(LinearTransformArray, self).__getitem__(idx)
+
+    def to_string(self):
+        code = int(self['type'])
+        header = [
+            'type      = {} # {}'.format(code, transform_codes.label[code]),
+            'nxforms   = {}'.format(self['nxforms'])]
+        xforms = [xfm.to_string() for xfm in self._xforms]
+        footer = [
+            'subject {}'.format(self['subject']),
+            'fscale {:.6f}'.format(float(self['fscale']))]
+        return '\n'.join(header + xforms + footer)
+
+    @classmethod
+    def from_string(klass, string):
+        lta = klass()
+        sa = lta.structarr
+        lines = string.splitlines()
+        for key in ('type', 'nxforms'):
+            label, valstring = lines.pop(0).split(' = ')
+            assert label.strip() == key
+
+            val = np.genfromtxt([valstring.encode()],
+                                dtype=klass.dtype[key])
+            sa[key] = val.reshape(sa[key].shape) if val.size else ''
+        for _ in range(sa['nxforms']):
+            lta._xforms.append(
+                LinearTransform.from_string('\n'.join(lines[:25])))
+            lines = lines[25:]
+        if lines:
+            for key in ('subject', 'fscale'):
+                # Optional keys
+                if not lines[0].startswith(key):
+                    continue
+                label, valstring = lines.pop(0).split(' ')
+                assert label.strip() == key
+
+                val = np.genfromtxt([valstring.encode()],
+                                    dtype=klass.dtype[key])
+                sa[key] = val.reshape(sa[key].shape) if val.size else ''
+
+        assert len(lta._xforms) == sa['nxforms']
+        return lta
+
+    @classmethod
+    def from_fileobj(klass, fileobj, check=True):
+        return klass.from_string(fileobj.read())
+
+    def set_type(self, target):
+        """
+        Convert the internal transformation matrix to a different type inplace
+
+        Parameters
+        ----------
+        target : str, int
+            Tranformation type
+        """
+        assert self['nxforms'] == 1, "Cannot convert multiple transformations"
+        xform = self['xforms'][0]
+        src = xform['src']
+        dst = xform['dst']
+        current = self['type']
+        if isinstance(target, str):
+            target = transform_codes.code[target]
+
+        # VOX2VOX -> RAS2RAS
+        if current == 0 and target == 1:
+            M = np.linalg.inv(src.as_affine()).dot(xform['m_L']).dot(dst.as_affine())
+        else:
+            raise NotImplementedError(
+                "Converting {0} to {1} is not yet available".format(
+                    transform_codes.label[current],
+                    transform_codes.label[target]
+                )
+            )
+        xform['m_L'] = M
+        self['type'] = target

nitransforms/linear.py

@@ -11,11 +11,13 @@
 import numpy as np
 from scipy import ndimage as ndi
 from pathlib import Path
+import warnings
 
 from nibabel.loadsave import load as loadimg
 from nibabel.affines import from_matvec, voxel_sizes, obliquity
 from .base import TransformBase
 from .patched import shape_zoom_affine
+from . import io
 
 
 LPS = np.diag([-1, -1, 1, 1])
@@ -126,8 +128,7 @@ def resample(self, moving, order=3, mode='constant', cval=0.0, prefilter=True,
         try:
             reference = self.reference
         except ValueError:
-            print('Warning: no reference space defined, using moving as reference',
-                  file=sys.stderr)
+            warnings.warn('No reference space defined, using moving as reference')
             reference = moving
 
         nvols = 1
@@ -150,8 +151,7 @@ def resample(self, moving, order=3, mode='constant', cval=0.0, prefilter=True,
             singlemat = np.linalg.inv(movaff).dot(self._matrix[0].dot(reference.affine))
 
         if singlemat is not None and nvols > nmats:
-            print('Warning: resampling a 4D volume with a single affine matrix',
-                  file=sys.stderr)
+            warnings.warn('Resampling a 4D volume with a single affine matrix')
 
         # Compose an index to index affine matrix
         moved = []
@@ -270,13 +270,13 @@ def to_filename(self, filename, fmt='X5', moving=None):
 3dvolreg matrices (DICOM-to-DICOM, row-by-row):""", fmt='%g')
             return filename
 
-        if fmt.lower() == 'fsl':
-            if not moving:
-                moving = self.reference
-
-            if isinstance(moving, str):
-                moving = loadimg(moving)
+        # for FSL / FS information
+        if not moving:
+            moving = self.reference
+        if isinstance(moving, str):
+            moving = loadimg(moving)
 
+        if fmt.lower() == 'fsl':
             # Adjust for reference image offset and orientation
             refswp, refspc = _fsl_aff_adapt(self.reference)
             pre = self.reference.affine.dot(
@@ -298,6 +298,22 @@ def to_filename(self, filename, fmt='X5', moving=None):
             else:
                 np.savetxt(filename, mat[0], delimiter=' ', fmt='%g')
             return filename
+        elif fmt.lower() == 'fs':
+            # xform info
+            lt = io.LinearTransform()
+            lt['sigma'] = 1.
+            lt['m_L'] = self.matrix
+            lt['src'] = io.VolumeGeometry.from_image(moving)
+            lt['dst'] = io.VolumeGeometry.from_image(self.reference)
+            # to make LTA file format
+            lta = io.LinearTransformArray()
+            lta['type'] = 1  # RAS2RAS
+            lta['xforms'].append(lt)
+
+            with open(filename, 'w') as f:
+                f.write(lta.to_string())
+            return filename
+
         return super(Affine, self).to_filename(filename, fmt=fmt)
 
 
@@ -326,6 +342,14 @@ def load(filename, fmt='X5', reference=None):
     # elif fmt.lower() == 'afni':
     #     parameters = LPS.dot(self.matrix.dot(LPS))
     #     parameters = parameters[:3, :].reshape(-1).tolist()
+    elif fmt.lower() == 'fs':
+        with open(filename) as ltafile:
+            lta = io.LinearTransformArray.from_fileobj(ltafile)
+        if lta['nxforms'] > 1:
+            raise NotImplementedError("Multiple transforms are not yet supported.")
+        if lta['type'] != 1:
+            lta.set_type(1)
+        matrix = lta['xforms'][0]['m_L']
     elif fmt.lower() in ('x5', 'bids'):
         raise NotImplementedError
     else:

nitransforms/patched.py

@@ -1,4 +1,5 @@
 import numpy as np
+from nibabel.wrapstruct import LabeledWrapStruct as LWS
 
 
 def shape_zoom_affine(shape, zooms, x_flip=True, y_flip=False):
@@ -63,3 +64,8 @@ def shape_zoom_affine(shape, zooms, x_flip=True, y_flip=False):
     aff[:3, :3] = np.diag(zooms)
     aff[:3, -1] = -origin * zooms
     return aff
+
+
+class LabeledWrapStruct(LWS):
+    def __setitem__(self, item, value):
+        self._structarr[item] = np.asanyarray(value)