FEAT: optimal graph-based spatial normalization

Author: balbasty

nitorch/cli/registration/__init__.py

@@ -1,9 +1,12 @@
 from . import affreg
+from . import affopt
 from . import autoreg
 from . import compose
 from . import orient
+from . import meanspace
 from . import register
 from . import reorient
 from . import reslice
 from . import resize
+from . import vopt
 from . import vexp

nitorch/cli/registration/affopt/__init__.py

@@ -0,0 +1 @@
+from . import cli as _

nitorch/cli/registration/affopt/cli.py

@@ -0,0 +1,185 @@
+import sys
+
+from nitorch.cli.cli import commands
+from nitorch.core.cli import ParseError
+from nitorch.io.transforms import loadf, savef
+from nitorch.spatial import optimal_affine
+
+
+def cli(args=None):
+    f"""Command-line interface for `affopt`
+
+    {help}
+
+    """
+
+    # Exceptions are dealt with here
+    try:
+        _cli(args)
+    except AskForHelp:
+        print(help)
+        return
+    except ParseError as e:
+        print(help)
+        print('[ERROR]', e)
+        return
+    # except Exception as e:
+    #     print(f'[ERROR] {str(e)}', file=sys.stderr)
+
+
+commands['affopt'] = cli
+
+help = r"""[nitorch] affopt
+
+Compute optimal "template-to-subject" matrices from "subject-to-subject" pairs.
+
+notes:
+    In Leung et al, all possible pairs of images are registered (in both
+    forward and backward directions), such that the "subject-to-template"
+    transforms can easily be computed as T[i,tpl] = expm(mean_j(logm(T[i,j]))).
+
+    Our implementation differs in several aspects:
+    - We allow some transformation pairs to be missing, at the cost of
+      introducing bias in the mean space estimation. This bias can be
+      overcome in the statistical sense if the number of subjects is large
+      and evaluated pairs are randomly sampled.
+    - Instead of first symmetrizing pairwise transforms, we fit the mean
+      space from all possible forward and backward transformations.
+    - Instead of minimizing the L2 norm in the matrix Lie algebra
+      (which is done implicitly by Leung et al's method), we add
+      the possibility to minimize the L2 norm in the embedding space (i.e.,
+      the Frobenius norm of affine matrices). This method is more accurate
+      when pairwise transformations are large, in which case affine
+      composition is badly approximated by log-matrix addition.
+
+usage:
+    nitorch affopt --input <fix> <mov> <path> [--input ...]
+
+arguments:
+    -i, --input         Affine transform for one pair of images
+                          <fix>   Index (or label) of fixed image
+                          <mov>   Index (or label) of moving image
+                          <path>  Path to an LTA file that warps <mov> to <fix>
+    -o, --output        Path to output transforms (default: {label}_optimal.lta)
+    -l, --log           Minimize L2 in Lie algebra (default: L2 in matrix space)
+    -a, --affine        Assume transforms are all affine (default)
+    -s, --similitude    Assume transforms are all similitude
+    -r, --rigid         Assume transforms are all rigid
+
+example:
+    nitorch affopt \
+      -i mtw pdw mtw_to_pdw.lta \
+      -i mtw t1w mtw_to_t1w.lta \
+      -i pdw mtw pdw_to_mtw.lta \
+      -i pdw t1w pdw_to_t1w.lta \
+      -i t1w mtw t1w_to_mtw.lta \
+      -i t1w pdw t1w_to_pdw.lta \
+      -o out/{label}_to_mean.lta
+
+references:
+    "Consistent multi-time-point brain atrophy estimation from the
+    boundary shift integral"
+    Leung, Ridgway, Ourselin, Fox
+    NeuroImage (2011)
+
+    "Symmetric Diffeomorphic Modeling of Longitudinal Structural MRI"
+    Ashburner, Ridgway
+    Front Neurosci. (2012)
+"""
+
+
+class AskForHelp(Exception):
+    pass
+
+
+def parse(args):
+
+    args = list(args)
+    if not args:
+        raise ParseError('No arguments')
+
+    inputs = {}
+    output = log = affine = similitude = rigid = None
+
+    tags = (
+        '-i', '--input',
+        '-o', '--output',
+        '-l', '--lie', '--log',
+        '-a', '--aff', '--affine',
+        '-s', '--sim', '--similitude',
+        '-r', '--rig', '--rigid'
+    )
+
+    while args:
+        tag = args.pop(0)
+        if tag in ('-h', '--help'):
+            raise AskForHelp
+        elif tag in ('-i', '--input'):
+            fix = args.pop(0)
+            if fix in tags:
+                raise ParseError(f'Expected <fix> <mov> <path> after {tag}')
+            mov = args.pop(0)
+            if mov in tags:
+                raise ParseError(f'Expected <fix> <mov> <path> after {tag}')
+            path = args.pop(0)
+            if path in tags:
+                raise ParseError(f'Expected <fix> <mov> <path> after {tag}')
+            inputs[(fix, mov)] = path
+        elif tag in ('-o', '--output'):
+            out = args.pop(0)
+            if out in tags:
+                raise ParseError(f'Expected <path> after {tag}')
+            if output is not None:
+                raise ParseError(f'Max one {tag} accepted')
+            output = out
+        elif tag in ('-l', '--log', '--lie'):
+            if log is not None:
+                raise ParseError(f'Max one {tag} accepted')
+            log = True
+        elif tag in ('-a', '--aff', '--affine'):
+            if affine is not None:
+                raise ParseError(f'Max one {tag} accepted')
+            affine = True
+        elif tag in ('-s', '--sim', '--similitude'):
+            if similitude is not None:
+                raise ParseError(f'Max one {tag} accepted')
+            similitude = True
+        elif tag in ('-r', '--rig', '--rigid'):
+            if rigid is not None:
+                raise ParseError(f'Max one {tag} accepted')
+            rigid = True
+        else:
+            raise ParseError(f'Unknown tag {tag}')
+
+    output = output or '{label}_optimal.lta'
+    affine = affine or False
+    similitude = similitude or False
+    rigid = rigid or False
+    if int(rigid) + int(similitude) + int(affine) > 1:
+        raise ParseError('Max one of --rigid, --similitude, --affine accepeted')
+    if int(rigid) + int(similitude) + int(affine) == 0:
+        affine = True
+    if affine:
+        basis = 'Aff+'
+    elif similitude:
+        basis = 'CSO'
+    else:
+        basis = 'SE'
+
+    return inputs, output, log, basis
+
+
+def _cli(args):
+
+    inputs, output, log, basis = parse(args)
+
+    # LinearTransformArray(v).matrix() returns array with shape
+    # [N, 4, 4], although N is always 1 in our case.
+    inputs = {k: loadf(v).squeeze()
+              for k, v in inputs.items()}
+    optimal = optimal_affine(inputs, basis=basis,
+                             loss='log' if log else 'exp')
+
+    labels = list(set([label for pair in inputs for label in pair]))
+    for i, label in enumerate(labels):
+        savef(optimal[i], output.format(label=label), type='ras')

