ENH: More comprehensive implementation of ITK affines I/O

Author: oesteban

nitransforms/io/base.py

@@ -1,23 +1,40 @@
 """Read/write linear transforms."""
-import numpy as np
-from nibabel.wrapstruct import LabeledWrapStruct as LWS
+from scipy.io.matlab.miobase import get_matfile_version
+from scipy.io.matlab.mio4 import MatFile4Reader  # , MatFile4Writer
+from scipy.io.matlab.mio5 import MatFile5Reader  # , MatFile5Writer
 
+from ..patched import LabeledWrapStruct
 
-class LabeledWrapStruct(LWS):
-    def __setitem__(self, item, value):
-        self._structarr[item] = np.asanyarray(value)
+
+class TransformFileError(Exception):
+    """A custom exception for transform files."""
 
 
 class StringBasedStruct(LabeledWrapStruct):
+    """File data structure from text files."""
+
     def __init__(self,
                  binaryblock=None,
                  endianness=None,
                  check=True):
-        if binaryblock is not None and getattr(binaryblock, 'dtype',
-                                               None) == self.dtype:
+        """Create a data structure based off of a string."""
+        _dtype = getattr(binaryblock, 'dtype', None)
+        if binaryblock is not None and _dtype == self.dtype:
             self._structarr = binaryblock.copy()
             return
         super(StringBasedStruct, self).__init__(binaryblock, endianness, check)
 
     def __array__(self):
+        """Return the internal structure array."""
         return self._structarr
+
+
+def _read_mat(byte_stream):
+    mjv, _ = get_matfile_version(byte_stream)
+    if mjv == 0:
+        reader = MatFile4Reader(byte_stream)
+    elif mjv == 1:
+        reader = MatFile5Reader(byte_stream)
+    elif mjv == 2:
+        raise TransformFileError('Please use HDF reader for matlab v7.3 files')
+    return reader.get_variables()

nitransforms/io/itk.py

@@ -1,6 +1,10 @@
 """Read/write ITK transforms."""
 import numpy as np
-from .base import StringBasedStruct
+from scipy.io import savemat as _save_mat
+from nibabel.affines import from_matvec
+from .base import StringBasedStruct, _read_mat
+
+LPS = np.diag([-1, -1, 1, 1])
 
 
 class ITKLinearTransform(StringBasedStruct):
@@ -13,20 +17,24 @@ class ITKLinearTransform(StringBasedStruct):
         ('offset', 'f4', 3),  # Center of rotation
     ])
     dtype = template_dtype
+    # files_types = (('string', '.tfm'), ('binary', '.mat'))
+    # valid_exts = ('.tfm', '.mat')
 
-    def __init__(self):
+    def __init__(self, parameters=None, offset=None):
         """Initialize with default offset and index."""
         super().__init__()
-        self.structarr['offset'] = [0, 0, 0]
         self.structarr['index'] = 1
+        self.structarr['offset'] = offset or [0, 0, 0]
         self.structarr['parameters'] = np.eye(4)
