enh: x5 implementation

Author: Julien Marabotto

nitransforms/base.py

@@ -248,17 +248,48 @@ def __ne__(self, other):
 class TransformBase:
     """Abstract image class to represent transforms."""
 
-    __slots__ = (
-        "_reference",
-        "_ndim",
-    )
-
-    def __init__(self, reference=None):
+    __slots__ = ("_reference", "_ndim", "_affine", "_shape", "_header",
+        "_grid", "_mapping", "_hdf5_dct", "_x5_dct")
+
+    x5_struct = {
+        'TransformGroup/0': {
+            'Type': None,
+            'Transform': None,
+            'Metadata': None,
+            'Inverse': None
+        },
+        'TransformGroup/0/Domain': {
+            'Grid': None,
+            'Size': None,
+            'Mapping': None
+        },
+        'TransformGroup/1': {},
+        'TransformChain': {}
+    }
+
+    def __init__(self, x5=None, hdf5=None, nifti=None, shape=None, affine=None, 
+        header=None, reference=None):
         """Instantiate a transform."""
+
         self._reference = None
         if reference:
             self.reference = reference
 
+        if nifti is not None:
+            self._x5_dct = self.init_x5_structure(nifti)
+        elif hdf5:
+            self.update_x5_structure(hdf5)
+        elif x5:
+            self.update_x5_structure(x5)
+        
+        self._shape = shape
+        self._affine = affine
+        self._header = header
+
+        # TO-DO
+        self._grid = None
+        self._mapping = None
+
     def __call__(self, x, inverse=False):
         """Apply y = f(x)."""
         return self.map(x, inverse=inverse)
@@ -295,6 +326,12 @@ def ndim(self):
         """Access the dimensions of the reference space."""
         raise TypeError("TransformBase has no dimensions")
 
+    def init_x5_structure(self, xfm_data=None):
+        self.x5_struct['TransformGroup/0/Transform'] = xfm_data
+    
+    def update_x5_structure(self, hdf5_struct=None):
+        self.x5_struct.update(hdf5_struct)
+
     def map(self, x, inverse=False):
         r"""
         Apply :math:`y = f(x)`.
@@ -316,20 +353,6 @@ def map(self, x, inverse=False):
         """
         return x
 
-    def to_filename(self, filename, fmt="X5"):
-        """Store the transform in BIDS-Transforms HDF5 file format (.x5)."""
-        with h5py.File(filename, "w") as out_file:
-            out_file.attrs["Format"] = "X5"
-            out_file.attrs["Version"] = np.uint16(1)
-            root = out_file.create_group("/0")
-            self._to_hdf5(root)
-
-        return filename
-
-    def _to_hdf5(self, x5_root):
-        """Serialize this object into the x5 file format."""
-        raise NotImplementedError
-
     def apply(self, *args, **kwargs):
         """Apply the transform to a dataset.
 

nitransforms/cli.py

@@ -2,11 +2,13 @@
 import os
 from textwrap import dedent
 
+from nitransforms.base import TransformBase
+from nitransforms.io.base import xfm_loader
+from nitransforms.linear import load as linload
+from nitransforms.nonlinear import load as nlinload
+from nitransforms.resampling import apply
 
-from .linear import load as linload
-from .nonlinear import load as nlinload
-from .resampling import apply
-
+import pprint
 
 def cli_apply(pargs):
     """
@@ -47,8 +49,26 @@ def cli_apply(pargs):
         cval=pargs.cval,
         prefilter=pargs.prefilter,
     )
-    moved.to_filename(pargs.out or f"nt_{os.path.basename(pargs.moving)}")
+    #moved.to_filename(pargs.out or f"nt_{os.path.basename(pargs.moving)}")
+
+
+def cli_xfm_util(pargs):
+    """
+    """
+
+    xfm_data = xfm_loader(pargs.transform)
+    xfm_x5 = TransformBase(**xfm_data)
+
+    if pargs.info:
+        pprint.pprint(xfm_x5.x5_struct)
+        print(f"Shape:\n{xfm_x5._shape}")
+        print(f"Affine:\n{xfm_x5._affine}")
 
