Merge pull request #35 from oesteban/maint/transfrom-io-refactor

oesteban · web-flow · commit 455f3374af5d · 2019-10-28T15:31:22.000-07:00
ENH: More comprehensive implementation of ITK affines I/O
diff --git a/nitransforms/io/base.py b/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
+from scipy.io.matlab.mio5 import MatFile5Reader
 
+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()
diff --git a/nitransforms/io/itk.py b/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, TransformFileError
+
+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['index'] = 0
+        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,21 +44,84 @@ 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'
 
         if banner is None:
-            banner = self.structarr['index'] == 1
+            banner = self.structarr['index'] == 0
 
         if banner:
             string = '#Insight Transform File V1.0\n%s'
         return string % self
 
     @classmethod
-    def from_string(klass, string):
+    def from_binary(cls, byte_stream, index=0):
+        """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=0):
+        """Read the struct from a matlab dictionary."""
+        tf = cls()
+        sa = tf.structarr
+
+        sa['index'] = index
+        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
+        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 +132,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 +155,74 @@ 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 = list(value)
+
+        # Update indexes
+        for i, val in enumerate(self.xforms):
+            val['index'] = i
 
     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 TransformFileError("Please use the ITK's new .h5 format.")
+
+        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):
-            xfm.structarr['index'] = i + 1
+        for i, xfm in enumerate(self.xforms):
+            xfm.structarr['index'] = i
             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."""
+        raise TransformFileError("Please use the ITK's new .h5 format.")
+
+    @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 +231,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())
diff --git a/nitransforms/linear.py b/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()
diff --git a/nitransforms/patched.py b/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)
diff --git a/nitransforms/tests/data/itktflist.tfm b/nitransforms/tests/data/itktflist.tfm
diff --git a/nitransforms/tests/test_io.py b/nitransforms/tests/test_io.py
