Skip to content

Commit

Permalink
Merge pull request #22 from yhuang43/master
Browse files Browse the repository at this point in the history 
reconstruct FrameImage, Affine, DeformField
  • Loading branch information
@jnolan14
jnolan14 authored Feb 14, 2024
2 parents 026cabe + 1622283 commit 19f3388
Show file tree
Hide file tree
Showing 4 changed files with 189 additions and 140 deletions.
135 changes: 29 additions & 106 deletions surfa/image/framed.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -383,18 +383,21 @@ def resample_like(self, target, method='linear', copy=True, fill=0):
method=method, affine=affine.matrix, fill=fill)
return self.new(interped, target_geom)

def transform(self, trf=None, method='linear', rotation='corner', resample=True, fill=0, affine=None):
def transform(self, trf=None, method='linear', rotation='corner', resample=True, fill=0):
"""
Apply an affine or non-linear transform.
**Note on deformation fields:** Until we come up with a reasonable way to represent
deformation fields, they can be implemented as multi-frame images. It is assumed that
they represent a *displacement* vector field in voxel space. So under the hood, images
will be moved into the space of the deformation field if the image geometries differ.
**The original implementation has been moved to Affine.transform and Warp.transform.
**The method is now impemeted to transform image using Affine.transform and Warp.transform.
**Note on trf argument:** It accepts Affine/Warp object, or deformation fields (4D numpy array).
Pass trf argument as a numpy array is deprecated and will be removed in the future.
It is assumed that the deformation fields represent a *displacement* vector field in voxel space.
So under the hood, images will be moved into the space of the deformation field if the image geometries differ.
Parameters
----------
trf : Affine or !class
trf : Affine/Warp or !class
Affine transform or nonlinear deformation (displacement) to apply to the image.
method : {'linear', 'nearest'}
Image interpolation method if resample is enabled.
Expand All @@ -406,8 +409,6 @@ def transform(self, trf=None, method='linear', rotation='corner', resample=True,
be updated (this is not possible if a displacement field is provided).
fill : scalar
Fill value for out-of-bounds voxels.
affine : Affine
Deprecated. Use the `trf` argument instead.
Returns
-------
Expand All @@ -418,107 +419,29 @@ def transform(self, trf=None, method='linear', rotation='corner', resample=True,
raise NotImplementedError('transform() is not yet implemented for 2D data, '
'contact andrew if you need this')

if affine is not None:
trf = affine
warnings.warn('The \'affine\' argument to transform() is deprecated. Just use '
'the first positional argument to specify a transform.',
image = self.copy()
transformer = trf
if isinstance(transformer, Affine):
return transformer.transform(image, method, rotation, resample, fill)

from surfa.transform.warp import Warp
if isinstance(transformer, np.ndarray):
warnings.warn('The option to pass \'trf\' argument as a numpy array is deprecated. '
'Pass \'trf\' as either an Affine or Warp object',
DeprecationWarning, stacklevel=2)

# one of these two will be set by the end of the function
disp_data = None
matrix_data = None

# first try to convert it to an affine matrix. if that fails
# we assume it has to be a deformation field
try:
trf = cast_affine(trf, allow_none=False)
except ValueError:
pass

if isinstance(trf, Affine):

# for clarity
affine = trf

# if not resampling, just change the image vox2world matrix and return
if not resample:

# TODO: if affine is missing geometry info, do we assume that the affine
# is in world space or voxel space? let's do world for now
if affine.source is not None and affine.target is not None:
affine = affine.convert(space='world', source=self)
# TODO: must try this again once I changed everything around!!
elif affine.space is None:
warnings.warn('Affine transform is missing metadata defining its coordinate '
'space or source and target geometry. Assuming matrix is a '
'world-space transform since resample=False, but this might '
'not always be the case. Best practice is to provide the '
'correct metadata in the affine')
elif affine.space != 'world':
raise ValueError('affine must contain source and target info '
'if not in world space')

# apply forward transform to the header
transformed = self.copy()
transformed.geom.update(vox2world=affine @ affine.source.vox2world)
return transformed

# sanity check and preprocess the affine if resampling
target_geom = self.geom

if affine.source is not None and affine.target is not None:
# it should be assumed that the default affine space is voxel
# when both source and target are set
if affine.space is None:
affine = affine.copy()
affine.space = 'voxel'
#
affine = affine.convert(space='voxel', source=self)
target_geom = affine.target
elif affine.space is not None and affine.space != 'voxel':
raise ValueError('affine must contain source and target info if '
'coordinate space is not \'voxel\'')

# ensure the rotation is around the image corner before interpolating
if rotation not in ('center', 'corner'):
raise ValueError("rotation must be 'center' or 'corner'")
elif rotation == 'center':
affine = center_to_corner_rotation(affine, self.baseshape)

# make sure the matrix is actually inverted since we want a target to
# source voxel mapping for resampling
matrix_data = affine.inv().matrix
source_data = self.framed_data

else:
if not resample:
raise ValueError('transform resampling must be enabled when deformation is used')

# cast deformation as a framed image data. important that the fallback geometry
# here is the current image space
deformation = cast_image(trf, fallback_geom=self.geom)
if deformation.nframes != self.basedim:
raise ValueError(f'deformation ({deformation.nframes}D) does not match '
f'dimensionality of image ({self.basedim}D)')

# since we only support deformations in the form of voxel displacement
# currently, must get the image in the space of the deformation
source_data = self.resample_like(deformation).framed_data

# make sure to use the deformation as the target geometry
target_geom = deformation.geom

# get displacement data
disp_data = deformation.data

# do the interpolation
interpolated = interpolate(source=source_data,
target_shape=target_geom.shape,
method=method,
affine=matrix_data,
disp=disp_data,
fill=fill)
return self.new(interpolated, target_geom)
image = image.resample_like(deformation)
transformer = Warp(data=trf,
source=image.geom,
target=deformation.geom,
format=Warp.Format.disp_crs)

if not isinstance(transformer, Warp):
raise ValueError("Pass \'trf\' as either an Affine or Warp object")

return transformer.transform(image, method, fill)


def reorient(self, orientation, copy=True):
"""
Expand Down
2 changes: 1 addition & 1 deletion surfa/transform/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,4 +14,4 @@
from .geometry import image_geometry2volgeom_dict
from .geometry import volgeom_dict2image_geometry