+    if pargs.x5:
+        filename = f"{os.path.basename(pargs.transform).split('.')[0]}.x5"
+        xfm_x5.to_filename(filename)
+        print(f"Writing out {filename}")
+        
 
 def get_parser():
     desc = dedent(
@@ -58,6 +78,7 @@ def get_parser():
         Commands:
 
             apply       Apply a transformation to an image
+            xfm_util    Assorted transform utilities
 
         For command specific information, use 'nt <command> -h'.
     """
@@ -122,6 +143,17 @@ def _add_subparser(name, description):
         help="Determines if the image's data array is prefiltered with a spline filter before "
         "interpolation (default: True)",
     )
+
+    xfm_util = _add_subparser("xfm_util", cli_xfm_util.__doc__)
+    xfm_util.set_defaults(func=cli_xfm_util)
+    xfm_util.add_argument("transform", help="The transform file")
+    xfm_util.add_argument("--info",
+        action="store_true",
+        help="Get information about the transform")
+    xfm_util.add_argument("--x5",
+        action="store_true",
+        help="Convert transform to .x5 file format.")
+
     return parser, subparsers
 
 
@@ -135,3 +167,7 @@ def main(pargs=None):
         subparser = subparsers.choices[pargs.command]
         subparser.print_help()
         raise (e)
+
+
+if __name__ == "__main__":
+    main()

nitransforms/io/__init__.py

@@ -22,7 +22,7 @@
     "fs": (lta, "FSLinearTransform"),
     "fsl": (fsl, "FSLLinearTransform"),
     "afni": (afni, "AFNILinearTransform"),
-    "x5": (x5, "X5LinearTransform")
+    "x5": (x5, "X5Transform")
 }
 
 

nitransforms/io/base.py

@@ -1,11 +1,137 @@
 """Read/write linear transforms."""
 from pathlib import Path
 import numpy as np
+import nibabel as nb
 from nibabel import load as loadimg
 
+import h5py
+
 from ..patched import LabeledWrapStruct
 
 
+def get_xfm_filetype(xfm_file):
+    path = Path(xfm_file)
+    ext = path.suffix
+    if ext == '.gz' and path.name.endswith('.nii.gz'):
+        return 'nifti'
+    
+    file_types = {
+        '.nii': 'nifti',
+        '.h5': 'hdf5',
+        '.x5': 'x5',
+        '.txt': 'txt',
+        '.mat': 'txt'
+    }
+    return file_types.get(ext, 'unknown')
+
+def gather_fields(x5=None, hdf5=None, nifti=None, shape=None, affine=None, header=None):
+    xfm_fields = {
+        "x5": x5,
+        "hdf5": hdf5,
+        "nifti": nifti,
+        "header": header,
+        "shape": shape,
+        "affine": affine
+    }
+    return xfm_fields
+
+def load_nifti(nifti_file):
+    nifti_xfm = nb.load(nifti_file)
+    xfm_data = nifti_xfm.get_fdata()
+    shape = nifti_xfm.shape
+    affine = nifti_xfm.affine
+    header = getattr(nifti_xfm, "header", None)
+    return gather_fields(nifti=xfm_data, shape=shape, affine=affine, header=header)
+
+def load_hdf5(hdf5_file):
+    storage = {}
+
+    def get_hdf5_items(name, x5_root):
+        if isinstance(x5_root, h5py.Dataset):
+            storage[name] = {
+                'type': 'dataset',
+                'attrs': dict(x5_root.attrs),
+                'shape': x5_root.shape,
+                'data': x5_root[()]
+            } 
+        elif isinstance(x5_root, h5py.Group):
+            storage[name] = {
+                'type': 'group',
+                'attrs': dict(x5_root.attrs),
+                'members': {}
+            }
+
+    with h5py.File(hdf5_file, 'r') as f:
+        f.visititems(get_hdf5_items)
+    if storage:
+        hdf5_storage = {'hdf5': storage}
+    return hdf5_storage
+
+def load_x5(x5_file):
+    load_hdf5(x5_file)
+
+def load_mat(mat_file):
+    affine_matrix = np.loadtxt(mat_file)
+    affine = nb.affines.from_matvec(affine_matrix[:,:3], affine_matrix[:,3])
+    return gather_fields(affine=affine)
+
+def xfm_loader(xfm_file):
+    loaders = {
+        'nifti': load_nifti,
+        'hdf5': load_hdf5,
+        'x5': load_x5,
+        'txt': load_mat,
+        'mat': load_mat
+    }
+    xfm_filetype = get_xfm_filetype(xfm_file)
+    loader = loaders.get(xfm_filetype)
+    if loader is None:
+        raise ValueError(f"Unsupported file type: {xfm_filetype}")
+    return loader(xfm_file)
+
+def to_filename(self, filename, fmt="X5"):
+    """Store the transform in BIDS-Transforms HDF5 file format (.x5)."""
+    with h5py.File(filename, "w") as out_file:
+        out_file.attrs["Format"] = "X5"
+        out_file.attrs["Version"] = np.uint16(1)
+        root = out_file.create_group("/0")
+        self._to_hdf5(root)
+
+    return filename
+
+def _to_hdf5(self, x5_root):
+    """Serialize this object into the x5 file format."""
+    transform_group = x5_root.create_group("TransformGroup")
+
+    """Group '0' containing Affine transform"""
+    transform_0 = transform_group.create_group("0")
+    
+    transform_0.attrs["Type"] = "Affine"
+    transform_0.create_dataset("Transform", data=self._affine)
+    transform_0.create_dataset("Inverse", data=np.linalg.inv(self._affine))
+
+    metadata = {"key": "value"}
+    transform_0.attrs["Metadata"] = str(metadata)
+
+    """sub-group 'Domain' contained within group '0' """
+    domain_group = transform_0.create_group("Domain")
+    #domain_group.attrs["Grid"] = self._grid
+    #domain_group.create_dataset("Size", data=_as_homogeneous(self._reference.shape))
+    #domain_group.create_dataset("Mapping", data=self.mapping)
+    
+def _from_x5(self, x5_root):
+    variables = {}
+
+    x5_root.visititems(lambda name, x5_root: loader(name, x5_root, variables))
+
+    _transform = variables["TransformGroup/0/Transform"]
+    _inverse = variables["TransformGroup/0/Inverse"]
+    _size = variables["TransformGroup/0/Domain/Size"]
+    _mapping = variables["TransformGroup/0/Domain/Mapping"]
+
+    return _transform, _inverse, _size, _map
+    
+
 class TransformIOError(IOError):
     """General I/O exception while reading/writing transforms."""
 

nitransforms/io/x5.py

@@ -13,46 +13,43 @@
     TransformFileError,
 )
 
-LPS = np.diag([-1, -1, 1, 1])
-
-class X5LinearTransform(LinearParameters):
+class X5Transform:
     """A string-based structure for X5 linear transforms."""
 
-    template_dtype = np.dtype(
-        [
-            ("type", "i4"),
-            ("index", "i4"),
-            ("parameters", "f8", (4, 4)),
-            ("offset", "f4", 3),  # Center of rotation
-        ]
-    )
-    dtype = template_dtype
+    _transform = None
 
     def __init__(self, parameters=None, offset=None):
         return
 
     def __str__(self):
         return
 
-    def to_filename(self, filename):
-        '''store this transform to a file with the X5 format'''
-        sa = self.structarr
-        affine = '''some affine that will return a 4x4 array'''
-        return
-    
     @classmethod
     def from_filename(cls, filename):
         """Read the struct from a X5 file given its path."""
         if str(filename).endswith(".h5"):
-            with H5File(str(filename)) as f:
-                return cls.from_h5obj(f)
+            with H5File(str(filename), 'r') as hdf:
+                return cls.from_h5obj(hdf)
+            
+    @classmethod
+    def from_h5obj(cls, h5obj):
+        """Read the transformations in an X5 file."""
+        xfm_list = list(h5obj.keys())
+
+        xfm = xfm_list["Transform"]
+        inv = xfm_list["Inverse"]
+        coords = xfm_list["Size"]
+        map = xfm_list["Mapping"]
+
+        return xfm, inv, coords, map
+
 
 class X5LinearTransformArray(BaseLinearTransformList):
     """A string-based structure for series of X5 linear transforms."""
 
-    _inner_type = X5LinearTransform
+    _inner_type = X5Transform
 
     @property
     def xforms(self):
-        """Get the list of internal ITKLinearTransforms."""
+        """Get the list of internal X5LinearTransforms."""
         return self._xforms