+        if parameters is not None:
+            self.structarr['parameters'] = parameters
 
     def __str__(self):
         """Generate a string representation."""
         sa = self.structarr
         lines = [
             '#Transform {:d}'.format(sa['index']),
-            'Transform: MatrixOffsetTransformBase_double_3_3',
+            'Transform: AffineTransform_float_3_3',
             'Parameters: {}'.format(' '.join(
                 ['%g' % p
                  for p in sa['parameters'][:3, :3].reshape(-1).tolist() +
@@ -36,6 +44,33 @@ def __str__(self):
         ]
         return '\n'.join(lines)
 
+    def to_filename(self, filename):
+        """Store this transform to a file with the appropriate format."""
+        if str(filename).endswith('.mat'):
+            sa = self.structarr
+            affine = np.array(np.hstack((
+                sa['parameters'][:3, :3].reshape(-1),
+                sa['parameters'][:3, 3]))[..., np.newaxis], dtype='f4')
+            fixed = np.array(sa['offset'][..., np.newaxis], dtype='f4')
+            mdict = {
+                'AffineTransform_float_3_3': affine,
+                'fixed': fixed,
+            }
+            _save_mat(str(filename), mdict, format='4')
+            return
+
+        with open(str(filename), 'w') as f:
+            f.write(self.to_string())
+
+    def to_ras(self):
+        """Return a nitransforms' internal RAS matrix."""
+        sa = self.structarr
+        matrix = sa['parameters']
+        offset = sa['offset']
+        c_neg = from_matvec(np.eye(3), offset * -1.0)
+        c_pos = from_matvec(np.eye(3), offset)
+        return LPS.dot(c_pos.dot(matrix.dot(c_neg.dot(LPS))))
+
     def to_string(self, banner=None):
         """Convert to a string directly writeable to file."""
         string = '%s'
@@ -48,9 +83,47 @@ def to_string(self, banner=None):
         return string % self
 
     @classmethod
-    def from_string(klass, string):
+    def from_binary(cls, byte_stream, index=None):
+        """Read the struct from a matlab binary file."""
+        mdict = _read_mat(byte_stream)
+        return cls.from_matlab_dict(mdict, index=index)
+
+    @classmethod
+    def from_fileobj(cls, fileobj, check=True):
+        """Read the struct from a file object."""
+        if fileobj.name.endswith('.mat'):
+            return cls.from_binary(fileobj)
+        return cls.from_string(fileobj.read())
+
+    @classmethod
+    def from_matlab_dict(cls, mdict, index=None):
+        """Read the struct from a matlab dictionary."""
+        tf = cls()
+        sa = tf.structarr
+        if index is not None:
+            raise NotImplementedError
+
+        sa['index'] = 1
+        parameters = np.eye(4, dtype='f4')
+        parameters[:3, :3] = mdict['AffineTransform_float_3_3'][:-3].reshape((3, 3))
+        parameters[:3, 3] = mdict['AffineTransform_float_3_3'][-3:].flatten()
+        sa['parameters'] = parameters
+        sa['offset'] = mdict['fixed'].flatten()
+        return tf
+
+    @classmethod
+    def from_ras(cls, ras, index=0):
+        """Create an ITK affine from a nitransform's RAS+ matrix."""
+        tf = cls()
+        sa = tf.structarr
+        sa['index'] = index + 1
+        sa['parameters'] = LPS.dot(ras.dot(LPS))
+        return tf
+
+    @classmethod
+    def from_string(cls, string):
         """Read the struct from string."""
-        tf = klass()
+        tf = cls()
         sa = tf.structarr
         lines = [l for l in string.splitlines()
                  if l.strip()]
@@ -61,19 +134,14 @@ def from_string(klass, string):
         parameters = np.eye(4, dtype='f4')
         sa['index'] = int(lines[0][lines[0].index('T'):].split()[1])
         sa['offset'] = np.genfromtxt([lines[3].split(':')[-1].encode()],
-                                     dtype=klass.dtype['offset'])
+                                     dtype=cls.dtype['offset'])
         vals = np.genfromtxt([lines[2].split(':')[-1].encode()],
                              dtype='f4')
         parameters[:3, :3] = vals[:-3].reshape((3, 3))
         parameters[:3, 3] = vals[-3:]
         sa['parameters'] = parameters
         return tf
 
-    @classmethod
-    def from_fileobj(klass, fileobj, check=True):
-        """Read the struct from a file object."""
-        return klass.from_string(fileobj.read())
-
 
 class ITKLinearTransformArray(StringBasedStruct):
     """A string-based structure for series of ITK linear transforms."""
@@ -89,33 +157,80 @@ def __init__(self,
                  check=True):
         """Initialize with (optionally) a list of transforms."""
         super().__init__(binaryblock, endianness, check)
-        self._xforms = []
-        for i, mat in enumerate(xforms or []):
-            xfm = ITKLinearTransform()
-            xfm['parameters'] = mat
-            xfm['index'] = i + 1
-            self._xforms.append(xfm)
+        self.xforms = [ITKLinearTransform(parameters=mat)
+                       for mat in xforms or []]
+
+    @property
+    def xforms(self):
+        """Get the list of internal ITKLinearTransforms."""
+        return self._xforms
+
+    @xforms.setter
+    def xforms(self, value):
+        self._xforms = value
+
+        # Update indexes
+        for i, val in enumerate(self._xforms):
+            val['index'] = i + 1
 
     def __getitem__(self, idx):
         """Allow dictionary access to the transforms."""
         if idx == 'xforms':
             return self._xforms
         if idx == 'nxforms':
             return len(self._xforms)
-        return super().__getitem__(idx)
+        raise KeyError(idx)
+
+    def to_filename(self, filename):
+        """Store this transform to a file with the appropriate format."""
+        if str(filename).endswith('.mat'):
+            raise NotImplementedError
+
+        with open(str(filename), 'w') as f:
+            f.write(self.to_string())
+
+    def to_ras(self):
+        """Return a nitransforms' internal RAS matrix."""
+        return np.stack([xfm.to_ras() for xfm in self._xforms])
 
     def to_string(self):
         """Convert to a string directly writeable to file."""
         strings = []
-        for i, xfm in enumerate(self._xforms):
+        for i, xfm in enumerate(self.xforms):
             xfm.structarr['index'] = i + 1
             strings.append(xfm.to_string())
         return '\n'.join(strings)
 
     @classmethod
-    def from_string(klass, string):
+    def from_binary(cls, byte_stream):
+        """Read the struct from a matlab binary file."""
+        mdict = _read_mat(byte_stream)
+        nxforms = mdict['fixed'].shape[0]
+
+        _self = cls()
+        _self.xforms = [ITKLinearTransform.from_matlab_dict(mdict, i)
+                        for i in range(nxforms)]
+        return _self
+
+    @classmethod
+    def from_fileobj(cls, fileobj, check=True):
+        """Read the struct from a file object."""
+        if fileobj.name.endswith('.mat'):
+            return cls.from_binary(fileobj)
+        return cls.from_string(fileobj.read())
+
+    @classmethod
+    def from_ras(cls, ras):
+        """Create an ITK affine from a nitransform's RAS+ matrix."""
+        _self = cls()
+        _self.xforms = [ITKLinearTransform.from_ras(ras[i, ...], i)
+                        for i in range(ras.shape[0])]
+        return _self
+
+    @classmethod
+    def from_string(cls, string):
         """Read the struct from string."""
-        _self = klass()
+        _self = cls()
         lines = [l.strip() for l in string.splitlines()
                  if l.strip()]
 
@@ -124,11 +239,6 @@ def from_string(klass, string):
 
         string = '\n'.join(lines[1:])
         for xfm in string.split('#')[1:]:
-            _self._xforms.append(ITKLinearTransform.from_string(
+            _self.xforms.append(ITKLinearTransform.from_string(
                 '#%s' % xfm))
         return _self
-
-    @classmethod
-    def from_fileobj(klass, fileobj, check=True):
-        """Read the struct from a file object."""
-        return klass.from_string(fileobj.read())

nitransforms/linear.py

@@ -12,7 +12,7 @@
 import numpy as np
 
 from nibabel.loadsave import load as loadimg
-from nibabel.affines import from_matvec, voxel_sizes, obliquity
+from nibabel.affines import voxel_sizes, obliquity
 from .base import TransformBase, _as_homogeneous, EQUALITY_TOL
 from .patched import shape_zoom_affine
 from . import io
@@ -140,10 +140,8 @@ def _to_hdf5(self, x5_root):
     def to_filename(self, filename, fmt='X5', moving=None):
         """Store the transform in BIDS-Transforms HDF5 file format (.x5)."""
         if fmt.lower() in ['itk', 'ants', 'elastix']:
-            itkobj = io.itk.ITKLinearTransformArray(
-                xforms=[LPS.dot(m.dot(LPS)) for m in self.matrix])
-            with open(filename, 'w') as f:
-                f.write(itkobj.to_string())
+            itkobj = io.itk.ITKLinearTransformArray.from_ras(self.matrix)
+            itkobj.to_filename(filename)
             return filename
 
         if fmt.lower() == 'afni':
@@ -235,19 +233,11 @@ def to_filename(self, filename, fmt='X5', moving=None):
 
 def load(filename, fmt='X5', reference=None):
     """Load a linear transform."""
-    if fmt.lower() in ['itk', 'ants', 'elastix', 'nifty']:
+    if fmt.lower() in ('itk', 'ants', 'elastix'):
         with open(filename) as itkfile:
             itkxfm = io.itk.ITKLinearTransformArray.from_fileobj(
                 itkfile)
-
-        matlist = []
-        for xfm in itkxfm['xforms']:
-            matrix = xfm['parameters']
-            offset = xfm['offset']
-            c_neg = from_matvec(np.eye(3), offset * -1.0)
-            c_pos = from_matvec(np.eye(3), offset)
-            matlist.append(LPS.dot(c_pos.dot(matrix.dot(c_neg.dot(LPS)))))
-        matrix = np.stack(matlist)
+        matrix = itkxfm.to_ras()
     # elif fmt.lower() == 'afni':
     #     parameters = LPS.dot(self.matrix.dot(LPS))
     #     parameters = parameters[:3, :].reshape(-1).tolist()

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)

nitransforms/tests/data/itktflist.tfm

@@ -1,45 +1,45 @@
 #Insight Transform File V1.0
 #Transform 1
-Transform: MatrixOffsetTransformBase_double_3_3
+Transform: AffineTransform_float_3_3
 Parameters: 1 0 0 0 1 0 0 0 1 0 0 0
 FixedParameters: 10 10 10
 
 #Transform 2
-Transform: MatrixOffsetTransformBase_double_3_3
+Transform: AffineTransform_float_3_3
 Parameters: 1 0 0 0 1 0 0 0 1 0 0 0
 FixedParameters: 10 10 10
 
 #Transform 3
-Transform: MatrixOffsetTransformBase_double_3_3
+Transform: AffineTransform_float_3_3
 Parameters: 1 0 0 0 1 0 0 0 1 0 0 0
 FixedParameters: 10 10 10
 
 #Transform 4
-Transform: MatrixOffsetTransformBase_double_3_3
+Transform: AffineTransform_float_3_3
 Parameters: 1 0 0 0 1 0 0 0 1 0 0 0
 FixedParameters: 10 10 10
 
 #Transform 5
-Transform: MatrixOffsetTransformBase_double_3_3
+Transform: AffineTransform_float_3_3
 Parameters: -1.53626 0.71973 -0.639856 -0.190759 -1.80082 -0.915885 0.502537 1.12532 0.275748 0.393413 1.13855 0.761131
 FixedParameters: -0.0993171 0.364984 1.99264
 
 #Transform 6
-Transform: MatrixOffsetTransformBase_double_3_3
+Transform: AffineTransform_float_3_3
 Parameters: -0.130507 -1.03017 2.08189 -1.51723 1.37849 -0.0890962 -0.656323 0.242694 2.15801 -1.26689 0.367131 1.23616
 FixedParameters: 0.626607 0.15351 1.24982
 
 #Transform 7
-Transform: MatrixOffsetTransformBase_double_3_3
+Transform: AffineTransform_float_3_3
 Parameters: -1.55395 -0.36383 -0.17749 1.3387 -0.384534 -0.901462 -1.06598 -0.448228 -1.07535 1.92599 0.454696 0.576697
 FixedParameters: -0.425602 0.333406 -1.14957
 
 #Transform 8
-Transform: MatrixOffsetTransformBase_double_3_3
+Transform: AffineTransform_float_3_3
 Parameters: 0.723719 -1.05617 -0.800562 -2.47048 -1.76301 -1.4447 -0.749896 1.29774 -1.48893 1.02789 0.65017 -1.48326
 FixedParameters: 0.800882 -1.20202 1.25495
 
 #Transform 9
-Transform: MatrixOffsetTransformBase_double_3_3
+Transform: AffineTransform_float_3_3
 Parameters: 1.24025 -0.77628 0.618013 -0.523829 1.09471 1.66921 0.73753 -1.33588 -0.627659 -0.449913 -0.00124181 0.21433
 FixedParameters: -0.226504 -0.877893 0.2608

nitransforms/tests/test_io.py

@@ -2,13 +2,22 @@
 # vi: set ft=python sts=4 ts=4 sw=4 et:
 """I/O test cases."""
 import numpy as np
+import pytest
 
+from nibabel.eulerangles import euler2mat
+from nibabel.affines import from_matvec
+from scipy.io import loadmat, savemat
 from ..io import (
     itk,
     VolumeGeometry as VG,
     LinearTransform as LT,
     LinearTransformArray as LTA,
 )
+from ..io.base import _read_mat, TransformFileError
+
+LPS = np.diag([-1, -1, 1, 1])
+ITK_MAT = LPS.dot(np.ones((4, 4)).dot(LPS))
+
 
 def test_VolumeGeometry(tmpdir, get_testdata):
     vg = VG()
@@ -23,7 +32,7 @@ def test_VolumeGeometry(tmpdir, get_testdata):
     assert len(vg.to_string().split('\n')) == 8
 
 
-def test_LinearTransform(tmpdir, get_testdata):
+def test_LinearTransform(tmpdir):
     lt = LT()
     assert lt['m_L'].shape == (4, 4)
     assert np.all(lt['m_L'] == 0)
@@ -57,6 +66,27 @@ def test_LinearTransformArray(tmpdir, data_path):
     assert np.allclose(lta['xforms'][0]['m_L'], lta2['xforms'][0]['m_L'])
 
 
+def test_ITKLinearTransform(tmpdir, data_path):
+    tmpdir.chdir()
+
+    matlabfile = str(data_path / 'ds-005_sub-01_from-T1_to-OASIS_affine.mat')
+    mat = loadmat(matlabfile)
+    with open(matlabfile, 'rb') as f:
+        itkxfm = itk.ITKLinearTransform.from_fileobj(f)
+    assert np.allclose(itkxfm['parameters'][:3, :3].flatten(),
+                       mat['AffineTransform_float_3_3'][:-3].flatten())
+    assert np.allclose(itkxfm['offset'], mat['fixed'].reshape((3, )))
+    itkxfm.to_filename('copy.mat')
+
+    with open('copy.mat', 'rb') as f:
+        itkxfm2 = itk.ITKLinearTransform.from_fileobj(f)
+    assert np.all(itkxfm['parameters'] == itkxfm2['parameters'])
+
+    rasmat = from_matvec(euler2mat(x=0.9, y=0.001, z=0.001), [4.0, 2.0, -1.0])
+    itkxfm = itk.ITKLinearTransform.from_ras(rasmat)
+    assert np.allclose(itkxfm['parameters'], ITK_MAT * rasmat)
+
+
 def test_ITKLinearTransformArray(tmpdir, data_path):
     tmpdir.chdir()
 
@@ -67,14 +97,20 @@ def test_ITKLinearTransformArray(tmpdir, data_path):
 
     assert itklist['nxforms'] == 9
     assert text == itklist.to_string()
+    with pytest.raises(ValueError):
+        itk.ITKLinearTransformArray.from_string(
+            '\n'.join(text.splitlines()[1:]))
 
     itklist = itk.ITKLinearTransformArray(
-        xforms=[np.around(np.random.normal(size=(4, 4)), decimals=5)
+        xforms=[np.random.normal(size=(4, 4))
                 for _ in range(4)])
 
     assert itklist['nxforms'] == 4
     assert itklist['xforms'][1].structarr['index'] == 2
 
+    with pytest.raises(KeyError):
+        itklist['invalid_key']
+
     xfm = itklist['xforms'][1]
     xfm['index'] = 1
     with open('extracted.tfm', 'w') as f:
@@ -84,3 +120,33 @@ def test_ITKLinearTransformArray(tmpdir, data_path):
         xfm2 = itk.ITKLinearTransform.from_fileobj(f)
     assert np.allclose(xfm.structarr['parameters'][:3, ...],
                        xfm2.structarr['parameters'][:3, ...])
+
+
+@pytest.mark.parametrize('matlab_ver', ['4', '5'])
+def test_read_mat1(tmpdir, matlab_ver):
+    """Test read from matlab."""
+    tmpdir.chdir()
+
+    savemat('val.mat', {'val': np.ones((3,))},
+            format=matlab_ver)
+    with open('val.mat', 'rb') as f:
+        mdict = _read_mat(f)
+
+    assert np.all(mdict['val'] == np.ones((3,)))
+
+
+def test_read_mat2(tmpdir, monkeypatch):
+    """Check read matlab raises adequate errors."""
+    from ..io import base
+
+    def _mockreturn(arg):
+        return (2, 0)
+
+    tmpdir.chdir()
+    savemat('val.mat', {'val': np.ones((3,))})
+
+    with monkeypatch.context() as m:
+        m.setattr(base, 'get_matfile_version', _mockreturn)
+        with pytest.raises(TransformFileError):
+            with open('val.mat', 'rb') as f:
+                _read_mat(f)