from .deformfield import DeformField
from .warp import Warp
148 changes: 126 additions & 22 deletions surfa/transform/affine.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -194,48 +194,70 @@ def __matmul__(self, other):
matrix = np.matmul(self.matrix, other.matrix)
return Affine(matrix)

def __call__(self, points):
def __call__(self, *args, **kwargs):
"""
Apply the affine transform matrix to a set of points. Calls `self.transform(points)`
under the hood.
Apply the affine transform matrix to a set of points, or an image Volume.
Calls `self.transform()` under the hood.
"""
return self.transform(points)
return self.transform(*args, **kwargs)

def transform(self, points):
def transform(self, data, method='linear', rotation='corner', resample=True, fill=0, points=None):
"""
Apply the affine transform matrix to an N-D point or set of points.
Apply the affine transform matrix to the input data.
Parameters
----------
points : (..., N) float
N-D point values to transform.
data : input data to transform
N-D point values, or image Volume
method : {'linear', 'nearest'}
Image interpolation method if resample is enabled.
rotation : {'corner', 'center'}
Apply affine with rotation axis at the image corner or center.
resample : bool
If enabled, voxel data will be interpolated and resampled, and geometry will be set
the target. If disabled, voxel data will not be modified, and only the geometry will
be updated (this is not possible if a displacement field is provided).
fill : scalar
Fill value for out-of-bounds voxels.
points : N-D point values
Deprecated. Use the `data` argument instead.
Returns
-------
(..., N) float
Transformed N-D point array.
"""
# a common mistake is to use this function for transforming an image
# or mesh, so run this check to help the user out a bit
if ismesh(points):
Transformed N-D point array if (input data is N-D point)
transformed : Volume
Transformed image if (input data is an image Volume)
"""
if points is not None:
data = points
warnings.warn('The \'points\' argument to transform() is deprecated. Just use '
'the first positional argument to specify set of points or an image to transform.',
DeprecationWarning, stacklevel=2)

# a common mistake is to use this function for transforming a mesh,
# so run this check to help the user out a bit
if ismesh(data):
raise ValueError('use mesh.transform(affine) to apply an affine to a mesh')
if isimage(points):
raise ValueError('use image.transform(affine) to apply an affine to an image')

if isimage(data):
return self.__transform_image(data, method, rotation, resample, fill)

# convert to array
points = np.ascontiguousarray(points)
data = np.ascontiguousarray(data)

# check correct dimensionality
if points.shape[-1] != self.ndim:
if data.shape[-1] != self.ndim:
raise ValueError(f'transform() method expected {self.ndim}D points, but got '
f'{points.shape[-1]}D input with shape {points.shape}')
f'{data.shape[-1]}D input with shape {data.shape}')

# account for multiple possible input axes and be sure to
# always return the same shape
shape = points.shape
points = points.reshape(-1, self.ndim)
points = np.c_[points, np.ones(points.shape[0])].T
moved = np.dot(self.matrix, points).T[:, :-1]
shape = data.shape
data = data.reshape(-1, self.ndim)
data = np.c_[data, np.ones(data.shape[0])].T
moved = np.dot(self.matrix, data).T[:, :-1]
return np.ascontiguousarray(moved).reshape(shape)

def inv(self):
Expand Down Expand Up @@ -365,6 +387,88 @@ def convert(self, source=None, target=None, space=None, copy=True):
return Affine(affine.matrix, source=source, target=target, space=space)


# the implementation is based on FramedImage.transform
def __transform_image(self, image, method='linear', rotation='corner', resample=True, fill=0):
"""
Apply the affine transform matrix to an image.
Parameters
----------
image : Volume
input image Volume
Returns
-------
transformed : Volume
transformed image
"""

if image.basedim == 2:
raise NotImplementedError('Affine.transform() is not yet implemented for 2D data')

affine = self.copy()

# if not resampling, just change the image vox2world matrix and return
if not resample:

# TODO: if affine is missing geometry info, do we assume that the affine
# is in world space or voxel space? let's do world for now
if affine.source is not None and affine.target is not None:
affine = affine.convert(space='world', source=image)
# TODO: must try this again once I changed everything around!!
elif affine.space is None:
warnings.warn('Affine transform is missing metadata defining its coordinate '
'space or source and target geometry. Assuming matrix is a '
'world-space transform since resample=False, but this might '
'not always be the case. Best practice is to provide the '
'correct metadata in the affine')
elif affine.space != 'world':
raise ValueError('affine must contain source and target info '
'if not in world space')

# apply forward transform to the header
transformed = image.copy()
transformed.geom.update(vox2world=affine @ affine.source.vox2world)
return transformed

# sanity check and preprocess the affine if resampling
target_geom = image.geom

if affine.source is not None and affine.target is not None:
# it should be assumed that the default affine space is voxel
# when both source and target are set
if affine.space is None:
affine = affine.copy()
affine.space = 'voxel'
#
affine = affine.convert(space='voxel', source=image)
target_geom = affine.target
elif affine.space is not None and affine.space != 'voxel':
raise ValueError('affine must contain source and target info if '
'coordinate space is not \'voxel\'')

# ensure the rotation is around the image corner before interpolating
if rotation not in ('center', 'corner'):
raise ValueError("rotation must be 'center' or 'corner'")
elif rotation == 'center':
affine = center_to_corner_rotation(affine, image.baseshape)

# make sure the matrix is actually inverted since we want a target to
# source voxel mapping for resampling
matrix_data = affine.inv().matrix
source_data = image.framed_data

# do the interpolation
from surfa.image.interp import interpolate
interpolated = interpolate(source=source_data,
target_shape=target_geom.shape,
method=method,
affine=matrix_data,
fill=fill)
return image.new(interpolated, target_geom)



def affine_equal(a, b, matrix_only=False, tol=0.0):
"""
Test whether two affine transforms are equivalent.
Expand Down
Loading

0 comments on commit 19f3388

Please sign in to comment.