Commit c68da9a3 authored by Paul McCarthy's avatar Paul McCarthy 🚵
Browse files

ENH: fileOrImage works with fsl.data.images.

parent 523b1f6e
......@@ -635,13 +635,25 @@ class _FileOrThing(object):
def fileOrImage(*imgargs):
"""Decorator which can be used to ensure that any NIfTI images are saved
to file, and output images can be loaded and returned as ``nibabel``
image objects.
image objects or :class:`.Image` objects.
"""
# keep track of the input argument
# types on each call, so we know
# whether to return a fsl.Image or
# a nibabel image
intypes = []
def prepIn(workdir, name, val):
infile = None
if isinstance(val, (fslimage.Image, nib.nifti1.Nifti1Image)):
intypes.append(type(val))
if isinstance(val, fslimage.Image):
val = val.nibImage
if isinstance(val, nib.nifti1.Nifti1Image):
infile = val.get_filename()
......@@ -661,13 +673,28 @@ def fileOrImage(*imgargs):
# create an independent in-memory
# copy of the image file
img = nib.load(path)
return nib.nifti1.Nifti1Image(img.get_data(), None, img.header)
# if any arguments were fsl images,
# that takes precedence.
if fslimage.Image in intypes:
return fslimage.Image(img.get_data(), header=img.header)
# but if all inputs were file names,
# nibabel takes precedence
elif nib.nifti1.Nifti1Image in intypes or len(intypes) == 0:
return nib.nifti1.Nifti1Image(img.get_data(), None, img.header)
# this function should not be called
# under any other circumstances
else:
raise RuntimeError('Cannot handle type: {}'.format(intypes))
def decorator(func):
fot = _FileOrThing(func, prepIn, prepOut, load, *imgargs)
def wrapper(*args, **kwargs):
return fot(*args, **kwargs)
result = fot(*args, **kwargs)
intypes[:] = []
return result
return _update_wrapper(wrapper, func)
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment