import nibabel as nib
import os.path as op
import os
import numpy as np
import logging
import pimms
from AFQ.tasks.decorators import as_file
from AFQ.tasks.utils import with_name, str_to_desc
import AFQ.data.fetch as afd
from AFQ.utils.path import drop_extension, write_json
from AFQ.definitions.mapping import SynMap
from AFQ.definitions.utils import Definition
from AFQ.definitions.image import ImageDefinition
from dipy.io.streamline import load_tractogram
from dipy.io.stateful_tractogram import Space
[docs]logger = logging.getLogger('AFQ')
@pimms.calc("b0_warped")
@as_file('_space-template_desc-b0_dwi.nii.gz')
[docs]def export_registered_b0(data_imap, mapping):
"""
full path to a nifti file containing
b0 transformed to template space
"""
mean_b0 = nib.load(data_imap["b0"]).get_fdata()
warped_b0 = mapping.transform(mean_b0)
warped_b0 = nib.Nifti1Image(warped_b0, data_imap["reg_template"].affine)
return warped_b0, dict(b0InSubject=data_imap["b0"])
@pimms.calc("template_xform")
@as_file('_space-subject_desc-template_dwi.nii.gz')
@pimms.calc("rois")
[docs]def export_rois(base_fname, output_dir, data_imap, mapping):
"""
dictionary of full paths to Nifti1Image files of ROIs
transformed to subject space
"""
bundle_dict = data_imap["bundle_dict"]
rois_dir = op.join(output_dir, 'ROIs')
os.makedirs(rois_dir, exist_ok=True)
roi_files = {}
base_roi_fname = op.join(rois_dir, op.split(base_fname)[1])
for bundle_name in bundle_dict:
roi_files[bundle_name] = []
for roi_fname in bundle_dict.transform_rois(
bundle_name, mapping, data_imap["dwi_affine"],
base_fname=base_roi_fname):
logger.info(f"Saving {roi_fname}")
roi_files[bundle_name].append(roi_fname)
meta = {}
meta_fname = f'{drop_extension(roi_fname)}.json'
write_json(meta_fname, meta)
return {'rois': roi_files}
@pimms.calc("mapping")
[docs]def mapping(base_fname, dwi_data_file, reg_subject, data_imap,
mapping_definition=None):
"""
mapping from subject to template space.
Parameters
----------
mapping_definition : instance of `AFQ.definitions.mapping`, optional
This defines how to either create a mapping from
each subject space to template space or load a mapping from
another software. If creating a map, will register reg_subject and
reg_template.
If None, use SynMap()
Default: None
"""
reg_template = data_imap["reg_template"]
if mapping_definition is None:
mapping_definition = SynMap()
if not isinstance(mapping_definition, Definition):
raise TypeError(
"mapping must be a mapping defined"
+ " in `AFQ.definitions.mapping`")
return mapping_definition.get_for_subses(
base_fname, data_imap["dwi"], dwi_data_file,
reg_subject, reg_template)
@pimms.calc("mapping")
[docs]def sls_mapping(base_fname, dwi_data_file, reg_subject, data_imap,
tractography_imap, mapping_definition=None):
"""
mapping from subject to template space.
Parameters
----------
mapping_definition : instance of `AFQ.definitions.mapping`, optional
This defines how to either create a mapping from
each subject space to template space or load a mapping from
another software. If creating a map, will register reg_subject and
reg_template.
If None, use SynMap()
Default: None
"""
reg_template = data_imap["reg_template"]
if mapping_definition is None:
mapping_definition = SynMap()
if not isinstance(mapping_definition, Definition):
raise TypeError(
"mapping must be a mapping defined"
+ " in `AFQ.definitions.mapping`")
streamlines_file = tractography_imap["streamlines"]
tg = load_tractogram(
streamlines_file, reg_subject,
Space.VOX, bbox_valid_check=False)
tg.to_rasmm()
atlas_fname = op.join(
afd.afq_home,
'hcp_atlas_16_bundles',
'Atlas_in_MNI_Space_16_bundles',
'whole_brain',
'whole_brain_MNI.trk')
if not op.exists(atlas_fname):
afd.fetch_hcp_atlas_16_bundles()
hcp_atlas = load_tractogram(
atlas_fname,
'same', bbox_valid_check=False)
return mapping_definition.get_for_subses(
base_fname, data_imap["dwi"],
dwi_data_file,
reg_subject, reg_template,
subject_sls=tg.streamlines,
template_sls=hcp_atlas.streamlines)
@pimms.calc("reg_subject")
[docs]def get_reg_subject(data_imap,
reg_subject_spec="power_map"):
"""
Nifti1Image which represents this subject
when registering the subject to the template
Parameters
----------
reg_subject_spec : str, instance of `AFQ.definitions.ImageDefinition`, optional # noqa
The source image data to be registered.
Can either be a Nifti1Image, an ImageFile, or str.
if "b0", "dti_fa_subject", "subject_sls", or "power_map,"
image data will be loaded automatically.
If "subject_sls" is used, slr registration will be used
and reg_template should be "hcp_atlas".
Default: "power_map"
"""
if not isinstance(reg_subject_spec, str)\
and not isinstance(reg_subject_spec, nib.Nifti1Image):
# Note the ImageDefinition case is handled in get_mapping_plan
raise TypeError(
"reg_subject must be a str, ImageDefinition, or Nifti1Image")
filename_dict = {
"b0": "b0",
"power_map": "csd_pmap",
"dti_fa_subject": "dti_fa",
"subject_sls": "b0",
}
bm = nib.load(data_imap["brain_mask"])
if reg_subject_spec in filename_dict:
reg_subject_spec = data_imap[filename_dict[reg_subject_spec]]
if isinstance(reg_subject_spec, str):
img = nib.load(reg_subject_spec)
bm = bm.get_fdata().astype(bool)
masked_data = img.get_fdata()
masked_data[~bm] = 0
img = nib.Nifti1Image(masked_data, img.affine)
return img
[docs]def get_mapping_plan(kwargs, use_sls=False):
mapping_tasks = with_name([
export_registered_b0, template_xform, export_rois, mapping,
get_reg_subject])
# add custom scalars
for scalar in kwargs["scalars"]:
if isinstance(scalar, Definition):
mapping_tasks[f"{scalar.get_name()}_res"] =\
pimms.calc(f"{scalar.get_name()}")(
as_file((
f'_desc-{str_to_desc(scalar.get_name())}'
'_dwi.nii.gz'))(
scalar.get_image_getter("mapping")))
if use_sls:
mapping_tasks["mapping_res"] = sls_mapping
reg_ss = kwargs.get("reg_subject_spec", None)
if isinstance(reg_ss, ImageDefinition):
del kwargs["reg_subject_spec"]
mapping_tasks["reg_subject_spec_res"] = pimms.calc("reg_subject_spec")(
as_file((
f'_desc-{str_to_desc(reg_ss.get_name())}'
'_dwi.nii.gz'))(reg_ss.get_image_getter("mapping")))
return pimms.plan(**mapping_tasks)
