import itertools
import multiprocessing as mp
import os
import platform
import shutil
import subprocess
import uuid
import warnings
from functools import partial
from typing import Dict, Sequence, Union
from nipype.interfaces.elastix import ApplyWarp, Registration
from nipype.interfaces.elastix.registration import RegistrationOutputSpec
from tqdm import tqdm
from tqdm.contrib.concurrent import process_map
from dosma import file_constants as fc
from dosma.data_io.med_volume import MedicalVolume
from dosma.data_io.nifti_io import NiftiReader, NiftiWriter
__all__ = ["register", "apply_warp", "symlink_elastix", "unlink_elastix"]
MedVolOrPath = Union[MedicalVolume, str]
[docs]def register(
target: MedVolOrPath,
moving: Union[MedVolOrPath, Sequence[MedVolOrPath]],
parameters: Union[str, Sequence[str]],
output_path: str,
target_mask: MedVolOrPath = None,
moving_masks: Union[MedVolOrPath, Sequence[MedVolOrPath]] = None,
sequential: bool = False,
collate: bool = True,
num_workers: int = 0,
num_threads: int = 1,
show_pbar: bool = False,
return_volumes: bool = False,
rtype: type = dict,
**kwargs,
):
"""Register moving image(s) to the target.
`MedVolOrPath` is a shorthand for `MedicalVolume` or `str`. It indicates the argument
can be either a `MedicalVolume` or a `str` path to a nifti file.
Args:
target (`MedicalVolume` or `str`): The target/fixed image.
moving (`MedicalVolume`(s) or `str`(s)): The moving/source image(s).
parameters (`str(s)`): Elastix parameter files to use.
output_path (`str`): Output directory to store files.
target_mask (`MedicalVolume` or `str`, optional): The target/fixed mask.
moving_masks (`MedicalVolume`(s) or `str`(s), optional): The moving mask(s).
If only one specified, the mask will be used for all moving images.
sequential (bool, optional): If `True`, apply parameter files sequentially.
collate (bool, optional): If `True`, will collate outputs from sequential registration
into single RegistrationOutputSpec instance. If `sequential=False`, this argument
is ignored.
num_workers (int, optional): Number of workers to use for reading/writing data and
for registration.
num_threads (int, optional): Number of threads to use for registration.
Note total number of threads used will be ``num_workers * num_threads``.
show_pbar (bool, optional): If `True`, show progress bar during registration.
Note the progress bar will not be shown for intermediate reading/writing.
return_volumes (bool, optional): If `True`, registered volumes will also be returned.
By default, only the output namespaces (RegistrationOutputSpec) of the registrations are
returned.
rtype (type, optional): The return type. Either `dict` or `tuple`.
kwargs: Keyword arguments used to initialize `nipype.interfaces.elastix.Registration`.
Returns:
Dict or Tuple: Type specified by `rtype`. If dict, with keys 'outputs'
(registration outputs) and 'volumes' (final volumes)
if `return_volumes=True`). If tuple, order is (`outputs`, `volumes` or `None`).
Length of `outputs` and `volumes` depends on number of images specified in `moving`.
outputs (Sequence[RegistrationOutputSpec]): The output objects from
elastix registration, one for each moving image. Each object is effectively
a namespace with four main attributes:
- 'transform' (List[str]): Paths to transform files produced using registration.
- 'warped_file' (str): Path to the final registered image.
- 'warped_files' (List[str]): Paths to all intermediate images created if
multiple parameter files used.
volumes (Sequence[MedicalVolume]): Registered volumes.
"""
assert issubclass(rtype, (Dict, Sequence)) # `rtype` must be dict or tuple
has_output_path = bool(output_path)
if not output_path:
output_path = os.path.join(fc.TEMP_FOLDER_PATH, "register")
moving = [moving] if isinstance(moving, (MedicalVolume, str)) else moving
moving_masks = (
[moving_masks]
if moving_masks is None or isinstance(moving_masks, (MedicalVolume, str))
else moving_masks
)
if len(moving_masks) > 1 and len(moving) != len(moving_masks):
raise ValueError(
"Got {} moving images but {} moving masks".format(len(moving), len(moving_masks))
)
files = [target, target_mask] + moving + moving_masks
# Write medical volumes (if any) to nifti file for use with elastix.
tmp_dir = os.path.join(output_path, "tmp")
default_files = (
["target", "target-mask"]
+ [f"moving-{idx}" for idx in range(len(moving))]
+ [f"moving-mask-{idx}" for idx in range(len(moving_masks))]
) # noqa
assert len(default_files) == len(files), default_files # should be 1-to-1 with # args provided
vols = [(idx, v) for idx, v in enumerate(files) if isinstance(v, MedicalVolume)]
idxs, vols = [x[0] for x in vols], [x[1] for x in vols]
if len(vols) > 0:
filepaths = [os.path.join(tmp_dir, f"{default_files[idx]}.nii.gz") for idx in idxs]
if num_workers > 0:
with mp.Pool(min(num_workers, len(vols))) as p:
out = p.starmap_async(_write, zip(vols, filepaths))
out.wait()
else:
for vol, fp in zip(vols, filepaths):
_write(vol, fp)
for idx, fp in zip(idxs, filepaths):
files[idx] = fp
# Assign file paths to respective variables.
target, moving = files[0], files[2 : 2 + len(moving)]
target_mask, moving_masks = files[1], files[2 + len(moving) :]
if len(moving_masks) == 1:
moving_masks = moving_masks * len(moving)
all_outputs = {}
# Perform registration.
reg_out_paths = [os.path.join(output_path, f"moving-{idx}") for idx in range(len(moving))]
reg_args = list(zip(moving, moving_masks, reg_out_paths))
if num_workers > 0:
func = partial(
_elastix_register_mp,
target=target,
parameters=parameters,
target_mask=target_mask,
sequential=sequential,
collate=collate,
num_threads=num_threads,
**kwargs,
)
max_workers = min(num_workers, len(reg_args))
out = process_map(
func, reg_args, max_workers=max_workers, tqdm_class=tqdm, disable=not show_pbar
)
else:
out = []
for mvg, mvg_mask, out_path in tqdm(reg_args, disable=not show_pbar):
_out = _elastix_register(
target,
mvg,
parameters,
out_path,
target_mask,
mvg_mask,
sequential,
collate,
num_threads,
**kwargs,
)
out.append(_out)
all_outputs["outputs"] = tuple(out)
# Load volumes.
if return_volumes:
filepaths = [x[-1].warped_file if isinstance(x, Sequence) else x.warped_file for x in out]
if num_workers > 0:
with mp.Pool(min(num_workers, len(filepaths))) as p:
vols = p.map(_read, filepaths)
else:
vols = []
for fp in filepaths:
vols.append(_read(fp))
all_outputs["volume"] = tuple(vols)
# Clean up.
for _dir in [tmp_dir, output_path if not has_output_path else None]:
if not _dir or not os.path.isdir(_dir):
continue
shutil.rmtree(_dir)
if issubclass(rtype, dict):
out = rtype(all_outputs)
elif issubclass(rtype, Sequence):
out = rtype([all_outputs["outputs"], all_outputs.get("volume", None)])
else:
assert False # Should have type checking earlier.
return out
[docs]def apply_warp(
moving: MedVolOrPath,
transform: Union[str, Sequence[str]] = None,
out_registration: RegistrationOutputSpec = None,
output_path: str = None,
rtype: type = MedicalVolume,
num_threads: int = 1,
show_pbar: bool = False,
) -> MedVolOrPath:
"""Apply transform(s) to moving image using transformix.
Use transformix to apply a transform on an input image. The transform(s) is/are
specified in the transform-parameter file(s).
Args:
moving (MedicalVolume(s) or str(s)): The moving/source image to transform.
transform (str(s)): Paths to transform files to be used by transformix.
If multiple files provided, transforms will be applied sequentially.
If `None`, will be determined by `out_registration.transform`.
out_registration (RegistrationOutputSpec(s)): Outputs from elastix registration
using nipype. Must be specified if `transform` is None.
output_path (str): Output directory to store files.
rtype (type, optional): Return type - either `MedicalVolume` or `str`.
If `str`, `output_path` must be specified. Defaults to `MedicalVolume`.
num_threads (int, optional): Number of threads to use for registration.
If `None`, defaults to 1.
show_pbar (bool, optional): If `True`, show progress bar when applying transforms.
Return:
MedVolOrPath: The medical volume or nifti file corresponding to the volume.
See `rtype` for details.
"""
assert rtype in [MedicalVolume, str], rtype # rtype must be MedicalVolume or str
has_output_path = bool(output_path)
if rtype == str and not has_output_path:
raise ValueError("`output_path` must be specified when `rtype=str`")
if not output_path:
output_path = os.path.join(fc.TEMP_FOLDER_PATH, f"apply_warp-{str(uuid.uuid1())}")
output_path = os.path.abspath(output_path)
os.makedirs(output_path, exist_ok=True)
if not transform:
transform = out_registration.transform
elif isinstance(transform, str):
transform = [transform]
transform = [os.path.abspath(t) for t in transform]
mv_filepath = os.path.join(output_path, "moving.nii.gz")
if isinstance(moving, MedicalVolume):
NiftiWriter().save(moving, mv_filepath)
moving = mv_filepath
transformix_path = _local_exe("transformix") # noqa
cwd = _local_lib_dir()
for tf in tqdm(transform, disable=not show_pbar):
reg = ApplyWarp()
reg.inputs.moving_image = moving
reg.inputs.transform_file = tf
reg.inputs.output_path = output_path
reg.terminal_output = fc.NIPYPE_LOGGING
reg.inputs.num_threads = num_threads
reg_output = reg.run(cwd=cwd)
moving = reg_output.outputs.warped_file
if rtype == MedicalVolume:
out = NiftiReader().load(moving)
else:
out = moving
if os.path.isfile(mv_filepath):
os.remove(mv_filepath)
if not has_output_path:
shutil.rmtree(output_path)
return out
[docs]def symlink_elastix(path: str = None, lib_only: bool = True, force: bool = False):
"""Symlinks elastix/transformix files to the dosma library.
Args:
path (str, optional): Path to elastix folder. This folder should
contain two folders `bin` and `lib`. If `None`, determined
using `which elastix`. This will overwrite existing linked
files. path cannot be automatically determined on Windows.
lib_only (bool, optional): If `True`, only links contents of `lib`
folder.
force (bool, optional): If `True`, unlinks existing files before relinking.
Note this operation is not atomic.
Note:
Setting elastix paths this way is not recommended unless you
are using a MacOS (Darwin) platform, where there are known
path issues with elastix (https://github.com/almarklein/pyelastix/issues/9).
For linux and windows machines, using the setup described in the elastix
guide is sufficient.
"""
system = platform.system().lower()
assert system in ["windows", "darwin", "linux"]
if system != "darwin":
warnings.warn(
f"Symlinking elastix/transformix paths not recommended for {system} " f"machines"
)
if path is None:
if system == "windows":
raise ValueError("`path` cannot be determined automatically on Windows")
try:
out = subprocess.check_output(["which", "elastix"]).decode("ascii").strip("\n")
path = os.path.dirname(os.path.dirname(out))
except subprocess.CalledProcessError:
raise ValueError(
"Path to `elastix` not intialized. "
"Use `export PATH=/path/to/elastix/folder:$PATH`"
)
assert os.path.isdir(path), path # must be a directory
dirs = {"lib": [x for x in os.listdir(os.path.join(path, "lib")) if x.startswith("libANNlib")]}
if not lib_only:
dirs["bin"] = ["elastix", "transformix"]
for dirname, files in dirs.items():
for file in files:
src = os.path.join(path, dirname, file)
tgt = os.path.join(fc._DOSMA_ELASTIX_FOLDER, file)
if os.path.exists(tgt):
if force:
os.remove(tgt)
else:
raise FileExistsError(
f"File {tgt} exists. "
f"Use `unlink_elastix` or `force` to unlink the file."
)
os.symlink(src, tgt)
[docs]def unlink_elastix():
"""Unlinks all elastix/transformix files in the dosma library."""
for x in os.listdir(fc._DOSMA_ELASTIX_FOLDER):
x = os.path.join(fc._DOSMA_ELASTIX_FOLDER, x)
if os.path.islink(x):
os.remove(x)
def _elastix_register(
target: str,
moving: str,
parameters: Sequence[str],
output_path: str,
target_mask: str = None,
moving_mask: str = None,
sequential=False,
collate=True,
num_threads=None,
use_mask: Sequence[bool] = None,
**kwargs,
):
def _register(_moving, _parameters, _output_path, _use_mask=None):
if isinstance(_parameters, str):
_parameters = [_parameters]
if _use_mask is None:
_use_mask = target_mask is not None or moving_mask is not None
_output_path = os.path.abspath(_output_path)
os.makedirs(_output_path, exist_ok=True)
elastix_path = _local_exe("elastix")
cwd = _local_lib_dir()
reg = Registration()
if elastix_path:
reg._cmd = elastix_path
reg.inputs.fixed_image = os.path.abspath(target)
reg.inputs.moving_image = os.path.abspath(_moving)
reg.inputs.parameters = [os.path.abspath(p) for p in _parameters]
reg.inputs.output_path = os.path.abspath(_output_path)
reg.terminal_output = fc.NIPYPE_LOGGING
if num_threads:
reg.inputs.num_threads = num_threads
if _use_mask and target_mask is not None:
reg.inputs.fixed_mask = os.path.abspath(target_mask)
if _use_mask and moving_mask is not None:
reg.inputs.target_mask = os.path.abspath(moving_mask)
for k, v in kwargs.items():
setattr(reg.inputs, k, v)
return reg.run(cwd=cwd).outputs
def _collate_outputs(_outs):
"""
Concatenates fields that are sequential and takes final output
for fields that are not.
"""
if len(_outs) == 1:
return _outs[0]
_result = _outs[0]
fields = list(_outs[0].__dict__.keys())
for _fld in fields:
_res_val = getattr(_result, _fld)
if not isinstance(_res_val, str) and isinstance(_res_val, Sequence):
val = list(itertools.chain.from_iterable([getattr(x, _fld) for x in _outs]))
else:
val = getattr(_outs[-1], _fld)
setattr(_result, _fld, val)
return _result
if use_mask is not None:
assert sequential # use_mask can only be specified when sequential is specified
if sequential:
outs, mvg = [], moving
for idx, param in enumerate(parameters):
_use_mask = None if use_mask is None else use_mask[idx]
_out = _register(mvg, param, os.path.join(output_path, f"param{idx}"), _use_mask)
outs.append(_out)
mvg = _out.warped_file
out = _collate_outputs(outs) if collate else outs
return out
else:
return _register(moving, parameters, output_path)
def _elastix_register_mp(args, **kwargs):
"""Reorder arguments for multiprocessing support."""
moving, moving_mask, output_path = args
return _elastix_register(
moving=moving, moving_mask=moving_mask, output_path=output_path, **kwargs
)
def _write(vol: MedicalVolume, path: str):
"""Extracted out for multiprocessing purposes."""
NiftiWriter().save(vol, path)
def _read(path: str):
return NiftiReader().load(path)
def _local_exe(exe):
"""Returns path to local executable if exists, else None."""
assert exe in ["elastix", "transformix"]
dosma_path = os.path.join(fc._DOSMA_ELASTIX_FOLDER, exe)
if os.path.isfile(dosma_path):
return os.path.abspath(dosma_path)
def _local_lib_dir():
"""Returns path to directory with local lib file if exists, else None."""
files = [x for x in os.listdir(fc._DOSMA_ELASTIX_FOLDER) if x.startswith("libANNlib")]
if len(files) > 0:
return fc._DOSMA_ELASTIX_FOLDER