enh: guess open linear transform formats

EAFP implementation of loading linear transforms without specifying the
format of the file.

Resolves: #86.
Resolves: #87.
Resolves: #107.

Author: oesteban

nitransforms/io/__init__.py

@@ -1,11 +1,33 @@
 # emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
 # vi: set ft=python sts=4 ts=4 sw=4 et:
 """Read and write transforms."""
-from . import afni, fsl, itk, lta
+from nitransforms.io import afni, fsl, itk, lta
+from nitransforms.io.base import TransformFileError
 
 __all__ = [
     "afni",
     "fsl",
     "itk",
     "lta",
+    "get_linear_factory",
+    "TransformFileError",
 ]
+
+_IO_TYPES = {
+    "itk": (itk, "ITKLinearTransform"),
+    "ants": (itk, "ITKLinearTransform"),
+    "elastix": (itk, "ITKLinearTransform"),
+    "lta": (lta, "FSLinearTransform"),
+    "fs": (lta, "FSLinearTransform"),
+    "fsl": (fsl, "FSLLinearTransform"),
+    "afni": (afni, "AFNILinearTransform"),
+}
+
+
+def get_linear_factory(fmt, is_array=True):
+    """Return the type required by a given format."""
+    if fmt.lower() not in _IO_TYPES:
+        raise TypeError(f"Unsupported transform format <{fmt}>.")
+
+    module, classname = _IO_TYPES[fmt.lower()]
+    return getattr(module, f"{classname}{'Array' * is_array}")

nitransforms/linear.py

@@ -14,14 +14,14 @@
 
 from nibabel.loadsave import load as _nbload
 
