peirong26's picture
Upload 187 files
2571f24 verified
"""High level interpolation API"""
__all__ = ['grid_pull', 'grid_push', 'grid_count', 'grid_grad',
'spline_coeff', 'spline_coeff_nd',
'identity_grid', 'add_identity_grid', 'add_identity_grid_']
import torch
from .utils import expanded_shape, matvec
from .jit_utils import movedim1, meshgrid
from .autograd import (GridPull, GridPush, GridCount, GridGrad,
SplineCoeff, SplineCoeffND)
from . import backend, jitfields
_doc_interpolation = \
"""`interpolation` can be an int, a string or an InterpolationType.
Possible values are:
- 0 or 'nearest'
- 1 or 'linear'
- 2 or 'quadratic'
- 3 or 'cubic'
- 4 or 'fourth'
- 5 or 'fifth'
- etc.
A list of values can be provided, in the order [W, H, D],
to specify dimension-specific interpolation orders."""
_doc_bound = \
"""`bound` can be an int, a string or a BoundType.
Possible values are:
- 'replicate' or 'nearest' : a a a | a b c d | d d d
- 'dct1' or 'mirror' : d c b | a b c d | c b a
- 'dct2' or 'reflect' : c b a | a b c d | d c b
- 'dst1' or 'antimirror' : -b -a 0 | a b c d | 0 -d -c
- 'dst2' or 'antireflect' : -c -b -a | a b c d | -d -c -b
- 'dft' or 'wrap' : b c d | a b c d | a b c
- 'zero' or 'zeros' : 0 0 0 | a b c d | 0 0 0
A list of values can be provided, in the order [W, H, D],
to specify dimension-specific boundary conditions.
Note that
- `dft` corresponds to circular padding
- `dct2` corresponds to Neumann boundary conditions (symmetric)
- `dst2` corresponds to Dirichlet boundary conditions (antisymmetric)
See https://en.wikipedia.org/wiki/Discrete_cosine_transform
https://en.wikipedia.org/wiki/Discrete_sine_transform"""
_doc_bound_coeff = \
"""`bound` can be an int, a string or a BoundType.
Possible values are:
- 'replicate' or 'nearest' : a a a | a b c d | d d d
- 'dct1' or 'mirror' : d c b | a b c d | c b a
- 'dct2' or 'reflect' : c b a | a b c d | d c b
- 'dst1' or 'antimirror' : -b -a 0 | a b c d | 0 -d -c
- 'dst2' or 'antireflect' : -c -b -a | a b c d | -d -c -b
- 'dft' or 'wrap' : b c d | a b c d | a b c
- 'zero' or 'zeros' : 0 0 0 | a b c d | 0 0 0
A list of values can be provided, in the order [W, H, D],
to specify dimension-specific boundary conditions.
Note that
- `dft` corresponds to circular padding
- `dct1` corresponds to mirroring about the center of the first/last voxel
- `dct2` corresponds to mirroring about the edge of the first/last voxel
See https://en.wikipedia.org/wiki/Discrete_cosine_transform
https://en.wikipedia.org/wiki/Discrete_sine_transform
/!\ Only 'dct1', 'dct2' and 'dft' are implemented for interpolation
orders >= 6."""
_ref_coeff = \
"""..[1] M. Unser, A. Aldroubi and M. Eden.
"B-Spline Signal Processing: Part I-Theory,"
IEEE Transactions on Signal Processing 41(2):821-832 (1993).
..[2] M. Unser, A. Aldroubi and M. Eden.
"B-Spline Signal Processing: Part II-Efficient Design and Applications,"
IEEE Transactions on Signal Processing 41(2):834-848 (1993).
..[3] M. Unser.
"Splines: A Perfect Fit for Signal and Image Processing,"
IEEE Signal Processing Magazine 16(6):22-38 (1999).
"""
def _preproc(grid, input=None, mode=None):
"""Preprocess tensors for pull/push/count/grad
Low level bindings expect inputs of shape
[batch, channel, *spatial] and [batch, *spatial, dim], whereas
the high level python API accepts inputs of shape
[..., [channel], *spatial] and [..., *spatial, dim].
This function broadcasts and reshapes the input tensors accordingly.
/!\\ This *can* trigger large allocations /!\\
"""
dim = grid.shape[-1]
if input is None:
spatial = grid.shape[-dim-1:-1]
batch = grid.shape[:-dim-1]
grid = grid.reshape([-1, *spatial, dim])
info = dict(batch=batch, channel=[1] if batch else [], dim=dim)
return grid, info
grid_spatial = grid.shape[-dim-1:-1]
grid_batch = grid.shape[:-dim-1]
input_spatial = input.shape[-dim:]
channel = 0 if input.dim() == dim else input.shape[-dim-1]
input_batch = input.shape[:-dim-1]
if mode == 'push':
grid_spatial = input_spatial = expanded_shape(grid_spatial, input_spatial)
# broadcast and reshape
batch = expanded_shape(grid_batch, input_batch)
grid = grid.expand([*batch, *grid_spatial, dim])
grid = grid.reshape([-1, *grid_spatial, dim])
input = input.expand([*batch, channel or 1, *input_spatial])
input = input.reshape([-1, channel or 1, *input_spatial])
out_channel = [channel] if channel else ([1] if batch else [])
info = dict(batch=batch, channel=out_channel, dim=dim)
return grid, input, info
def _postproc(out, shape_info, mode):
"""Postprocess tensors for pull/push/count/grad"""
dim = shape_info['dim']
if mode != 'grad':
spatial = out.shape[-dim:]
feat = []
else:
spatial = out.shape[-dim-1:-1]
feat = [out.shape[-1]]
batch = shape_info['batch']
channel = shape_info['channel']
out = out.reshape([*batch, *channel, *spatial, *feat])
return out
def grid_pull(input, grid, interpolation='linear', bound='zero',
extrapolate=False, prefilter=False):
"""Sample an image with respect to a deformation field.
Notes
-----
{interpolation}
{bound}
If the input dtype is not a floating point type, the input image is
assumed to contain labels. Then, unique labels are extracted
and resampled individually, making them soft labels. Finally,
the label map is reconstructed from the individual soft labels by
assigning the label with maximum soft value.
Parameters
----------
input : (..., [channel], *inshape) tensor
Input image.
grid : (..., *outshape, dim) tensor
Transformation field.
interpolation : int or sequence[int], default=1
Interpolation order.
bound : BoundType or sequence[BoundType], default='zero'
Boundary conditions.
extrapolate : bool or int, default=True
Extrapolate out-of-bound data.
prefilter : bool, default=False
Apply spline pre-filter (= interpolates the input)
Returns
-------
output : (..., [channel], *outshape) tensor
Deformed image.
"""
if backend.jitfields and jitfields.available:
return jitfields.grid_pull(input, grid, interpolation, bound,
extrapolate, prefilter)
grid, input, shape_info = _preproc(grid, input)
batch, channel = input.shape[:2]
dim = grid.shape[-1]
if not input.dtype.is_floating_point:
# label map -> specific processing
out = input.new_zeros([batch, channel, *grid.shape[1:-1]])
pmax = grid.new_zeros([batch, channel, *grid.shape[1:-1]])
for label in input.unique():
soft = (input == label).to(grid.dtype)
if prefilter:
input = spline_coeff_nd(soft, interpolation=interpolation,
bound=bound, dim=dim, inplace=True)
soft = GridPull.apply(soft, grid, interpolation, bound, extrapolate)
out[soft > pmax] = label
pmax = torch.max(pmax, soft)
else:
if prefilter:
input = spline_coeff_nd(input, interpolation=interpolation,
bound=bound, dim=dim)
out = GridPull.apply(input, grid, interpolation, bound, extrapolate)
return _postproc(out, shape_info, mode='pull')
def grid_push(input, grid, shape=None, interpolation='linear', bound='zero',
extrapolate=False, prefilter=False):
"""Splat an image with respect to a deformation field (pull adjoint).
Notes
-----
{interpolation}
{bound}
Parameters
----------
input : (..., [channel], *inshape) tensor
Input image.
grid : (..., *inshape, dim) tensor
Transformation field.
shape : sequence[int], default=inshape
Output shape
interpolation : int or sequence[int], default=1
Interpolation order.
bound : BoundType, or sequence[BoundType], default='zero'
Boundary conditions.
extrapolate : bool or int, default=True
Extrapolate out-of-bound data.
prefilter : bool, default=False
Apply spline pre-filter.
Returns
-------
output : (..., [channel], *shape) tensor
Spatted image.
"""
if backend.jitfields and jitfields.available:
return jitfields.grid_push(input, grid, shape, interpolation, bound,
extrapolate, prefilter)
grid, input, shape_info = _preproc(grid, input, mode='push')
dim = grid.shape[-1]
if shape is None:
shape = tuple(input.shape[2:])
out = GridPush.apply(input, grid, shape, interpolation, bound, extrapolate)
if prefilter:
out = spline_coeff_nd(out, interpolation=interpolation, bound=bound,
dim=dim, inplace=True)
return _postproc(out, shape_info, mode='push')
def grid_count(grid, shape=None, interpolation='linear', bound='zero',
extrapolate=False):
"""Splatting weights with respect to a deformation field (pull adjoint).
Notes
-----
{interpolation}
{bound}
Parameters
----------
grid : (..., *inshape, dim) tensor
Transformation field.
shape : sequence[int], default=inshape
Output shape
interpolation : int or sequence[int], default=1
Interpolation order.
bound : BoundType, or sequence[BoundType], default='zero'
Boundary conditions.
extrapolate : bool or int, default=True
Extrapolate out-of-bound data.
Returns
-------
output : (..., [1], *shape) tensor
Splatted weights.
"""
if backend.jitfields and jitfields.available:
return jitfields.grid_count(grid, shape, interpolation, bound, extrapolate)
grid, shape_info = _preproc(grid)
out = GridCount.apply(grid, shape, interpolation, bound, extrapolate)
return _postproc(out, shape_info, mode='count')
def grid_grad(input, grid, interpolation='linear', bound='zero',
extrapolate=False, prefilter=False):
"""Sample spatial gradients of an image with respect to a deformation field.
Notes
-----
{interpolation}
{bound}
Parameters
----------
input : (..., [channel], *inshape) tensor
Input image.
grid : (..., *inshape, dim) tensor
Transformation field.
shape : sequence[int], default=inshape
Output shape
interpolation : int or sequence[int], default=1
Interpolation order.
bound : BoundType, or sequence[BoundType], default='zero'
Boundary conditions.
extrapolate : bool or int, default=True
Extrapolate out-of-bound data.
prefilter : bool, default=False
Apply spline pre-filter (= interpolates the input)
Returns
-------
output : (..., [channel], *shape, dim) tensor
Sampled gradients.
"""
if backend.jitfields and jitfields.available:
return jitfields.grid_grad(input, grid, interpolation, bound,
extrapolate, prefilter)
grid, input, shape_info = _preproc(grid, input)
dim = grid.shape[-1]
if prefilter:
input = spline_coeff_nd(input, interpolation, bound, dim)
out = GridGrad.apply(input, grid, interpolation, bound, extrapolate)
return _postproc(out, shape_info, mode='grad')
def spline_coeff(input, interpolation='linear', bound='dct2', dim=-1,
inplace=False):
"""Compute the interpolating spline coefficients, for a given spline order
and boundary conditions, along a single dimension.
Notes
-----
{interpolation}
{bound}
References
----------
{ref}
Parameters
----------
input : tensor
Input image.
interpolation : int or sequence[int], default=1
Interpolation order.
bound : BoundType or sequence[BoundType], default='dct1'
Boundary conditions.
dim : int, default=-1
Dimension along which to process
inplace : bool, default=False
Process the volume in place.
Returns
-------
output : tensor
Coefficient image.
"""
# This implementation is based on the file bsplines.c in SPM12, written
# by John Ashburner, which is itself based on the file coeff.c,
# written by Philippe Thevenaz: http://bigwww.epfl.ch/thevenaz/interpolation
# . DCT1 boundary conditions were derived by Thevenaz and Unser.
# . DFT boundary conditions were derived by John Ashburner.
# SPM12 is released under the GNU-GPL v2 license.
# Philippe Thevenaz's code does not have an explicit license as far
# as we know.
if backend.jitfields and jitfields.available:
return jitfields.spline_coeff(input, interpolation, bound,
dim, inplace)
out = SplineCoeff.apply(input, bound, interpolation, dim, inplace)
return out
def spline_coeff_nd(input, interpolation='linear', bound='dct2', dim=None,
inplace=False):
"""Compute the interpolating spline coefficients, for a given spline order
and boundary conditions, along the last `dim` dimensions.
Notes
-----
{interpolation}
{bound}
References
----------
{ref}
Parameters
----------
input : (..., *spatial) tensor
Input image.
interpolation : int or sequence[int], default=1
Interpolation order.
bound : BoundType or sequence[BoundType], default='dct1'
Boundary conditions.
dim : int, default=-1
Number of spatial dimensions
inplace : bool, default=False
Process the volume in place.
Returns
-------
output : (..., *spatial) tensor
Coefficient image.
"""
# This implementation is based on the file bsplines.c in SPM12, written
# by John Ashburner, which is itself based on the file coeff.c,
# written by Philippe Thevenaz: http://bigwww.epfl.ch/thevenaz/interpolation
# . DCT1 boundary conditions were derived by Thevenaz and Unser.
# . DFT boundary conditions were derived by John Ashburner.
# SPM12 is released under the GNU-GPL v2 license.
# Philippe Thevenaz's code does not have an explicit license as far
# as we know.
if backend.jitfields and jitfields.available:
return jitfields.spline_coeff_nd(input, interpolation, bound,
dim, inplace)
out = SplineCoeffND.apply(input, bound, interpolation, dim, inplace)
return out
grid_pull.__doc__ = grid_pull.__doc__.format(
interpolation=_doc_interpolation, bound=_doc_bound)
grid_push.__doc__ = grid_push.__doc__.format(
interpolation=_doc_interpolation, bound=_doc_bound)
grid_count.__doc__ = grid_count.__doc__.format(
interpolation=_doc_interpolation, bound=_doc_bound)
grid_grad.__doc__ = grid_grad.__doc__.format(
interpolation=_doc_interpolation, bound=_doc_bound)
spline_coeff.__doc__ = spline_coeff.__doc__.format(
interpolation=_doc_interpolation, bound=_doc_bound_coeff, ref=_ref_coeff)
spline_coeff_nd.__doc__ = spline_coeff_nd.__doc__.format(
interpolation=_doc_interpolation, bound=_doc_bound_coeff, ref=_ref_coeff)
# aliases
pull = grid_pull
push = grid_push
count = grid_count
def identity_grid(shape, dtype=None, device=None):
"""Returns an identity deformation field.
Parameters
----------
shape : (dim,) sequence of int
Spatial dimension of the field.
dtype : torch.dtype, default=`get_default_dtype()`
Data type.
device torch.device, optional
Device.
Returns
-------
grid : (*shape, dim) tensor
Transformation field
"""
mesh1d = [torch.arange(float(s), dtype=dtype, device=device)
for s in shape]
grid = torch.stack(meshgrid(mesh1d), dim=-1)
return grid
@torch.jit.script
def add_identity_grid_(disp):
"""Adds the identity grid to a displacement field, inplace.
Parameters
----------
disp : (..., *spatial, dim) tensor
Displacement field
Returns
-------
grid : (..., *spatial, dim) tensor
Transformation field
"""
dim = disp.shape[-1]
spatial = disp.shape[-dim-1:-1]
mesh1d = [torch.arange(s, dtype=disp.dtype, device=disp.device)
for s in spatial]
grid = meshgrid(mesh1d)
disp = movedim1(disp, -1, 0)
for i, grid1 in enumerate(grid):
disp[i].add_(grid1)
disp = movedim1(disp, 0, -1)
return disp
@torch.jit.script
def add_identity_grid(disp):
"""Adds the identity grid to a displacement field.
Parameters
----------
disp : (..., *spatial, dim) tensor
Displacement field
Returns
-------
grid : (..., *spatial, dim) tensor
Transformation field
"""
return add_identity_grid_(disp.clone())
def affine_grid(mat, shape):
"""Create a dense transformation grid from an affine matrix.
Parameters
----------
mat : (..., D[+1], D+1) tensor
Affine matrix (or matrices).
shape : (D,) sequence[int]
Shape of the grid, with length D.
Returns
-------
grid : (..., *shape, D) tensor
Dense transformation grid
"""
mat = torch.as_tensor(mat)
shape = list(shape)
nb_dim = mat.shape[-1] - 1
if nb_dim != len(shape):
raise ValueError('Dimension of the affine matrix ({}) and shape ({}) '
'are not the same.'.format(nb_dim, len(shape)))
if mat.shape[-2] not in (nb_dim, nb_dim+1):
raise ValueError('First argument should be matrces of shape '
'(..., {0}, {1}) or (..., {1], {1}) but got {2}.'
.format(nb_dim, nb_dim+1, mat.shape))
batch_shape = mat.shape[:-2]
grid = identity_grid(shape, mat.dtype, mat.device)
if batch_shape:
for _ in range(len(batch_shape)):
grid = grid.unsqueeze(0)
for _ in range(nb_dim):
mat = mat.unsqueeze(-1)
lin = mat[..., :nb_dim, :nb_dim]
off = mat[..., :nb_dim, -1]
grid = matvec(lin, grid) + off
return grid