def _warp_image1(image,
target,
shape=None,
affine=None,
nonlin=None,
backward=False,
reslice=False):
"""Returns the warped image, with channel dimension last"""
# build transform
aff_right = target
aff_left = spatial.affine_inv(image.affine)
aff = None
if affine:
# exp = affine.iexp if backward else affine.exp
exp = affine.exp
aff = exp(recompute=False, cache_result=True)
if backward:
aff = spatial.affine_inv(aff)
if nonlin:
if affine:
if affine.position[0].lower() in ('ms' if backward else 'fs'):
aff_right = spatial.affine_matmul(aff, aff_right)
if affine.position[0].lower() in ('fs' if backward else 'ms'):
aff_left = spatial.affine_matmul(aff_left, aff)
exp = nonlin.iexp if backward else nonlin.exp
phi = exp(recompute=False, cache_result=True)
aff_left = spatial.affine_matmul(aff_left, nonlin.affine)
aff_right = spatial.affine_lmdiv(nonlin.affine, aff_right)
if _almost_identity(aff_right) and nonlin.shape == shape:
phi = nonlin.add_identity(phi)
else:
tmp = spatial.affine_grid(aff_right, shape)
phi = regutils.smart_pull_grid(phi, tmp).add_(tmp)
del tmp
if not _almost_identity(aff_left):
phi = spatial.affine_matvec(aff_left, phi)
else:
# no nonlin: single affine even if position == 'symmetric'
if reslice:
aff = spatial.affine_matmul(aff, aff_right)
aff = spatial.affine_matmul(aff_left, aff)
phi = spatial.affine_grid(aff, shape)
else:
phi = None
# warp image
if phi is not None:
warped = image.pull(phi)
else:
warped = image.dat
# write to disk
if len(warped) == 1:
warped = warped[0]
else:
warped = utils.movedim(warped, 0, -1)
return warped
def affine_to_fs(affine, shape, source='voxel', dest='ras'):
"""Convert an affine matrix into FS parameters (vx/cosine/shift)
Parameters
----------
affine : (4, 4) tensor
shape : (int, int, int)
source : {'voxel', 'physical', 'ras'}, default='voxel'
dest : {'voxel', 'physical', 'ras'}, default='ras'
Returns
-------
voxel_size : (float, float, float)
x : (float, float, float)
y : (float, float, float)
z: (float, float, float)
c : (float, float, float)
"""
affine = torch.as_tensor(affine)
backend = dict(dtype=affine.dtype, device=affine.device)
vx = get_voxel_size(affine)
shape = torch.as_tensor(shape, **backend)
source = source.lower()[0]
dest = dest.lower()[0]
shift = shape / 2.
shift = -shift * vx
vox2phys = Orientation(shift, vx).affine()
if (source, dest) in (('v', 'p'), ('p', 'v')):
phys2ras = torch.eye(4, **backend)
elif (source, dest) in (('v', 'r'), ('r', 'v')):
if source == 'r':
affine = affine_inv(affine)
phys2vox = affine_inv(vox2phys)
phys2ras = affine_matmul(affine, phys2vox)
else:
assert (source, dest) in (('p', 'r'), ('r', 'p'))
if source == 'r':
affine = affine_inv(affine)
phys2ras = affine
phys2ras = HomogeneousAffineMatrix(phys2ras)
return (vx.tolist(), phys2ras.xras().tolist(), phys2ras.yras().tolist(),
phys2ras.zras().tolist(), phys2ras.cras().tolist())
def build_from_target(target):
"""Compose all transformations, starting from the final orientation"""
grid = spatial.affine_grid(target.affine.to(**backend), target.shape)
for trf in reversed(options.transformations):
if isinstance(trf, Linear):
grid = spatial.affine_matvec(trf.affine.to(**backend), grid)
else:
mat = trf.affine.to(**backend)
if trf.inv:
vx0 = spatial.voxel_size(mat)
vx1 = spatial.voxel_size(target.affine.to(**backend))
factor = vx0 / vx1
disp, mat = spatial.resize_grid(trf.dat[None],
factor,
affine=mat,
interpolation=trf.spline)
disp = spatial.grid_inv(disp[0], type='disp')
order = 1
else:
disp = trf.dat
order = trf.spline
imat = spatial.affine_inv(mat)
grid = spatial.affine_matvec(imat, grid)
grid += helpers.pull_grid(disp, grid, interpolation=order)
grid = spatial.affine_matvec(mat, grid)
return grid
def collapse_transforms(options):
"""Pre-invert affines and combine sequential affines"""
trfs = []
last_trf = None
for trf in options.transformations:
if isinstance(trf, Linear):
if trf.inv:
trf.affine = spatial.affine_inv(trf.affine)
trf.inv = False
if isinstance(last_trf, Linear):
last_trf.affine = spatial.affine_matmul(
last_trf.affine, trf.affine)
else:
last_trf = trf
else:
if isinstance(last_trf, Linear):
trfs.append(last_trf)
last_trf = None
trfs.append(trf)
if isinstance(last_trf, Linear):
trfs.append(last_trf)
options.transformations = trfs
def crop(inp,
size=None,
center=None,
space='vx',
like=None,
bbox=False,
output=None,
transform=None):
"""Crop a ND volume, while preserving the orientation matrices.
Parameters
----------
inp : str or (tensor, tensor)
Either a path to a volume file or a tuple `(dat, affine)`, where
the first element contains the volume data and the second contains
the orientation matrix.
size : [sequence of] int, optional
Size of the patch to extract.
Its unit and axes are defined by `units` and `layout`.
center : [sequence of] int, optional
Coordinate of the center of the patch.
Its unit and axes are defined by `units` and `layout`.
By default, the center of the FOV is used.
space : [sequence of] {'vox', 'ras'}, default='vox'
The space in which the `size` and `center` parameters are expressed.
bbox : bool or float, default=False
Crop at the bounding box of `inp > threshold`.
If `bbox` is a float, it is the threshold to use.
If `bbox` is `True`, the threshold is 0.
like : str or (tensor, tensor), optional
Reference patch.
Either a path to a volume file or a tuple `(dat, affine)`, where
the first element contains the volume data and the second contains
the orientation matrix.
output : [sequence of] str, optional
Output filename(s).
If the input is not a path, the unstacked data is not written
on disk by default.
If the input is a path, the default output filename is
'{dir}/{base}.{i}{ext}', where `dir`, `base` and `ext`
are the directory, base name and extension of the input file,
`i` is the coordinate (starting at 1) of the slice.
transform : [sequence of] str, optional
Input or output filename(s) of the corresponding transforms.
Not written by default.
If a transform is provided and all other parameters
(i.e., `size` and `like`) are None, the transform is considered
as an input transform to apply.
Returns
-------
output : list[str or (tensor, tensor)]
If the input is a path, the output paths are returned.
Else, the unstacked data and orientation matrices are returned.
"""
dir = ''
base = ''
ext = ''
fname = None
transform_in = False
use_bbox = bool(bbox or isinstance(bbox, float))
# --- Open input ---
is_file = isinstance(inp, str)
if is_file:
fname = inp
f = io.volumes.map(inp)
inp = (f.data(numpy=True) if use_bbox else f, f.affine)
if output is None:
output = '{dir}{sep}{base}.crop{ext}'
dir, base, ext = py.fileparts(fname)
dat, aff0 = inp
dim = aff0.shape[-1] - 1
shape0 = dat.shape[:dim]
layout0 = spatial.affine_to_layout(aff0)
# save input space in case we reorient later
aff00 = aff0
shape00 = shape0
if bool(size) + bool(like) + bool(bbox or isinstance(bbox, float)) > 1:
raise ValueError('Can only use one of `size`, `like` and `bbox`.')
# --- Open reference and compute size/center ---
if like:
like_is_file = isinstance(like, str)
if like_is_file:
f = io.volumes.map(like)
like = (f.shape, f.affine)
like_shape, like_aff = like
like_layout = spatial.affine_to_layout(like_aff)
if (layout0 != like_layout).any():
aff0, dat = spatial.affine_reorient(aff0, dat, like_layout)
shape0 = dat.shape[:dim]
if torch.is_tensor(like_shape):
like_shape = like_shape.shape
size, center, unit, layout = _crop_to_param(aff0, like_aff, like_shape)
space = 'vox'
elif bbox or isinstance(bbox, float):
if bbox is True:
bbox = 0.
mask = torch.as_tensor(dat > bbox)
while mask.dim() > 3:
mask = mask.any(dim=-1)
mins = []
maxs = []
for d in range(dim):
n = mask.shape[d]
idx = utils.movedim(mask, d,
0).reshape([n, -1
]).any(-1).nonzero(as_tuple=False)
mins.append(idx.min())
maxs.append(idx.max())
mins = utils.as_tensor(mins)
maxs = utils.as_tensor(maxs)
size = maxs + 1 - mins
center = (maxs + 1 + mins).float() / 2
space = 'vox'
del mask
# --- Open transformation file and compute size/center ---
elif not size:
if not transform:
raise ValueError('At least one of size/like/transform must '
'be provided')
transform_in = True
t = io.transforms.map(transform)
if not isinstance(t, io.transforms.LinearTransformArray):
raise TypeError('Expected an LTA file')
like_aff, like_shape = t.destination_space()
size, center, unit, layout = _crop_to_param(aff0, like_aff, like_shape)
# --- use center of the FOV ---
if not torch.is_tensor(center) and not center:
center = torch.as_tensor(shape0[:dim], dtype=torch.float)
center = center.sub_(1).mul_(0.5)
# --- convert size/center to voxels ---
size = utils.make_vector(size, dim, dtype=torch.long)
center = utils.make_vector(center, dim, dtype=torch.float)
space_size, space_center = py.make_list(space, 2)
if space_center.lower() == 'ras':
center = spatial.affine_matvec(spatial.affine_inv(aff0), center)
if space_size.lower() == 'ras':
perm = spatial.affine_to_layout(aff0)[:, 0]
size = size[perm.long()]
size = size / spatial.voxel_size(aff0)
# --- compute first/last ---
center = center.float()
size = (size.ceil() if size.dtype.is_floating_point else size).long()
first = center - size.float().sub_(1).mul_(0.5)
first = first.round().long()
last = (first + size).tolist()
first = [max(f, 0) for f in first.tolist()]
last = [min(l, s) for l, s in zip(last, shape0[:dim])]
verb = 'Cropping patch ['
verb += ', '.join([f'{f}:{l}' for f, l in zip(first, last)])
verb += f'] from volume with shape {shape0[:dim]}'
print(verb)
slicer = tuple(slice(f, l) for f, l in zip(first, last))
# --- do crop ---
if use_bbox:
dat = dat.numpy()
dat = dat[slicer]
if not torch.is_tensor(dat):
dat = dat.data(numpy=True)
aff, _ = spatial.affine_sub(aff0, shape0[:dim], slicer)
shape = dat.shape[:dim]
if output:
if is_file:
output = output.format(dir=dir or '.',
base=base,
ext=ext,
sep=os.path.sep)
io.volumes.save(dat, output, like=fname, affine=aff)
else:
output = output.format(sep=os.path.sep)
io.volumes.save(dat, output, affine=aff)
if transform and not transform_in:
if is_file:
transform = transform.format(dir=dir or '.',
base=base,
ext=ext,
sep=os.path.sep)
else:
transform = transform.format(sep=os.path.sep)
trf = io.transforms.LinearTransformArray(transform, 'w')
trf.set_source_space(aff00, shape00)
trf.set_destination_space(aff, shape)
trf.set_metadata({
'src': {
'filename': fname
},
'dst': {
'filename': output
},
'type': 1
}) # RAS_TO_RAS
trf.set_fdata(torch.eye(4))
trf.save()
if is_file:
return output
else:
return dat, aff
def fs_to_affine(shape,
voxel_size=1.,
x=None,
y=None,
z=None,
c=0.,
source='voxel',
dest='ras'):
"""Transform FreeSurfer orientation parameters into an affine matrix.
The returned matrix is effectively a "<source> to <dest>" transform.
Parameters
----------
shape : sequence of int
voxel_size : [sequence of] float, default=1
x : [sequence of] float, default=[1, 0, 0]
y: [sequence of] float, default=[0, 1, 0]
z: [sequence of] float, default=[0, 0, 1]
c: [sequence of] float, default=0
source : {'voxel', 'physical', 'ras'}, default='voxel'
dest : {'voxel', 'physical', 'ras'}, default='ras'
Returns
-------
affine : (4, 4) tensor
"""
dim = len(shape)
shape, voxel_size, x, y, z, c \
= utils.to_max_backend(shape, voxel_size, x, y, z, c)
backend = dict(dtype=shape.dtype, device=shape.device)
voxel_size = utils.make_vector(voxel_size, dim)
if x is None:
x = [1, 0, 0]
if y is None:
y = [0, 1, 0]
if z is None:
z = [0, 0, 1]
x = utils.make_vector(x, dim)
y = utils.make_vector(y, dim)
z = utils.make_vector(z, dim)
c = utils.make_vector(c, dim)
shift = shape / 2.
shift = -shift * voxel_size
vox2phys = Orientation(shift, voxel_size).affine()
phys2ras = XYZC(x, y, z, c).affine()
affines = []
if source.lower().startswith('vox'):
affines.append(vox2phys)
middle_space = 'phys'
elif source.lower().startswith('phys'):
if dest.lower().startswith('vox'):
affines.append(affine_inv(vox2phys))
middle_space = 'vox'
else:
affines.append(phys2ras)
middle_space = 'ras'
elif source.lower() == 'ras':
affines.append(affine_inv(phys2ras))
middle_space = 'phys'
else:
# We need a matrix to switch orientations
affines.append(layout_matrix(source, **backend))
middle_space = 'ras'
if dest.lower().startswith('phys'):
if middle_space == 'vox':
affines.append(vox2phys)
elif middle_space == 'ras':
affines.append(affine_inv(phys2ras))
elif dest.lower().startswith('vox'):
if middle_space == 'phys':
affines.append(affine_inv(vox2phys))
elif middle_space == 'ras':
affines.append(affine_inv(phys2ras))
affines.append(affine_inv(vox2phys))
elif dest.lower().startswith('ras'):
if middle_space == 'phys':
affines.append(phys2ras)
elif middle_space.lower().startswith('vox'):
affines.append(vox2phys)
affines.append(phys2ras)
else:
if middle_space == 'phys':
affines.append(affine_inv(phys2ras))
elif middle_space == 'vox':
affines.append(vox2phys)
affines.append(phys2ras)
layout = layout_matrix(dest, **backend)
affines.append(affine_inv(layout))
affine, *affines = affines
for aff in affines:
affine = affine_matmul(aff, affine)
return affine
def crop(inp,
size=None,
center=None,
space='vx',
like=None,
output=None,
transform=None):
"""Crop a ND volume, while preserving the orientation matrices.
Parameters
----------
inp : str or (tensor, tensor)
Either a path to a volume file or a tuple `(dat, affine)`, where
the first element contains the volume data and the second contains
the orientation matrix.
size : [sequence of] int, optional
Size of the patch to extract.
Its unit and axes are defined by `units` and `layout`.
center : [sequence of] int, optional
Coordinate of the center of the patch.
Its unit and axes are defined by `units` and `layout`.
By default, the center of the FOV is used.
space : [sequence of] {'vox', 'ras'}, default='vox'
The space in which the `size` and `center` parameters are expressed.
like : str or (tensor, tensor), optional
Reference patch.
Either a path to a volume file or a tuple `(dat, affine)`, where
the first element contains the volume data and the second contains
the orientation matrix.
output : [sequence of] str, optional
Output filename(s).
If the input is not a path, the unstacked data is not written
on disk by default.
If the input is a path, the default output filename is
'{dir}/{base}.{i}{ext}', where `dir`, `base` and `ext`
are the directory, base name and extension of the input file,
`i` is the coordinate (starting at 1) of the slice.
transform : [sequence of] str, optional
Input or output filename(s) of the corresponding transforms.
Not written by default.
If a transform is provided and all other parameters
(i.e., `size` and `like`) are None, the transform is considered
as an input transform to apply.
Returns
-------
output : list[str or (tensor, tensor)]
If the input is a path, the output paths are returned.
Else, the unstacked data and orientation matrices are returned.
"""
dir = ''
base = ''
ext = ''
fname = None
transform_in = False
# --- Open input ---
is_file = isinstance(inp, str)
if is_file:
fname = inp
f = io.volumes.map(inp)
inp = (f.data(numpy=True), f.affine)
if output is None:
output = '{dir}{sep}{base}.crop{ext}'
dir, base, ext = py.fileparts(fname)
dat, aff0 = inp
dim = aff0.shape[-1] - 1
shape0 = dat.shape[:dim]
if size and like:
raise ValueError('Cannot use both `size` and `like`.')
# --- Open reference and compute size/center ---
if like:
like_is_file = isinstance(like, str)
if like_is_file:
f = io.volumes.map(like)
like = (f.shape, f.affine)
like_shape, like_aff = like
if torch.is_tensor(like_shape):
like_shape = like_shape.shape
size, center, unit, layout = _crop_to_param(aff0, like_aff, like_shape)
# --- Open transformation file and compute size/center ---
elif not size:
if not transform:
raise ValueError('At least one of size/like/transform must '
'be provided')
transform_in = True
t = io.transforms.map(transform)
if not isinstance(t, io.transforms.LinearTransformArray):
raise TypeError('Expected an LTA file')
like_aff, like_shape = t.destination_space()
size, center, unit, layout = _crop_to_param(aff0, like_aff, like_shape)
# --- use center of the FOV ---
if not torch.is_tensor(center) and not center:
center = torch.as_tensor(shape0[:dim], dtype=torch.float) * 0.5
# --- convert size/center to voxels ---
size = utils.make_vector(size, dim, dtype=torch.long)
center = utils.make_vector(center, dim, dtype=torch.float)
space_size, space_center = py.make_list(space, 2)
if space_center.lower() == 'ras':
center = spatial.affine_matvec(spatial.affine_inv(aff0), center)
if space_size.lower() == 'ras':
perm = spatial.affine_to_layout(aff0)[:, 0]
size = size[perm.long()]
size = size / spatial.voxel_size(aff0)
# --- compute first/last ---
center = center.float()
size = size.ceil().long()
first = (center - size.float() / 2).round().long()
last = (first + size).tolist()
first = [max(f, 0) for f in first.tolist()]
last = [min(l, s) for l, s in zip(last, shape0[:dim])]
verb = 'Cropping patch ['
verb += ', '.join([f'{f}:{l}' for f, l in zip(first, last)])
verb += f'] from volume with shape {shape0[:dim]}'
print(verb)
slicer = tuple(slice(f, l) for f, l in zip(first, last))
# --- do crop ---
dat = dat[slicer]
aff, _ = spatial.affine_sub(aff0, shape0[:dim], slicer)
shape = dat.shape[:dim]
if output:
if is_file:
output = output.format(dir=dir or '.',
base=base,
ext=ext,
sep=os.path.sep)
io.volumes.save(dat, output, like=fname, affine=aff)
else:
output = output.format(sep=os.path.sep)
io.volumes.save(dat, output, affine=aff)
if transform and not transform_in:
if is_file:
transform = transform.format(dir=dir or '.',
base=base,
ext=ext,
sep=os.path.sep)
else:
transform = transform.format(sep=os.path.sep)
trf = io.transforms.LinearTransformArray(transform, 'w')
trf.set_source_space(aff0, shape0)
trf.set_destination_space(aff, shape)
trf.set_metadata({
'src': {
'filename': fname
},
'dst': {
'filename': output
},
'type': 1
}) # RAS_TO_RAS
trf.set_fdata(torch.eye(4))
trf.save()
if is_file:
return output
else:
return dat, aff
def write_data(options):
backend = dict(dtype=torch.float32, device=options.device)
# Pre-exponentiate velocities
for trf in options.transformations:
if isinstance(trf, Velocity):
f = io.volumes.map(trf.file)
trf.affine = f.affine
trf.shape = squeeze_to_nd(f.shape, 3, 1)
trf.dat = f.fdata(**backend).reshape(trf.shape)
trf.shape = trf.shape[:3]
trf.dat = spatial.exp(trf.dat[None],
displacement=True,
inverse=trf.inv)[0]
trf.inv = False
trf.order = 1
elif isinstance(trf, Displacement):
f = io.volumes.map(trf.file)
trf.affine = f.affine
trf.shape = squeeze_to_nd(f.shape, 3, 1)
trf.dat = f.fdata(**backend).reshape(trf.shape)
trf.shape = trf.shape[:3]
if trf.unit == 'mm':
# convert mm displacement to vox displacement
trf.dat = spatial.affine_lmdiv(trf.affine, trf.dat[..., None])
trf.dat = trf.dat[..., 0]
trf.unit = 'vox'
def build_from_target(target):
"""Compose all transformations, starting from the final orientation"""
grid = spatial.affine_grid(target.affine.to(**backend), target.shape)
for trf in reversed(options.transformations):
if isinstance(trf, Linear):
grid = spatial.affine_matvec(trf.affine.to(**backend), grid)
else:
mat = trf.affine.to(**backend)
if trf.inv:
vx0 = spatial.voxel_size(mat)
vx1 = spatial.voxel_size(target.affine.to(**backend))
factor = vx0 / vx1
disp, mat = spatial.resize_grid(trf.dat[None],
factor,
affine=mat,
interpolation=trf.spline)
disp = spatial.grid_inv(disp[0], type='disp')
order = 1
else:
disp = trf.dat
order = trf.spline
imat = spatial.affine_inv(mat)
grid = spatial.affine_matvec(imat, grid)
grid += helpers.pull_grid(disp, grid, interpolation=order)
grid = spatial.affine_matvec(mat, grid)
return grid
if options.target:
# If target is provided, we build a dense transformation field
grid = build_from_target(options.target)
oaffine = options.target.affine
if options.output_unit[0] == 'v':
grid = spatial.affine_matvec(spatial.affine_inv(oaffine), grid)
grid = grid - spatial.identity_grid(grid.shape[:-1],
**utils.backend(grid))
else:
grid = grid - spatial.affine_grid(
oaffine.to(**utils.backend(grid)), grid.shape[:-1])
io.volumes.savef(grid,
options.output.format(ext='.nii.gz'),
affine=oaffine)
else:
if len(options.transformations) > 1:
raise RuntimeError('Something weird happened: '
'multiple transforms and no target')
io.transforms.savef(options.transformations[0].affine,
options.output.format(ext='.lta'))
def compose(self,
orient_in,
deformation,
orient_mean,
affine=None,
orient_out=None,
shape_out=None):
"""Composes a deformation defined in a mean space to an image space.
Parameters
----------
orient_in : (4, 4) tensor
Orientation of the input image
deformation : (*shape_mean, 3) tensor
Random deformation
orient_mean : (4, 4) tensor
Orientation of the mean space (where the deformation is)
affine : (4, 4) tensor, default=identity
Random affine
orient_out : (4, 4) tensor, default=orient_in
Orientation of the output image
shape_out : sequence[int], default=shape_mean
Shape of the output image
Returns
-------
grid : (*shape_out, 3)
Voxel-to-voxel transform
"""
if orient_out is None:
orient_out = orient_in
if shape_out is None:
shape_out = deformation.shape[:-1]
if affine is None:
affine = torch.eye(4,
4,
device=orient_in.device,
dtype=orient_in.dtype)
shape_mean = deformation.shape[:-1]
orient_in, affine, deformation, orient_mean, orient_out \
= utils.to_max_backend(orient_in, affine, deformation, orient_mean, orient_out)
backend = utils.backend(deformation)
eye = torch.eye(4, **backend)
# Compose deformation on the right
right_affine = spatial.affine_lmdiv(orient_mean, orient_out)
if not (shape_mean == shape_out and right_affine.all_close(eye)):
# the mean space and native space are not the same
# we must compose the diffeo with a dense affine transform
# we write the diffeo as an identity plus a displacement
# (id + disp)(aff) = aff + disp(aff)
# -------
# to displacement
deformation = deformation - spatial.identity_grid(
deformation.shape[:-1], **backend)
trf = spatial.affine_grid(right_affine, shape_out)
deformation = spatial.grid_pull(utils.movedim(deformation, -1,
0)[None],
trf[None],
bound='dft',
extrapolate=True)
deformation = utils.movedim(deformation[0], 0, -1)
trf = trf + deformation # add displacement
# Compose deformation on the left
# the output of the diffeo(right) are mean_space voxels
# we must compose on the left with `in\(aff(mean))`
# -------
left_affine = spatial.affine_matmul(spatial.affine_inv(orient_in),
affine)
left_affine = spatial.affine_matmul(left_affine, orient_mean)
trf = spatial.affine_matvec(left_affine, trf)
return trf
def update(moving, fname, inv=False, lin=None, nonlin=None,
interpolation=1, bound='dct2', extrapolate=False, device='cpu',
verbose=True):
"""Apply the linear transform to the header of a file +
apply the non-linear component in the original space.
Notes
-----
.. The shape and general orientation of the moving image is kept untouched.
.. The linear transform is composed with the original orientation matrix.
.. The non-linear component is "wrapped" in the input space, where it
is applied.
.. This function writes a new file (it does not modify the input files
in place).
Parameters
----------
moving : ImageFile
An object describing an image to wrap.
fname : list of str
Output filename for each input file of the moving image
(since images can be encoded over multiple volumes)
inv : bool, default=False
True if we are warping the fixed image to the moving space.
In the case, `moving` should be a `FixedImageFile`.
Else it should be a `MovingImageFile`.
lin : (4, 4) tensor, optional
Linear (or rather affine) transformation
nonlin : dict, optional
Non-linear displacement field, with keys:
disp : (..., 3) tensor
Displacement field (in voxels)
affine : (4, 4) tensor
Orientation matrix of the displacement field
interpolation : int, default=1
bound : str, default='dct2'
extrapolate : bool, default = False
device : default='cpu'
"""
nonlin = nonlin or dict(disp=None, affine=None)
prm = dict(interpolation=interpolation, bound=bound, extrapolate=extrapolate)
moving_affine = moving.affine.to(device)
if inv:
# affine-corrected fixed space
if lin is not None:
new_affine = affine_lmdiv(lin, moving_affine)
else:
new_affine = moving.affine
if nonlin['disp'] is not None:
# moving voxels to param voxels (warps param to moving)
mov2nlin = affine_lmdiv(nonlin['affine'].to(device), moving_affine)
if samespace(mov2nlin, nonlin['disp'].shape[:-1], moving.shape):
g = smalldef(nonlin['disp'].to(device))
else:
g = affine_grid(mov2nlin, moving.shape)
g += pull_grid(nonlin['disp'].to(device), g)
# param to moving
nonlin2mov = affine_inv(mov2nlin)
g = affine_matvec(nonlin2mov, g)
else:
g = None
else:
# affine-corrected moving space
if lin is not None:
new_affine = affine_matmul(lin, moving_affine)
else:
new_affine = moving_affine
if nonlin['disp'] is not None:
# moving voxels to param voxels (warps param to moving)
mov2nlin = affine_lmdiv(nonlin['affine'].to(device), new_affine)
if samespace(mov2nlin, nonlin['disp'].shape[:-1], moving.shape):
g = smalldef(nonlin['disp'].to(device))
else:
g = affine_grid(mov2nlin, moving.shape)
g += pull_grid(nonlin['disp'].to(device), g)
# param to moving
nonlin2mov = affine_inv(mov2nlin)
g = affine_matvec(nonlin2mov, g)
else:
g = None
for file, ofname in zip(moving.files, fname):
if verbose:
print(f'Registered: {file.fname}\n'
f' -> {ofname}')
dat = io.volumes.loadf(file.fname, rand=True, device=device)
dat = dat.reshape([*file.shape, file.channels])
if g is not None:
dat = utils.movedim(dat, -1, 0)
dat = pull(dat, g, **prm)
dat = utils.movedim(dat, 0, -1)
io.savef(dat.cpu(), ofname, like=file.fname, affine=new_affine.cpu())
def write_data(options):
backend = dict(dtype=torch.float32, device=options.device)
# 1) Pre-exponentiate velocities
for trf in options.transformations:
if isinstance(trf, struct.Velocity):
f = io.volumes.map(trf.file)
trf.affine = f.affine
trf.shape = squeeze_to_nd(f.shape, 3, 1)
trf.dat = f.fdata(**backend).reshape(trf.shape)
trf.shape = trf.shape[:3]
trf.dat = spatial.exp(trf.dat[None],
displacement=True,
inverse=trf.inv)[0]
trf.inv = False
trf.order = 1
elif isinstance(trf, struct.Displacement):
f = io.volumes.map(trf.file)
trf.affine = f.affine
trf.shape = squeeze_to_nd(f.shape, 3, 1)
trf.dat = f.fdata(**backend).reshape(trf.shape)
trf.shape = trf.shape[:3]
# 2) If the first transform is linear, compose it with the input
# orientation matrix
if (options.transformations
and isinstance(options.transformations[0], struct.Linear)):
trf = options.transformations[0]
for file in options.files:
mat = file.affine.to(**backend)
aff = trf.affine.to(**backend)
file.affine = spatial.affine_lmdiv(aff, mat)
options.transformations = options.transformations[1:]
def build_from_target(target):
"""Compose all transformations, starting from the final orientation"""
grid = spatial.affine_grid(target.affine.to(**backend), target.shape)
for trf in reversed(options.transformations):
if isinstance(trf, struct.Linear):
grid = spatial.affine_matvec(trf.affine.to(**backend), grid)
else:
mat = trf.affine.to(**backend)
if trf.inv:
vx0 = spatial.voxel_size(mat)
vx1 = spatial.voxel_size(target.affine.to(**backend))
factor = vx0 / vx1
disp, mat = spatial.resize_grid(trf.dat[None],
factor,
affine=mat,
interpolation=trf.order)
disp = spatial.grid_inv(disp[0], type='disp')
order = 1
else:
disp = trf.dat
order = trf.order
imat = spatial.affine_inv(mat)
grid = spatial.affine_matvec(imat, grid)
grid += helpers.pull_grid(disp, grid, interpolation=order)
grid = spatial.affine_matvec(mat, grid)
return grid
# 3) If target is provided, we can build most of the transform once
# and just multiply it with a input-wise affine matrix.
if options.target:
grid = build_from_target(options.target)
oaffine = options.target.affine
# 4) Loop across input files
opt = dict(interpolation=options.interpolation,
bound=options.bound,
extrapolate=options.extrapolate)
output = utils.make_list(options.output, len(options.files))
for file, ofname in zip(options.files, output):
ofname = ofname.format(dir=file.dir, base=file.base, ext=file.ext)
print(f'Reslicing: {file.fname}\n' f' -> {ofname}')
dat = io.volumes.loadf(file.fname, rand=True, **backend)
dat = dat.reshape([*file.shape, file.channels])
dat = utils.movedim(dat, -1, 0)
if not options.target:
grid = build_from_target(file)
oaffine = file.affine
mat = file.affine.to(**backend)
imat = spatial.affine_inv(mat)
dat = helpers.pull(dat, spatial.affine_matvec(imat, grid), **opt)
dat = utils.movedim(dat, 0, -1)
io.volumes.savef(dat, ofname, like=file.fname, affine=oaffine)
def get_oriented_slice(image, dim=-1, index=None, affine=None,
space=None, bbox=None, interpolation=1,
transpose_sagittal=False, return_index=False,
return_mat=False):
"""Sample a slice in a RAS system
Parameters
----------
image : (..., *shape3)
dim : int, default=-1
Index of spatial dimension to sample in the visualization space
If RAS: -1 = axial / -2 = coronal / -3 = sagittal
index : int, default=shape//2
Coordinate (in voxel) of the slice to extract
affine : (4, 4) tensor, optional
Orientation matrix of the image
space : (4, 4) tensor, optional
Orientation matrix of the visualisation space.
Default: RAS with minimum voxel size of all inputs.
bbox : (2, D) tensor_like, optional
Bounding box: min and max coordinates (in millimetric
visualisation space). Default: bounding box of the input image.
interpolation : {0, 1}, default=1
Interpolation order.
Returns
-------
slice : (..., *shape2) tensor
Slice in the visualisation space.
"""
# preproc dim
if isinstance(dim, str):
dim = dim.lower()[0]
if dim == 'a':
dim = -1
if dim == 'c':
dim = -2
if dim == 's':
dim = -3
backend = utils.backend(image)
# compute default space (mn/mx are in voxels)
affine, shape = _get_default_native(affine, image.shape[-3:])
space, mn, mx = _get_default_space(affine, [shape], space, bbox)
affine, shape = (affine[0], shape[0])
# compute default cursor (in voxels)
if index is None:
index = (mx + mn) / 2
else:
index = torch.as_tensor(index)
index = spatial.affine_matvec(spatial.affine_inv(space), index)
# include slice to volume matrix
shape = tuple(((mx-mn) + 1).round().int().tolist())
if dim == -1:
# axial
shift = [[1, 0, 0, - mn[0] + 1],
[0, 1, 0, - mn[1] + 1],
[0, 0, 1, - index[2]],
[0, 0, 0, 1]]
shift = utils.as_tensor(shift, **backend)
shape = shape[:-1]
index = (index[0] - mn[0] + 1, index[1] - mn[1] + 1)
elif dim == -2:
# coronal
shift = [[1, 0, 0, - mn[0] + 1],
[0, 0, 1, - mn[2] + 1],
[0, 1, 0, - index[1]],
[0, 0, 0, 1]]
shift = utils.as_tensor(shift, **backend)
shape = (shape[0], shape[2])
index = (index[0] - mn[0] + 1, index[2] - mn[2] + 1)
elif dim == -3:
# sagittal
if not transpose_sagittal:
shift = [[0, 0, 1, - mn[2] + 1],
[0, 1, 0, - mn[1] + 1],
[1, 0, 0, - index[0]],
[0, 0, 0, 1]]
shift = utils.as_tensor(shift, **backend)
shape = (shape[2], shape[1])
index = (index[2] - mn[2] + 1, index[1] - mn[1] + 1)
else:
shift = [[0, -1, 0, mx[1] + 1],
[0, 0, 1, - mn[2] + 1],
[1, 0, 0, - index[0]],
[0, 0, 0, 1]]
shift = utils.as_tensor(shift, **backend)
shape = (shape[1], shape[2])
index = (mx[1] + 1 - index[1], index[2] - mn[2] + 1)
else:
raise ValueError(f'Unknown dimension {dim}')
# sample
space = spatial.affine_rmdiv(space, shift)
affine = spatial.affine_lmdiv(affine, space)
affine = affine.to(**backend)
grid = spatial.affine_grid(affine, [*shape, 1])
*channel, s0, s1, s2 = image.shape
imshape = (s0, s1, s2)
image = image.reshape([1, -1, *imshape])
image = spatial.grid_pull(image, grid[None], interpolation=interpolation,
bound='dct2', extrapolate=False)
image = image.reshape([*channel, *shape])
return ((image, index, space) if return_index and return_mat else
(image, index) if return_index else
(image, space) if return_mat else image)
def write_data(options):
backend = dict(dtype=torch.float32, device=options.device)
# 1) Pre-exponentiate velocities
for trf in options.transformations:
if isinstance(trf, struct.Velocity):
f = io.volumes.map(trf.file)
trf.affine = f.affine
trf.shape = squeeze_to_nd(f.shape, 3, 1)
trf.dat = f.fdata(**backend).reshape(trf.shape)
trf.shape = trf.shape[:3]
if trf.json:
with open(trf.json) as f:
prm = json.load(f)
prm['voxel_size'] = spatial.voxel_size(trf.affine)
trf.dat = spatial.shoot(trf.dat[None],
displacement=True,
return_inverse=trf.inv)
if trf.inv:
trf.dat = trf.dat[-1]
else:
trf.dat = spatial.exp(trf.dat[None],
displacement=True,
inverse=trf.inv)
trf.dat = trf.dat[0] # drop batch dimension
trf.inv = False
trf.order = 1
elif isinstance(trf, struct.Displacement):
f = io.volumes.map(trf.file)
trf.affine = f.affine
trf.shape = squeeze_to_nd(f.shape, 3, 1)
trf.dat = f.fdata(**backend).reshape(trf.shape)
trf.shape = trf.shape[:3]
if trf.unit == 'mm':
# convert mm displacement to vox displacement
trf.dat = spatial.affine_lmdiv(trf.affine, trf.dat[..., None])
trf.dat = trf.dat[..., 0]
trf.unit = 'vox'
# 2) If the first transform is linear, compose it with the input
# orientation matrix
if (options.transformations
and isinstance(options.transformations[0], struct.Linear)):
trf = options.transformations[0]
for file in options.files:
mat = file.affine.to(**backend)
aff = trf.affine.to(**backend)
file.affine = spatial.affine_lmdiv(aff, mat)
options.transformations = options.transformations[1:]
def build_from_target(affine, shape):
"""Compose all transformations, starting from the final orientation"""
grid = spatial.affine_grid(affine.to(**backend), shape)
for trf in reversed(options.transformations):
if isinstance(trf, struct.Linear):
grid = spatial.affine_matvec(trf.affine.to(**backend), grid)
else:
mat = trf.affine.to(**backend)
if trf.inv:
vx0 = spatial.voxel_size(mat)
vx1 = spatial.voxel_size(affine.to(**backend))
factor = vx0 / vx1
disp, mat = spatial.resize_grid(trf.dat[None],
factor,
affine=mat,
interpolation=trf.order)
disp = spatial.grid_inv(disp[0], type='disp')
order = 1
else:
disp = trf.dat
order = trf.order
imat = spatial.affine_inv(mat)
grid = spatial.affine_matvec(imat, grid)
grid += helpers.pull_grid(disp, grid, interpolation=order)
grid = spatial.affine_matvec(mat, grid)
return grid
# 3) If target is provided, we can build most of the transform once
# and just multiply it with a input-wise affine matrix.
if options.target:
grid = build_from_target(options.target.affine, options.target.shape)
oaffine = options.target.affine
# 4) Loop across input files
opt_pull0 = dict(interpolation=options.interpolation,
bound=options.bound,
extrapolate=options.extrapolate)
opt_coeff = dict(interpolation=options.interpolation,
bound=options.bound,
dim=3,
inplace=True)
output = py.make_list(options.output, len(options.files))
for file, ofname in zip(options.files, output):
is_label = isinstance(options.interpolation,
str) and options.interpolation == 'l'
ofname = ofname.format(dir=file.dir, base=file.base, ext=file.ext)
print(f'Reslicing: {file.fname}\n' f' -> {ofname}')
if is_label:
backend_int = dict(dtype=torch.long, device=backend['device'])
dat = io.volumes.load(file.fname, **backend_int)
opt_pull = dict(opt_pull0)
opt_pull['interpolation'] = 1
else:
dat = io.volumes.loadf(file.fname,
rand=options.interpolation > 0,
**backend)
opt_pull = opt_pull0
dat = dat.reshape([*file.shape, file.channels])
dat = utils.movedim(dat, -1, 0)
if not options.target:
oaffine = file.affine
oshape = file.shape
if options.voxel_size:
ovx = utils.make_vector(options.voxel_size,
3,
dtype=oaffine.dtype)
factor = spatial.voxel_size(oaffine) / ovx
oaffine, oshape = spatial.affine_resize(oaffine,
oshape,
factor=factor,
anchor='f')
grid = build_from_target(oaffine, oshape)
mat = file.affine.to(**backend)
imat = spatial.affine_inv(mat)
if options.prefilter and not is_label:
dat = spatial.spline_coeff_nd(dat, **opt_coeff)
dat = helpers.pull(dat, spatial.affine_matvec(imat, grid), **opt_pull)
dat = utils.movedim(dat, 0, -1)
if is_label:
io.volumes.save(dat, ofname, like=file.fname, affine=oaffine)
else:
io.volumes.savef(dat, ofname, like=file.fname, affine=oaffine)