-from .base import (
+from nitransforms.base import (
     ImageGrid,
     TransformBase,
     SpatialReference,
     _as_homogeneous,
     EQUALITY_TOL,
 )
-from . import io
+from nitransforms.io import get_linear_factory, TransformFileError
 
 
 class Affine(TransformBase):
@@ -183,51 +183,40 @@ def _to_hdf5(self, x5_root):
             self.reference._to_hdf5(x5_root.create_group("Reference"))
 
     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.ITKLinearTransform.from_ras(self.matrix)
-            itkobj.to_filename(filename)
-            return filename
-
-        # Rest of the formats peek into moving and reference image grids
-        moving = ImageGrid(moving) if moving is not None else self.reference
-
-        _factory = {
-            "afni": io.afni.AFNILinearTransform,
-            "fsl": io.fsl.FSLLinearTransform,
-            "lta": io.lta.FSLinearTransform,
-            "fs": io.lta.FSLinearTransform,
-        }
-
-        if fmt not in _factory:
-            raise NotImplementedError(f"Unsupported format <{fmt}>")
-
-        _factory[fmt].from_ras(
-            self.matrix, moving=moving, reference=self.reference
-        ).to_filename(filename)
-        return filename
+        """Store the transform in the requested output format."""
+        writer = get_linear_factory(fmt, is_array=False)
 
-    @classmethod
-    def from_filename(cls, filename, fmt="X5", reference=None, moving=None):
-        """Create an affine from a transform file."""
         if fmt.lower() in ("itk", "ants", "elastix"):
-            _factory = io.itk.ITKLinearTransformArray
-        elif fmt.lower() in ("lta", "fs"):
-            _factory = io.lta.FSLinearTransformArray
-        elif fmt.lower() == "fsl":
-            _factory = io.fsl.FSLLinearTransformArray
-        elif fmt.lower() == "afni":
-            _factory = io.afni.AFNILinearTransformArray
+            writer.from_ras(self.matrix).to_filename(filename)
         else:
-            raise NotImplementedError
+            # Rest of the formats peek into moving and reference image grids
+            writer.from_ras(
+                self.matrix,
+                reference=self.reference,
+                moving=ImageGrid(moving) if moving is not None else self.reference,
+            ).to_filename(filename)
+        return filename
 
-        struct = _factory.from_filename(filename)
-        matrix = struct.to_ras(reference=reference, moving=moving)
-        if cls == Affine:
-            if np.shape(matrix)[0] != 1:
-                raise TypeError("Cannot load transform array '%s'" % filename)
-            matrix = matrix[0]
-        return cls(matrix, reference=reference)
+    @classmethod
+    def from_filename(cls, filename, fmt=None, reference=None, moving=None):
+        """Create an affine from a transform file."""
+        fmtlist = [fmt] if fmt is not None else ("itk", "lta", "afni", "fsl")
+
+        for potential_fmt in fmtlist:
+            try:
+                struct = get_linear_factory(potential_fmt).from_filename(filename)
+                matrix = struct.to_ras(reference=reference, moving=moving)
+                if cls == Affine:
+                    if np.shape(matrix)[0] != 1:
+                        raise TypeError("Cannot load transform array '%s'" % filename)
+                    matrix = matrix[0]
+                return cls(matrix, reference=reference)
+            except TransformFileError:
+                continue
+
+        raise TransformFileError(
+            f"Could not open <{filename}> (formats tried: {', '.join(fmtlist)}."
+        )
 
     def __repr__(self):
         """
@@ -353,31 +342,18 @@ def map(self, x, inverse=False):
         return np.swapaxes(affine.dot(coords), 1, 2)
 
     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.from_ras(self.matrix)
-            itkobj.to_filename(filename)
-            return filename
+        """Store the transform in the requested output format."""
+        writer = get_linear_factory(fmt, is_array=True)
 
-        # Rest of the formats peek into moving and reference image grids
-        if moving is not None:
-            moving = ImageGrid(moving)
+        if fmt.lower() in ("itk", "ants", "elastix"):
+            writer.from_ras(self.matrix).to_filename(filename)
         else:
-            moving = self.reference
-
-        _factory = {
-            "afni": io.afni.AFNILinearTransformArray,
-            "fsl": io.fsl.FSLLinearTransformArray,
-            "lta": io.lta.FSLinearTransformArray,
-            "fs": io.lta.FSLinearTransformArray,
-        }
-
-        if fmt not in _factory:
-            raise NotImplementedError(f"Unsupported format <{fmt}>")
-
-        _factory[fmt].from_ras(
-            self.matrix, moving=moving, reference=self.reference
-        ).to_filename(filename)
+            # Rest of the formats peek into moving and reference image grids
+            writer.from_ras(
+                self.matrix,
+                reference=self.reference,
+                moving=ImageGrid(moving) if moving is not None else self.reference,
+            ).to_filename(filename)
         return filename
 
     def apply(
@@ -486,7 +462,7 @@ def apply(
         return resampled
 
 
-def load(filename, fmt="X5", reference=None, moving=None):
+def load(filename, fmt=None, reference=None, moving=None):
     """
     Load a linear transform file.
 

nitransforms/tests/test_linear.py

@@ -11,7 +11,8 @@
 import nibabel as nb
 from nibabel.eulerangles import euler2mat
 from nibabel.affines import from_matvec
-from .. import linear as nitl
+from nitransforms import linear as nitl
+from nitransforms import io
 from .utils import assert_affines_by_filename
 
 RMSE_TOL = 0.1
@@ -152,7 +153,7 @@ def test_linear_save(tmpdir, data_path, get_testdata, image_orientation, sw_tool
     xfm = (
         nitl.Affine(T) if (sw_tool, image_orientation) != ("afni", "oblique") else
         # AFNI is special when moving or reference are oblique - let io do the magic
-        nitl.Affine(nitl.io.afni.AFNILinearTransform.from_ras(T).to_ras(
+        nitl.Affine(io.afni.AFNILinearTransform.from_ras(T).to_ras(
             reference=img,
             moving=img,
         ))
@@ -199,7 +200,7 @@ def test_apply_linear_transform(tmpdir, get_testdata, get_testmask, image_orient
     xfm_fname = "M.%s%s" % (sw_tool, ext)
     # Change reference dataset for AFNI & oblique
     if (sw_tool, image_orientation) == ("afni", "oblique"):
-        nitl.io.afni.AFNILinearTransform.from_ras(
+        io.afni.AFNILinearTransform.from_ras(
             T,
             moving=img,
             reference=img,