nitorch/cli/registration/meanspace/__init__.py

@@ -0,0 +1 @@
+from . import cli as _

nitorch/cli/registration/meanspace/cli.py

@@ -0,0 +1,156 @@
+import sys
+import torch
+import os.path as op
+from tempfile import NamedTemporaryFile
+
+from nitorch.cli.cli import commands
+from nitorch.core.cli import ParseError
+from nitorch import io
+from nitorch.cli.registration.reslice.main import reslice
+from nitorch.spatial import mean_space
+from nitorch.core.py import make_list, fileparts
+from .parser import parser
+
+
+def cli(args=None):
+    f"""Command-line interface for `meanspace`
+
+    {help}
+
+    """
+
+    # Exceptions are dealt with here
+    try:
+        _cli(args)
+    except ParseError as e:
+        print(help)
+        print('[ERROR]', e)
+        return
+    # except Exception as e:
+    #     print(f'[ERROR] {str(e)}', file=sys.stderr)
+
+
+commands['meanspace'] = cli
+
+help = r"""
+Compute an average "oriented voxel space" from a series of oriented images.
+
+usage:
+    nitorch meanspace --input <path> [--input ...]
+
+arguments:
+    -i, --input         Affine transform for one pair of images
+                          <fix>   Index (or label) of fixed image
+                          <mov>   Index (or label) of moving image
+                          <path>  Path to an LTA file that warps <mov> to <fix>
+    -o, --output        Path to output transforms (default: {label}_optimal.lta)
+    -l, --log           Minimize L2 in Lie algebra (default: L2 in matrix space)
+    -a, --affine        Assume transforms are all affine (default)
+    -s, --similitude    Assume transforms are all similitude
+    -r, --rigid         Assume transforms are all rigid
+
+example:
+    optimal_affine \
+      -i mtw pdw mtw_to_pdw.lta \
+      -i mtw t1w mtw_to_t1w.lta \
+      -i pdw mtw pdw_to_mtw.lta \
+      -i pdw t1w pdw_to_t1w.lta \
+      -i t1w mtw t1w_to_mtw.lta \
+      -i t1w pdw t1w_to_pdw.lta \
+      -o out/{label}_to_mean.lta
+
+references:
+    "Consistent multi-time-point brain atrophy estimation from the
+    boundary shift integral"
+    Leung, Ridgway, Ourselin, Fox
+    NeuroImage (2011)
+
+    "Symmetric Diffeomorphic Modeling of Longitudinal Structural MRI"
+    Ashburner, Ridgway
+    Front Neurosci. (2012)
+"""
+
+
+def _cli(args):
+    args = args or sys.argv[1:]
+
+    options = parser.parse(args)
+    if not options:
+        return
+    if options.help:
+        print(help)
+        return
+
+    # get all shape and matrices
+    shapes = []
+    affines = []
+    for path in options.input:
+        f = io.map(path)
+        shapes += [f.shape[:3]]
+        affines += [f.affine]
+
+    ndim = max(map(len, shapes))
+
+    # parse voxel size
+    voxel_size = options.voxel_size
+    voxel_size = make_list(voxel_size or [])
+    if voxel_size and isinstance(voxel_size[-1], str):
+        *voxel_size, vx_unit = voxel_size
+    else:
+        vx_unit = 'mm'
+    if voxel_size:
+        voxel_size = make_list(voxel_size, ndim)
+    else:
+        voxel_size = None
+
+    # parse padding
+    pad = options.pad
+    pad = make_list(pad or [])
+    if pad and isinstance(pad[-1], str):
+        *pad, pad_unit = pad
+    else:
+        pad_unit = '%'
+    if pad:
+        pad = make_list(pad, ndim)
+    else:
+        pad = None
+
+    # compute mean space
+    affine, shape = mean_space(
+        affines,
+        shapes,
+        voxel_size=voxel_size,
+        vx_unit=vx_unit,
+        pad=pad,
+        pad_unit=pad_unit
+    )
+    print(shape)
+    print(affine.numpy())
+
+    dir, base, ext = fileparts(op.abspath(options.input[0]))
+
+    # write mean space image
+    write_meanspace = options.output is not False
+    if not write_meanspace:
+        tmp = NamedTemporaryFile("wb", delete=False, suffix=ext)
+        options.output = tmp.name
+
+    options.output = options.output.format(dir=dir, base=base, ext=ext)
+
+    io.savef(
+        torch.zeros(shape),
+        options.output,
+        affine=affine
+    )
+
+    # write resliced images
+    if options.resliced:
+        for input in options.input:
+            input = op.abspath(input)
+            dir, base, ext = fileparts(input)
+            output = options.resliced.format(dir=dir, base=base, ext=ext)
+            print(output)
+            reslice([input, '-o', output, '-t', options.output])
+
+        if not write_meanspace:
+            tmp.close()