import logging
import nibabel as nib
import numpy as np
import scipy.sparse as sp
import talon
import talon.cli.utils
DESCRIPTION = """
Transform a tractogram into the `indices` and `weights` matrices that are used
in the definition of the linear operator used by TALON.
"""
def add_parser(subparsers):
parser = subparsers.add_parser(
'voxelize',
description=DESCRIPTION,
help='Voxelize a tractogram using TALON.')
parser.add_argument(
'in_tracks',
type=str,
help='Tractogram file to be voxelized in RAS+ and mm space. '
'The streamline coordinate (0,0,0) refers to the center of the '
'voxel. '
'Must be in NiBabel-readable format (.trk or .tck).'
)
parser.add_argument(
'in_img',
type=str,
help='Image serving as space reference. '
'Must be in NiBabel-readable format (.nii or .nii.gz).'
)
parser.add_argument(
'out_ind',
type=str,
help='Path where the indices will be saved in .npz format.'
)
parser.add_argument(
'out_wei',
type=str,
help='Path where the weights will be saved in .npz format.'
)
talon.cli.utils.add_ndir_to_input(parser)
talon.cli.utils.add_verbosity_and_force_to_parser(parser)
parser.set_defaults(func=run)
[docs]def run(in_tracks, in_img, out_ind, out_wei, force, ndir, **kwargs):
"""
This function reads tractogram files and writes the corresponding indices
and weights files.
Args:
in_tracks: str
Tractogram file to be voxelized in RAS+ and mm space. The
streamline coordinate (0,0,0) refers to the center of the voxel.
Must be in NiBabel-readable format (.trk or .tck).
in_img: str
Image serving as space reference. Must be in NiBabel-readable
format (.nii or .nii.gz).
out_ind: str
Path where the indices will be saved in .npz format.
out_wei: str
Path where the weights will be saved in .npz format.
force: bool
True if the file can be overwritten, False otherwise.
ndir:
Number of directions for the voxelization.
"""
for f in [out_ind, out_wei]:
talon.cli.utils.check_can_write_file(f, force=force)
logging.info(f'Loading reference image: {in_img}')
img = nib.load(in_img)
affine = np.linalg.inv(img.affine)
image_shape = img.shape
logging.info('Image affine:')
logging.info(img.affine)
logging.info(f'Image shape: {image_shape}')
logging.info(f'Number of directions: {ndir}')
vertices = talon.utils.directions(number_of_points=ndir)
logging.info(f'Loading tractogram: {in_tracks}')
tracto = nib.streamlines.load(in_tracks)
logging.info('Applying inverse affine to streamlines to get them in voxel '
'coordinates.')
tracto.tractogram.apply_affine(affine)
streamlines = tracto.streamlines
logging.info(f'Loaded {len(streamlines)}')
logging.info(f'Doing voxelization. This might take some time...')
i, w = talon.voxelize(streamlines, vertices, image_shape)
logging.info(f'Saving indices in {out_ind}')
sp.save_npz(out_ind, i)
logging.info(f'Saving weights in {out_wei}')
sp.save_npz(out_wei, w)
logging.info('Voxelization complete')