Set up I/O functions to be expandable, deferring some design decisions to the future. Started the beginning of some xfm_utils for quicker testing and validation.

sgiavasis · sgiavasis · commit 7ba96c3bb559 · 2024-07-18T00:58:43.000-04:00
diff --git a/nitransforms/base.py b/nitransforms/base.py
@@ -178,14 +178,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)
@@ -222,6 +256,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 apply(
         self,
         spatialimage,
@@ -338,65 +378,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."""
-        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._matrix)
-        transform_0.create_dataset("Inverse", data=np.linalg.inv(self._matrix))
-
-        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.map)
-        
-        raise NotImplementedError
-    
-    def read_x5(x5_root):
-        variables = {}
-        with h5py.File(x5_root, "r") as f:
-            f.visititems(lambda name, x5_root: _from_hdf5(name, x5_root, variables))
-
-        _transform = variables["TransformGroup/0/Transform"]
-        _inverse = variables["TransformGroup/0/Inverse"]
-        _size = variables["TransformGroup/0/Domain/Size"]
-        _map = variables["TransformGroup/0/Domain/Mapping"]
-
-        return _transform, _inverse, _size, _map
-        
-    def _from_hdf5(name, x5_root, storage):
-        if isinstance(x5_root, h5py.Dataset):
-            storage[name] = {
-                'type': 'dataset',
-                'attrs': dict(x5_root.attrs),
-                'shape': x5_root.shape,
-                'data': x5_root[()]  # Read the data
-            } 
-        elif isinstance(x5_root, h5py.Group):
-            storage[name] = {
-                'type': 'group',
-                'attrs': dict(x5_root.attrs),
-                'members': {}
-            }
-
 
 def _as_homogeneous(xyz, dtype="float32", dim=3):
     """
diff --git a/nitransforms/cli.py b/nitransforms/cli.py
@@ -2,10 +2,13 @@
 import os
 from textwrap import dedent
 
+from nitransforms.io.base import xfm_loader
+from nitransforms.linear import load as linload
+from nitransforms.nonlinear import load as nlinload
 
-from .linear import load as linload
-from .nonlinear import load as nlinload
+from nitransforms.base import TransformBase
 
+import pprint
 
 def cli_apply(pargs):
     """
@@ -45,9 +48,27 @@ 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(
         """
@@ -56,6 +77,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'.
     """
@@ -120,6 +142,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
 
 
@@ -133,3 +166,7 @@ def main(pargs=None):
         subparser = subparsers.choices[pargs.command]
         subparser.print_help()
         raise (e)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/nitransforms/io/base.py b/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."""
 
