# -*- coding: utf-8 -*-
from copy import deepcopy
import json
import logging
import os
import shutil
from typing import Any, List, Tuple, Type, Union, Optional
import zipfile
import nibabel as nib
from nibabel.affines import voxel_sizes
from nibabel.nifti1 import Nifti1Header, Nifti1Image
from nibabel.orientations import aff2axcodes
from nibabel.streamlines.array_sequence import ArraySequence
from nibabel.streamlines.trk import TrkFile
from nibabel.streamlines.tractogram import Tractogram, LazyTractogram
import numpy as np
from trx.io import get_trx_tmp_dir
from trx.utils import (append_generator_to_dict,
close_or_delete_mmap,
convert_data_dict_to_tractogram,
get_reference_info_wrapper)
try:
import dipy
except ImportError:
dipy_available = False
[docs]
def _append_last_offsets(nib_offsets: np.ndarray, nb_vertices: int) -> np.ndarray:
"""Appends the last element of offsets from header information
Keyword arguments:
nib_offsets -- np.ndarray
Array of offsets with the last element being the start of the last
streamline (nibabel convention)
nb_vertices -- int
Total number of vertices in the streamlines
Returns:
Offsets -- np.ndarray (VTK convention)
"""
def is_sorted(a): return np.all(a[:-1] <= a[1:])
if not is_sorted(nib_offsets):
raise ValueError('Offsets must be sorted values.')
return np.append(nib_offsets, nb_vertices).astype(nib_offsets.dtype)
[docs]
def _generate_filename_from_data(arr: np.ndarray, filename: str) -> str:
"""Determines the data type from array data and generates the appropriate
filename
Keyword arguments:
arr -- a NumPy array (1-2D, otherwise ValueError raised)
filename -- the original filename
Returns:
An updated filename
"""
base, ext = os.path.splitext(filename)
if ext:
logging.warning("Will overwrite provided extension if needed.")
dtype = arr.dtype
dtype = "bit" if dtype == bool else dtype.name
if arr.ndim == 1:
new_filename = "{}.{}".format(base, dtype)
elif arr.ndim == 2:
dim = arr.shape[-1]
if dim == 1:
new_filename = "{}.{}".format(base, dtype)
else:
new_filename = "{}.{}.{}".format(base, arr.shape[-1], dtype)
else:
raise ValueError("Invalid dimensionality.")
return new_filename
[docs]
def _split_ext_with_dimensionality(filename: str) -> Tuple[str, int, str]:
"""Takes a filename and splits it into its components
Keyword arguments:
filename -- Input filename
Returns:
tuple of strings (basename, dimension, extension)
"""
basename = os.path.basename(filename)
split = basename.split(".")
if len(split) != 2 and len(split) != 3:
raise ValueError("Invalid filename.")
basename = split[0]
ext = ".{}".format(split[-1])
dim = 1 if len(split) == 2 else split[1]
_is_dtype_valid(ext)
return basename, int(dim), ext
[docs]
def _compute_lengths(offsets: np.ndarray) -> np.ndarray:
"""Compute lengths from offsets
Keyword arguments:
offsets -- An np.ndarray of offsets
Returns:
lengths -- An np.ndarray of lengths
"""
if len(offsets) > 0:
last_elem_pos = _dichotomic_search(offsets)
lengths = np.ediff1d(offsets)
if len(lengths) > last_elem_pos:
lengths[last_elem_pos] = 0
else:
lengths = np.array([0])
return lengths.astype(np.uint32)
[docs]
def _is_dtype_valid(ext: str) -> bool:
"""Verifies that filename extension is a valid datatype
Keyword arguments:
ext -- filename extension
Returns:
boolean representing if provided datatype is valid
"""
if ext.replace(".", "") == "bit":
return True
try:
isinstance(np.dtype(ext.replace(".", "")), np.dtype)
return True
except TypeError:
return False
[docs]
def _dichotomic_search(
x: np.ndarray, l_bound: Optional[int] = None, r_bound: Optional[int] = None
) -> int:
"""Find where data of a contiguous array is actually ending
Keyword arguments:
x -- np.ndarray of values
l_bound -- lower bound index for search
r_bound -- upper bound index for search
Returns:
index at which array value is 0 (if possible), otherwise returns -1"""
if l_bound is None and r_bound is None:
l_bound = 0
r_bound = len(x) - 1
if l_bound == r_bound:
val = l_bound if x[l_bound] != 0 else -1
return val
mid_bound = (l_bound + r_bound + 1) // 2
if x[mid_bound] == 0:
return _dichotomic_search(x, l_bound, mid_bound - 1)
else:
return _dichotomic_search(x, mid_bound, r_bound)
[docs]
def _create_memmap(
filename: str,
mode: str = "r",
shape: Tuple = (1,),
dtype: np.dtype = np.float32,
offset: int = 0,
order: str = "C",
) -> np.ndarray:
"""Wrapper to support empty array as memmaps
Keyword arguments:
filename -- filename where the empty memmap should be created
mode -- file open mode (see: np.memmap for options)
shape -- shape of memmapped np.ndarray
dtype -- datatype of memmapped np.ndarray
offset -- offset of the data within the file
order -- data representation on disk (C or Fortran)
Returns:
mmapped np.ndarray or a zero-filled Numpy array if array has a shape of 0
in the first dimension
"""
if np.dtype(dtype) == bool:
filename = filename.replace(".bool", ".bit")
if shape[0]:
return np.memmap(
filename, mode=mode, offset=offset, shape=shape, dtype=dtype,
order=order
)
else:
if not os.path.isfile(filename):
f = open(filename, "wb")
f.close()
return np.zeros(shape, dtype=dtype)
[docs]
def load(input_obj: str, check_dpg: bool = True) -> Type["TrxFile"]:
"""Load a TrxFile (compressed or not)
Keyword arguments:
input_obj -- A directory name or filepath to the trx data
check_dpg -- Boolean denoting if group metadata should be checked
Returns:
TrxFile object representing the read data
"""
# TODO Check if 0 streamlines, then 0 vertices is expected (vice-versa)
# TODO 4x4 affine matrices should contains values (no all-zeros)
# TODO 3x1 dimensions array should contains values at each position (int)
if os.path.isfile(input_obj):
was_compressed = False
with zipfile.ZipFile(input_obj, "r") as zf:
for info in zf.infolist():
if info.compress_type != 0:
was_compressed = True
break
if was_compressed:
with zipfile.ZipFile(input_obj, "r") as zf:
tmp_dir = get_trx_tmp_dir()
zf.extractall(tmp_dir.name)
trx = load_from_directory(tmp_dir.name)
trx._uncompressed_folder_handle = tmp_dir
logging.info(
"File was compressed, call the close() function before "
"exiting."
)
else:
trx = load_from_zip(input_obj)
elif os.path.isdir(input_obj):
trx = load_from_directory(input_obj)
else:
raise ValueError("File/Folder does not exist")
# Example of robust check for metadata
if check_dpg:
for dpg in trx.data_per_group.keys():
if dpg not in trx.groups.keys():
raise ValueError(
"An undeclared group ({}) has " "data_per_group.".format(
dpg)
)
return trx
[docs]
def load_from_zip(filename: str) -> Type["TrxFile"]:
"""Load a TrxFile from a single zipfile. Note: does not work with
compressed zipfiles
Keyword arguments:
filename -- path of the zipped TrxFile
Returns:
TrxFile representing the read data
"""
with zipfile.ZipFile(filename, mode="r") as zf:
with zf.open("header.json") as zf_header:
header = json.load(zf_header)
header["VOXEL_TO_RASMM"] = np.reshape(
header["VOXEL_TO_RASMM"], (4, 4)
).astype(np.float32)
header["DIMENSIONS"] = np.array(
header["DIMENSIONS"], dtype=np.uint16)
files_pointer_size = {}
for zip_info in zf.filelist:
elem_filename = zip_info.filename
_, ext = os.path.splitext(elem_filename)
if ext == ".json" or zip_info.is_dir():
continue
if not _is_dtype_valid(ext):
continue
raise ValueError(
"The dtype {} is not supported".format(elem_filename))
if ext == ".bit":
ext = ".bool"
mem_adress = zip_info.header_offset + len(zip_info.FileHeader())
dtype_size = np.dtype(ext[1:]).itemsize
size = zip_info.file_size / dtype_size
if len(zip_info.extra):
mem_adress -= len(zip_info.extra)
if size.is_integer():
files_pointer_size[elem_filename] = mem_adress, int(size)
else:
raise ValueError("Wrong size or datatype")
return TrxFile._create_trx_from_pointer(
header, files_pointer_size, root_zip=filename
)
[docs]
def load_from_directory(directory: str) -> Type["TrxFile"]:
"""Load a TrxFile from a folder containing memmaps
Keyword arguments:
filename -- path of the zipped TrxFile
Returns:
TrxFile representing the read data
"""
directory = os.path.abspath(directory)
with open(os.path.join(directory, "header.json")) as header:
header = json.load(header)
header["VOXEL_TO_RASMM"] = np.reshape(header["VOXEL_TO_RASMM"],
(4, 4)).astype(np.float32)
header["DIMENSIONS"] = np.array(header["DIMENSIONS"], dtype=np.uint16)
files_pointer_size = {}
for root, dirs, files in os.walk(directory):
for name in files:
elem_filename = os.path.join(root, name)
_, ext = os.path.splitext(elem_filename)
if ext == ".json":
continue
if not _is_dtype_valid(ext):
raise ValueError(
"The dtype of {} is not supported".format(elem_filename)
)
if ext == ".bit":
ext = ".bool"
dtype_size = np.dtype(ext[1:]).itemsize
size = os.path.getsize(elem_filename) / dtype_size
if size.is_integer():
files_pointer_size[elem_filename] = 0, int(size)
elif os.path.getsize(elem_filename) == 1:
files_pointer_size[elem_filename] = 0, 0
else:
raise ValueError("Wrong size or datatype")
return TrxFile._create_trx_from_pointer(header, files_pointer_size,
root=directory)
[docs]
def concatenate(
trx_list: List["TrxFile"],
delete_dpv: bool = False,
delete_dps: bool = False,
delete_groups: bool = False,
check_space_attributes: bool = True,
preallocation: bool = False,
) -> "TrxFile":
"""Concatenate multiple TrxFile together, support preallocation
Keyword arguments:
trx_list -- A list containing TrxFiles to concatenate
delete_dpv -- Delete dpv keys that do not exist in all the provided
TrxFiles
delete_dps -- Delete dps keys that do not exist in all the provided
TrxFile
delete_groups -- Delete all the groups that currently exist in the
TrxFiles
check_space_attributes -- Verify that dimensions and size of data are
similar between all the TrxFiles
preallocation -- Preallocated TrxFile has already been generated and
is the first element in trx_list
(Note: delete_groups must be set to True as well)
Returns:
TrxFile representing the concatenated data
"""
trx_list = [
curr_trx for curr_trx in trx_list if curr_trx.header["NB_STREAMLINES"] > 0
]
if len(trx_list) == 0:
logging.warning("Inputs of concatenation were empty.")
return TrxFile()
ref_trx = trx_list[0]
all_dps = []
all_dpv = []
for curr_trx in trx_list:
all_dps.extend(list(curr_trx.data_per_streamline.keys()))
all_dpv.extend(list(curr_trx.data_per_vertex.keys()))
all_dps, all_dpv = set(all_dps), set(all_dpv)
if check_space_attributes:
for curr_trx in trx_list[1:]:
if not np.allclose(
ref_trx.header["VOXEL_TO_RASMM"], curr_trx.header["VOXEL_TO_RASMM"]
) or not np.array_equal(
ref_trx.header["DIMENSIONS"], curr_trx.header["DIMENSIONS"]
):
raise ValueError("Wrong space attributes.")
if preallocation and not delete_groups:
raise ValueError(
"Groups are variables, cannot be handled with " "preallocation"
)
# Verifying the validity of fixed-size arrays, coherence between inputs
for curr_trx in trx_list:
for key in all_dpv:
if key not in ref_trx.data_per_vertex.keys() \
or key not in curr_trx.data_per_vertex.keys():
if not delete_dpv:
logging.debug(
"{} dpv key does not exist in all TrxFile.".format(key)
)
raise ValueError(
"TrxFile must be sharing identical dpv " "keys.")
elif (
ref_trx.data_per_vertex[key]._data.dtype
!= curr_trx.data_per_vertex[key]._data.dtype
):
logging.debug(
"{} dpv key is not declared with the same dtype "
"in all TrxFile.".format(key)
)
raise ValueError("Shared dpv key, has different dtype.")
for curr_trx in trx_list:
for key in all_dps:
if key not in ref_trx.data_per_streamline.keys() \
or key not in curr_trx.data_per_streamline.keys():
if not delete_dps:
logging.debug(
"{} dps key does not exist in all " "TrxFile.".format(
key)
)
raise ValueError(
"TrxFile must be sharing identical dps " "keys.")
elif (
ref_trx.data_per_streamline[key].dtype
!= curr_trx.data_per_streamline[key].dtype
):
logging.debug(
"{} dps key is not declared with the same dtype "
"in all TrxFile.".format(key)
)
raise ValueError("Shared dps key, has different dtype.")
all_groups_len = {}
all_groups_dtype = {}
# Variable-size arrays do not have to exist in all TrxFile
if not delete_groups:
for trx_1 in trx_list:
for group_key in trx_1.groups.keys():
# Concatenating groups together
if group_key in all_groups_len:
all_groups_len[group_key] += len(trx_1.groups[group_key])
else:
all_groups_len[group_key] = len(trx_1.groups[group_key])
if (
group_key in all_groups_dtype
and trx_1.groups[group_key].dtype != all_groups_dtype[group_key]
):
raise ValueError("Shared group key, has different dtype.")
else:
all_groups_dtype[group_key] = trx_1.groups[group_key].dtype
# Once the checks are done, actually concatenate
to_concat_list = trx_list[1:] if preallocation else trx_list
if not preallocation:
nb_vertices = 0
nb_streamlines = 0
for curr_trx in to_concat_list:
curr_strs_len, curr_pts_len = curr_trx._get_real_len()
nb_streamlines += curr_strs_len
nb_vertices += curr_pts_len
new_trx = TrxFile(
nb_vertices=nb_vertices, nb_streamlines=nb_streamlines,
init_as=ref_trx
)
if delete_dps:
new_trx.data_per_streamline = {}
if delete_dpv:
new_trx.data_per_vertex = {}
if delete_groups:
new_trx.groups = {}
tmp_dir = new_trx._uncompressed_folder_handle.name
# When memory is allocated on the spot, groups and data_per_group can
# be concatenated together
for group_key in all_groups_len.keys():
if not os.path.isdir(os.path.join(tmp_dir, "groups/")):
os.mkdir(os.path.join(tmp_dir, "groups/"))
dtype = all_groups_dtype[group_key]
group_filename = os.path.join(
tmp_dir, "groups/" "{}.{}".format(group_key, dtype.name)
)
group_len = all_groups_len[group_key]
new_trx.groups[group_key] = _create_memmap(
group_filename, mode="w+", shape=(group_len,), dtype=dtype
)
if delete_groups:
continue
pos = 0
count = 0
for curr_trx in trx_list:
curr_len = len(curr_trx.groups[group_key])
new_trx.groups[group_key][pos: pos + curr_len] = \
curr_trx.groups[group_key] + count
pos += curr_len
count += curr_trx.header["NB_STREAMLINES"]
strs_end, pts_end = 0, 0
else:
new_trx = ref_trx
strs_end, pts_end = new_trx._get_real_len()
for curr_trx in to_concat_list:
# Copy the TrxFile fixed-size info (the right chunk)
strs_end, pts_end = new_trx._copy_fixed_arrays_from(
curr_trx, strs_start=strs_end, pts_start=pts_end
)
return new_trx
[docs]
def save(
trx: "TrxFile", filename: str, compression_standard: Any = zipfile.ZIP_STORED
) -> None:
"""Save a TrxFile (compressed or not)
Keyword arguments:
trx -- The TrxFile to save
filename -- The path to save the TrxFile to
compression_standard -- The compression standard to use, as defined by
the ZipFile library
"""
_, ext = os.path.splitext(filename)
if ext not in [".zip", ".trx", ""]:
raise ValueError("Unsupported extension.")
copy_trx = trx.deepcopy()
copy_trx.resize()
tmp_dir_name = copy_trx._uncompressed_folder_handle.name
if ext in [".zip", ".trx"]:
zip_from_folder(tmp_dir_name, filename, compression_standard)
else:
if os.path.isdir(filename):
shutil.rmtree(filename)
shutil.copytree(tmp_dir_name, filename)
copy_trx.close()
[docs]
def zip_from_folder(
directory: str, filename: str, compression_standard: Any = zipfile.ZIP_STORED
) -> None:
"""Utils function to zip on-disk memmaps
Keyword arguments
directory -- The path to the on-disk memmap
filename -- The path where the zip file should be created
compression_standard -- The compression standard to use, as defined by
the ZipFile library
"""
with zipfile.ZipFile(filename, mode="w", compression=compression_standard) as zf:
for root, _, files in os.walk(directory):
for name in files:
curr_filename = os.path.join(root, name)
tmp_filename = curr_filename.replace(directory, "")[1:]
zf.write(curr_filename, tmp_filename)
[docs]
class TrxFile:
"""Core class of the TrxFile"""
[docs]
streamlines: Type[ArraySequence]
[docs]
data_per_streamline: dict
def __init__(
self,
nb_vertices: Optional[int] = None,
nb_streamlines: Optional[int] = None,
init_as: Optional[Type["TrxFile"]] = None,
reference: Union[
str,
dict,
Type[Nifti1Image],
Type[TrkFile],
Type[Nifti1Header],
None,
] = None,
) -> None:
"""Initialize an empty TrxFile, support preallocation
Keyword Arguments:
nb_vertices -- The number of vertices to use in the new TrxFile
nb_streamlines -- The number of streamlines in the new TrxFile
init_as -- A TrxFile to use as reference
reference -- A Nifti or Trk file/obj to use as reference
"""
if init_as is not None:
affine = init_as.header["VOXEL_TO_RASMM"]
dimensions = init_as.header["DIMENSIONS"]
elif reference is not None:
affine, dimensions, _, _ = get_reference_info_wrapper(reference)
else:
logging.debug(
"No reference provided, using blank space "
"attributes, please update them later."
)
affine = np.eye(4).astype(np.float32)
dimensions = np.array([1, 1, 1], dtype=np.uint16)
if nb_vertices is None and nb_streamlines is None:
if init_as is not None:
raise ValueError(
"Cant use init_as without declaring "
"nb_vertices AND nb_streamlines"
)
logging.debug("Intializing empty TrxFile.")
self.header = {}
# Using the new format default type
tmp_strs = ArraySequence()
tmp_strs._data = tmp_strs._data.astype(np.float32)
tmp_strs._offsets = tmp_strs._offsets.astype(np.uint32)
tmp_strs._lengths = tmp_strs._lengths.astype(np.uint32)
self.streamlines = tmp_strs
self.groups = {}
self.data_per_streamline = {}
self.data_per_vertex = {}
self.data_per_group = {}
self._uncompressed_folder_handle = None
nb_vertices = 0
nb_streamlines = 0
elif nb_vertices is not None and nb_streamlines is not None:
logging.debug(
"Preallocating TrxFile with size {} streamlines"
"and {} vertices.".format(nb_streamlines, nb_vertices)
)
trx = self._initialize_empty_trx(
nb_streamlines, nb_vertices, init_as=init_as
)
self.__dict__ = trx.__dict__
else:
raise ValueError(
"You must declare both nb_vertices AND " "NB_STREAMLINES")
self.header["VOXEL_TO_RASMM"] = affine
self.header["DIMENSIONS"] = dimensions
self.header["NB_VERTICES"] = nb_vertices
self.header["NB_STREAMLINES"] = nb_streamlines
self._copy_safe = True
[docs]
def __str__(self) -> str:
"""Generate the string for printing"""
affine = np.array(self.header["VOXEL_TO_RASMM"], dtype=np.float32)
dimensions = np.array(self.header["DIMENSIONS"], dtype=np.uint16)
vox_sizes = np.array(voxel_sizes(affine), dtype=np.float32)
vox_order = "".join(aff2axcodes(affine))
text = "VOXEL_TO_RASMM: \n{}".format(
np.array2string(affine, formatter={
"float_kind": lambda x: "%.6f" % x})
)
text += "\nDIMENSIONS: {}".format(np.array2string(dimensions))
text += "\nVOX_SIZES: {}".format(
np.array2string(vox_sizes, formatter={
"float_kind": lambda x: "%.2f" % x})
)
text += "\nVOX_ORDER: {}".format(vox_order)
strs_size = self.header["NB_STREAMLINES"]
pts_size = self.header["NB_VERTICES"]
strs_len, pts_len = self._get_real_len()
if strs_size != strs_len or pts_size != pts_len:
text += "\nstreamline_size: {}".format(strs_size)
text += "\nvertex_size: {}".format(pts_size)
text += "\nstreamline_count: {}".format(strs_len)
text += "\nvertex_count: {}".format(pts_len)
text += "\ndata_per_vertex keys: {}".format(
list(self.data_per_vertex.keys()))
text += "\ndata_per_streamline keys: {}".format(
list(self.data_per_streamline.keys())
)
text += "\ngroups keys: {}".format(list(self.groups.keys()))
for group_key in self.groups.keys():
if group_key in self.data_per_group:
text += "\ndata_per_groups ({}) keys: {}".format(
group_key, list(self.data_per_group[group_key].keys())
)
text += "\ncopy_safe: {}".format(self._copy_safe)
return text
[docs]
def __len__(self) -> int:
"""Define the length of the object"""
return len(self.streamlines)
[docs]
def __getitem__(self, key) -> Any:
"""Slice all data in a consistent way"""
if isinstance(key, int):
if key < 0:
key += len(self)
key = [key]
elif isinstance(key, slice):
key = [ii for ii in range(*key.indices(len(self)))]
return self.select(key, keep_group=False)
[docs]
def __deepcopy__(self) -> Type["TrxFile"]:
return self.deepcopy()
[docs]
def deepcopy(self) -> Type["TrxFile"]:
"""Create a deepcopy of the TrxFile
Returns
A deepcopied TrxFile of the current TrxFile
"""
tmp_dir = get_trx_tmp_dir()
out_json = open(os.path.join(tmp_dir.name, "header.json"), "w")
tmp_header = deepcopy(self.header)
if not isinstance(tmp_header["VOXEL_TO_RASMM"], list):
tmp_header["VOXEL_TO_RASMM"] = tmp_header["VOXEL_TO_RASMM"].tolist()
if not isinstance(tmp_header["DIMENSIONS"], list):
tmp_header["DIMENSIONS"] = tmp_header["DIMENSIONS"].tolist()
# tofile() alway write in C-order
if not self._copy_safe:
to_dump = self.streamlines.copy()._data
tmp_header["NB_STREAMLINES"] = len(self.streamlines)
tmp_header["NB_VERTICES"] = len(to_dump)
else:
to_dump = self.streamlines._data
json.dump(tmp_header, out_json)
out_json.close()
positions_filename = _generate_filename_from_data(
to_dump, os.path.join(tmp_dir.name, "positions")
)
to_dump.tofile(positions_filename)
if not self._copy_safe:
to_dump = _append_last_offsets(self.streamlines.copy()._offsets,
self.header["NB_VERTICES"])
else:
to_dump = _append_last_offsets(self.streamlines._offsets,
self.header["NB_VERTICES"])
offsets_filename = _generate_filename_from_data(
self.streamlines._offsets, os.path.join(tmp_dir.name, "offsets")
)
to_dump.tofile(offsets_filename)
if len(self.data_per_vertex.keys()) > 0:
os.mkdir(os.path.join(tmp_dir.name, "dpv/"))
for dpv_key in self.data_per_vertex.keys():
if not self._copy_safe:
to_dump = self.data_per_vertex[dpv_key].copy()._data
else:
to_dump = self.data_per_vertex[dpv_key]._data
dpv_filename = _generate_filename_from_data(
to_dump, os.path.join(tmp_dir.name, "dpv/", dpv_key)
)
to_dump.tofile(dpv_filename)
if len(self.data_per_streamline.keys()) > 0:
os.mkdir(os.path.join(tmp_dir.name, "dps/"))
for dps_key in self.data_per_streamline.keys():
to_dump = self.data_per_streamline[dps_key]
dps_filename = _generate_filename_from_data(
to_dump, os.path.join(tmp_dir.name, "dps/", dps_key)
)
to_dump.tofile(dps_filename)
if len(self.groups.keys()) > 0:
os.mkdir(os.path.join(tmp_dir.name, "groups/"))
for group_key in self.groups.keys():
to_dump = self.groups[group_key]
group_filename = _generate_filename_from_data(
to_dump, os.path.join(tmp_dir.name, "groups/", group_key)
)
to_dump.tofile(group_filename)
if group_key not in self.data_per_group:
continue
for dpg_key in self.data_per_group[group_key].keys():
# Creates 'dpg/' only if required
if not os.path.isdir(os.path.join(tmp_dir.name, "dpg/")):
os.mkdir(os.path.join(tmp_dir.name, "dpg/"))
if not os.path.isdir(os.path.join(tmp_dir.name, "dpg/", group_key)):
os.mkdir(os.path.join(tmp_dir.name, "dpg/", group_key))
to_dump = self.data_per_group[group_key][dpg_key]
dpg_filename = _generate_filename_from_data(
to_dump, os.path.join(
tmp_dir.name, "dpg/", group_key, dpg_key)
)
to_dump.tofile(dpg_filename)
copy_trx = load_from_directory(tmp_dir.name)
copy_trx._uncompressed_folder_handle = tmp_dir
return copy_trx
[docs]
def _get_real_len(self) -> Tuple[int, int]:
"""Get the real size of data (ignoring zeros of preallocation)
Returns
A tuple representing the index of the last streamline and the total
length of all the streamlines
"""
if len(self.streamlines._lengths) == 0:
return 0, 0
last_elem_pos = _dichotomic_search(self.streamlines._lengths)
if last_elem_pos != -1:
strs_end = int(last_elem_pos + 1)
pts_end = int(np.sum(self.streamlines._lengths[0:strs_end]))
return strs_end, pts_end
return 0, 0
[docs]
def _copy_fixed_arrays_from(
self,
trx: Type["TrxFile"],
strs_start: int = 0,
pts_start: int = 0,
nb_strs_to_copy: Optional[int] = None,
) -> Tuple[int, int]:
"""Fill a TrxFile using another and start indexes (preallocation)
Keyword arguments:
trx -- TrxFile to copy data from
strs_start -- The start index of the streamline
pts_start -- The start index of the point
nb_strs_to_copy -- The number of streamlines to copy. If not set
will copy all
Returns
A tuple representing the end of the copied streamlines and end of
copied points
"""
if nb_strs_to_copy is None:
curr_strs_len, curr_pts_len = trx._get_real_len()
else:
curr_strs_len = int(nb_strs_to_copy)
curr_pts_len = np.sum(trx.streamlines._lengths[0:curr_strs_len])
curr_pts_len = int(curr_pts_len)
strs_end = strs_start + curr_strs_len
pts_end = pts_start + curr_pts_len
if curr_pts_len == 0:
return strs_start, pts_start
# Mandatory arrays
self.streamlines._data[pts_start:pts_end] = \
trx.streamlines._data[0:curr_pts_len]
self.streamlines._offsets[strs_start:strs_end] = \
(trx.streamlines._offsets[0:curr_strs_len] + pts_start)
self.streamlines._lengths[strs_start:strs_end] = \
trx.streamlines._lengths[0:curr_strs_len]
# Optional fixed-sized arrays
for dpv_key in self.data_per_vertex.keys():
self.data_per_vertex[dpv_key]._data[
pts_start:pts_end
] = trx.data_per_vertex[dpv_key]._data[0:curr_pts_len]
self.data_per_vertex[dpv_key]._offsets = self.streamlines._offsets
self.data_per_vertex[dpv_key]._lengths = self.streamlines._lengths
for dps_key in self.data_per_streamline.keys():
self.data_per_streamline[dps_key][
strs_start:strs_end
] = trx.data_per_streamline[dps_key][0:curr_strs_len]
return strs_end, pts_end
@staticmethod
[docs]
def _initialize_empty_trx(
nb_streamlines: int, nb_vertices: int,
init_as: Optional[Type["TrxFile"]] = None) -> Type["TrxFile"]:
"""Create on-disk memmaps of a certain size (preallocation)
Keyword arguments:
nb_streamlines -- The number of streamlines that the empty TrxFile
will be initialized with
nb_vertices -- The number of vertices that the empty TrxFile will
be initialized with
init_as -- A TrxFile to initialize the empty TrxFile with
Returns:
An empty TrxFile preallocated with a certain size
"""
trx = TrxFile()
tmp_dir = get_trx_tmp_dir()
logging.info("Temporary folder for memmaps: {}".format(tmp_dir.name))
trx.header["NB_VERTICES"] = nb_vertices
trx.header["NB_STREAMLINES"] = nb_streamlines
if init_as is not None:
trx.header["VOXEL_TO_RASMM"] = init_as.header["VOXEL_TO_RASMM"]
trx.header["DIMENSIONS"] = init_as.header["DIMENSIONS"]
positions_dtype = init_as.streamlines._data.dtype
offsets_dtype = init_as.streamlines._offsets.dtype
lengths_dtype = init_as.streamlines._lengths.dtype
else:
positions_dtype = np.dtype(np.float16)
offsets_dtype = np.dtype(np.uint32)
lengths_dtype = np.dtype(np.uint32)
logging.debug(
"Initializing positions with dtype: {}".format(
positions_dtype.name)
)
logging.debug(
"Initializing offsets with dtype: {}".format(offsets_dtype.name))
logging.debug(
"Initializing lengths with dtype: {}".format(lengths_dtype.name))
# A TrxFile without init_as only contain the essential arrays
positions_filename = os.path.join(
tmp_dir.name, "positions.3.{}".format(positions_dtype.name)
)
trx.streamlines._data = _create_memmap(
positions_filename, mode="w+", shape=(nb_vertices, 3),
dtype=positions_dtype
)
offsets_filename = os.path.join(
tmp_dir.name, "offsets.{}".format(offsets_dtype.name)
)
trx.streamlines._offsets = _create_memmap(
offsets_filename, mode="w+", shape=(nb_streamlines,),
dtype=offsets_dtype
)
trx.streamlines._lengths = np.zeros(
shape=(nb_streamlines,), dtype=lengths_dtype
)
# Only the structure of fixed-size arrays is copied
if init_as is not None:
if len(init_as.data_per_vertex.keys()) > 0:
os.mkdir(os.path.join(tmp_dir.name, "dpv/"))
if len(init_as.data_per_streamline.keys()) > 0:
os.mkdir(os.path.join(tmp_dir.name, "dps/"))
for dpv_key in init_as.data_per_vertex.keys():
dtype = init_as.data_per_vertex[dpv_key]._data.dtype
tmp_as = init_as.data_per_vertex[dpv_key]._data
if tmp_as.ndim == 1:
dpv_filename = os.path.join(
tmp_dir.name, "dpv/" "{}.{}".format(
dpv_key, dtype.name)
)
shape = (nb_vertices, 1)
elif tmp_as.ndim == 2:
dim = tmp_as.shape[-1]
shape = (nb_vertices, dim)
dpv_filename = os.path.join(
tmp_dir.name, "dpv/" "{}.{}.{}".format(
dpv_key, dim, dtype.name)
)
else:
raise ValueError("Invalid dimensionality.")
logging.debug(
"Initializing {} (dpv) with dtype: "
"{}".format(dpv_key, dtype.name)
)
trx.data_per_vertex[dpv_key] = ArraySequence()
trx.data_per_vertex[dpv_key]._data = _create_memmap(
dpv_filename, mode="w+", shape=shape, dtype=dtype
)
trx.data_per_vertex[dpv_key]._offsets = trx.streamlines._offsets
trx.data_per_vertex[dpv_key]._lengths = trx.streamlines._lengths
for dps_key in init_as.data_per_streamline.keys():
dtype = init_as.data_per_streamline[dps_key].dtype
tmp_as = init_as.data_per_streamline[dps_key]
if tmp_as.ndim == 1:
dps_filename = os.path.join(
tmp_dir.name, "dps/" "{}.{}".format(
dps_key, dtype.name)
)
shape = (nb_streamlines,)
elif tmp_as.ndim == 2:
dim = tmp_as.shape[-1]
shape = (nb_streamlines, dim)
dps_filename = os.path.join(
tmp_dir.name, "dps/" "{}.{}.{}".format(
dps_key, dim, dtype.name)
)
else:
raise ValueError("Invalid dimensionality.")
logging.debug(
"Initializing {} (dps) with and dtype: "
"{}".format(dps_key, dtype.name)
)
trx.data_per_streamline[dps_key] = _create_memmap(
dps_filename, mode="w+", shape=shape, dtype=dtype
)
trx._uncompressed_folder_handle = tmp_dir
return trx
[docs]
def _create_trx_from_pointer(
header: dict,
dict_pointer_size: dict,
root_zip: Optional[str] = None,
root: Optional[str] = None,
) -> Type["TrxFile"]:
"""After reading the structure of a zip/folder, create a TrxFile
Keyword arguments:
header -- A TrxFile header dictionary which will be used for the
new TrxFile
dict_pointer_size -- A dictionary containing the filenames of all
the files within the TrxFile disk file/folder
root_zip -- The path of the ZipFile pointer
root -- The dirname of the ZipFile pointer
Returns:
A TrxFile constructer from the pointer provided
"""
# TODO support empty positions, using optional tag?
trx = TrxFile()
trx.header = header
positions, offsets = None, None
for elem_filename in dict_pointer_size.keys():
if root_zip:
filename = root_zip
else:
filename = elem_filename
folder = os.path.dirname(elem_filename)
base, dim, ext = _split_ext_with_dimensionality(elem_filename)
if ext == ".bit":
ext = ".bool"
mem_adress, size = dict_pointer_size[elem_filename]
if root is not None:
# This is for Unix
if os.name != 'nt' and folder.startswith(root.rstrip("/")):
folder = folder.replace(root, "").lstrip("/")
# These three are for Windows
elif os.path.isdir(folder) and os.path.basename(folder) in ['dpv', 'dps', 'groups']:
folder = os.path.basename(folder)
elif os.path.basename(os.path.dirname(folder)) == 'dpg':
folder = os.path.join('dpg', os.path.basename(folder))
else:
folder = ''
# Parse/walk the directory tree
if base == "positions" and folder == "":
if size != trx.header["NB_VERTICES"] * 3 or dim != 3:
raise ValueError("Wrong data size/dimensionality.")
positions = _create_memmap(
filename,
mode="r+",
offset=mem_adress,
shape=(trx.header["NB_VERTICES"], 3),
dtype=ext[1:],
)
elif base == "offsets" and folder == "":
if size != trx.header["NB_STREAMLINES"]+1 or dim != 1:
raise ValueError("Wrong offsets size/dimensionality.")
offsets = _create_memmap(
filename,
mode="r+",
offset=mem_adress,
shape=(trx.header["NB_STREAMLINES"]+1,),
dtype=ext[1:],
)
if offsets[-1] != 0:
lengths = _compute_lengths(offsets)
else:
lengths = [0]
elif folder == "dps":
nb_scalar = size / trx.header["NB_STREAMLINES"]
if not nb_scalar.is_integer() or nb_scalar != dim:
raise ValueError("Wrong dps size/dimensionality.")
else:
shape = (trx.header["NB_STREAMLINES"], int(nb_scalar))
trx.data_per_streamline[base] = _create_memmap(
filename, mode="r+", offset=mem_adress, shape=shape,
dtype=ext[1:]
)
elif folder == "dpv":
nb_scalar = size / trx.header["NB_VERTICES"]
if not nb_scalar.is_integer() or nb_scalar != dim:
raise ValueError("Wrong dpv size/dimensionality.")
else:
shape = (trx.header["NB_VERTICES"], int(nb_scalar))
trx.data_per_vertex[base] = _create_memmap(
filename, mode="r+", offset=mem_adress, shape=shape,
dtype=ext[1:]
)
elif folder.startswith("dpg"):
if int(size) != dim:
raise ValueError("Wrong dpg size/dimensionality.")
else:
shape = (1, int(size))
# Handle the two-layers architecture
data_name = os.path.basename(base)
sub_folder = os.path.basename(folder)
if sub_folder not in trx.data_per_group:
trx.data_per_group[sub_folder] = {}
trx.data_per_group[sub_folder][data_name] = _create_memmap(
filename, mode="r+", offset=mem_adress, shape=shape,
dtype=ext[1:]
)
elif folder == "groups":
# Groups are simply indices, nothing else
# TODO Crash if not uint?
if dim != 1:
raise ValueError("Wrong group dimensionality.")
else:
shape = (int(size),)
trx.groups[base] = _create_memmap(
filename, mode="r+", offset=mem_adress, shape=shape,
dtype=ext[1:]
)
else:
logging.error(
"{} is not part of a valid structure.".format(
elem_filename)
)
# All essential array must be declared
if positions is not None and offsets is not None:
trx.streamlines._data = positions
trx.streamlines._offsets = offsets[:-1]
trx.streamlines._lengths = lengths
else:
raise ValueError("Missing essential data.")
for dpv_key in trx.data_per_vertex:
tmp = trx.data_per_vertex[dpv_key]
trx.data_per_vertex[dpv_key] = ArraySequence()
trx.data_per_vertex[dpv_key]._data = tmp
trx.data_per_vertex[dpv_key]._offsets = offsets[:-1]
trx.data_per_vertex[dpv_key]._lengths = lengths
return trx
[docs]
def resize(
self,
nb_streamlines: Optional[int] = None,
nb_vertices: Optional[int] = None,
delete_dpg: bool = False,
) -> None:
"""Remove the ununsed portion of preallocated memmaps
Keyword arguments:
nb_streamlines -- The number of streamlines to keep
nb_vertices -- The number of vertices to keep
delete_dpg -- Remove data_per_group when resizing
"""
if not self._copy_safe:
raise ValueError("Cannot resize a sliced datasets.")
strs_end, pts_end = self._get_real_len()
if nb_streamlines is not None and nb_streamlines < strs_end:
strs_end = nb_streamlines
logging.info(
"Resizing (down) memmaps, less streamlines than it "
"actually contains."
)
if nb_vertices is None:
pts_end = int(np.sum(self.streamlines._lengths[0:nb_streamlines]))
nb_vertices = pts_end
elif nb_vertices < pts_end:
# Resizing vertices only is too dangerous, not allowed
logging.warning("Cannot resize (down) vertices for consistency.")
return
if nb_streamlines is None:
nb_streamlines = strs_end
if (
nb_streamlines == self.header["NB_STREAMLINES"]
and nb_vertices == self.header["NB_VERTICES"]
):
logging.debug("TrxFile of the right size, no resizing.")
return
trx = self._initialize_empty_trx(
nb_streamlines, nb_vertices, init_as=self)
logging.info(
"Resizing streamlines from size {} to {}".format(
len(self.streamlines), nb_streamlines
)
)
logging.info(
"Resizing vertices from size {} to {}".format(
len(self.streamlines._data), nb_vertices
)
)
# Copy the fixed-sized info from the original TrxFile to the new
# (resized) one.
if nb_streamlines < self.header["NB_STREAMLINES"]:
trx._copy_fixed_arrays_from(self, nb_strs_to_copy=nb_streamlines)
else:
trx._copy_fixed_arrays_from(self)
tmp_dir = trx._uncompressed_folder_handle.name
if len(self.groups.keys()) > 0:
os.mkdir(os.path.join(tmp_dir, "groups/"))
for group_key in self.groups.keys():
group_dtype = self.groups[group_key].dtype
group_name = os.path.join(
tmp_dir, "groups/", "{}.{}".format(group_key, group_dtype.name)
)
ori_len = len(self.groups[group_key])
# Remove groups indices if resizing down
tmp = self.groups[group_key][self.groups[group_key] < strs_end]
trx.groups[group_key] = _create_memmap(
group_name, mode="w+", shape=(len(tmp),), dtype=group_dtype
)
logging.debug(
"{} group went from {} items to {}".format(
group_key, ori_len, len(tmp))
)
trx.groups[group_key][:] = tmp
if delete_dpg:
self.close()
self.__dict__ = trx.__dict__
return
if len(self.data_per_group.keys()) > 0:
os.mkdir(os.path.join(tmp_dir, "dpg/"))
for group_key in self.data_per_group:
if not os.path.isdir(os.path.join(tmp_dir, "dpg/", group_key)):
os.mkdir(os.path.join(tmp_dir, "dpg/", group_key))
if group_key not in trx.data_per_group:
trx.data_per_group[group_key] = {}
for dpg_key in self.data_per_group[group_key].keys():
dpg_dtype = self.data_per_group[group_key][dpg_key].dtype
dpg_filename = _generate_filename_from_data(
self.data_per_group[group_key][dpg_key],
os.path.join(tmp_dir, "dpg/", group_key, dpg_key),
)
shape = self.data_per_group[group_key][dpg_key].shape
if dpg_key not in trx.data_per_group[group_key]:
trx.data_per_group[group_key][dpg_key] = {}
trx.data_per_group[group_key][dpg_key] = _create_memmap(
dpg_filename, mode="w+", shape=shape, dtype=dpg_dtype
)
trx.data_per_group[group_key][dpg_key][:] = \
self.data_per_group[group_key][dpg_key]
self.close()
self.__dict__ = trx.__dict__
[docs]
def get_dtype_dict(self):
"""Get the dtype dictionary for the TrxFile
Returns
A dictionary containing the dtype for each data element
"""
dtype_dict = {"positions": self.streamlines._data.dtype,
"offsets": self.streamlines._offsets.dtype,
"dpv": {}, "dps": {}, "dpg": {}, "groups": {}}
for key in self.data_per_vertex.keys():
dtype_dict['dpv'][key] = self.data_per_vertex[key]._data.dtype
for key in self.data_per_streamline.keys():
dtype_dict['dps'][key] = self.data_per_streamline[key].dtype
for group_key in self.data_per_group.keys():
dtype_dict['groups'][group_key] = self.groups[group_key].dtype
for group_key in self.data_per_group.keys():
dtype_dict['dpg'][group_key] = {}
for dpg_key in self.data_per_group[group_key].keys():
dtype_dict['dpg'][group_key][dpg_key] = \
self.data_per_group[group_key][dpg_key].dtype
return dtype_dict
[docs]
def append(self, obj, extra_buffer: int = 0) -> None:
curr_dtype_dict = self.get_dtype_dict()
if dipy_available:
from dipy.io.stateful_tractogram import StatefulTractogram
if not isinstance(obj, (TrxFile, Tractogram)) \
and (dipy_available and not isinstance(obj, StatefulTractogram)):
raise TypeError(
"{} is not a supported object type for appending.".format(type(obj)))
elif isinstance(obj, Tractogram):
obj = self.from_tractogram(obj, reference=self.header,
dtype_dict=curr_dtype_dict)
elif dipy_available and isinstance(obj, StatefulTractogram):
obj = self.from_sft(obj, dtype_dict=curr_dtype_dict)
self._append_trx(obj, extra_buffer=extra_buffer)
[docs]
def _append_trx(self, trx: Type["TrxFile"],
extra_buffer: int = 0) -> None:
"""Append a TrxFile to another (support buffer)
Keyword arguments:
trx -- The TrxFile to append to the current TrxFile
extra_buffer -- The additional buffer space required to append data
"""
strs_end, pts_end = self._get_real_len()
nb_streamlines = strs_end + trx.header["NB_STREAMLINES"]
nb_vertices = pts_end + trx.header["NB_VERTICES"]
if (
self.header["NB_STREAMLINES"] < nb_streamlines
or self.header["NB_VERTICES"] < nb_vertices
):
self.resize(
nb_streamlines=nb_streamlines + extra_buffer,
nb_vertices=nb_vertices + extra_buffer * 100,
)
_ = concatenate([self, trx], preallocation=True, delete_groups=True)
[docs]
def get_group(
self, key: str, keep_group: bool = True, copy_safe: bool = False
) -> Type["TrxFile"]:
"""Get a particular group from the TrxFile
Keyword arguments:
key -- The group name to select
keep_group -- Make sure group exists in returned TrxFile
copy_safe -- Perform a deepcopy
Returns
A TrxFile exclusively containing data from said group
"""
return self.select(self.groups[key], keep_group=keep_group, copy_safe=copy_safe)
[docs]
def select(
self, indices: np.ndarray, keep_group: bool = True, copy_safe: bool = False
) -> Type["TrxFile"]:
"""Get a subset of items, always vertices to the same memmaps
Keyword arguments:
indices -- The list of indices of elements to return
keep_group -- Ensure group is returned in output TrxFile
copy_safe -- Perform a deep-copy
Returns:
A TrxFile containing data originating from the selected indices
"""
indices = np.array(indices, dtype=np.uint32)
new_trx = TrxFile()
new_trx._copy_safe = copy_safe
new_trx.header = deepcopy(self.header)
if isinstance(indices, np.ndarray) and len(indices) == 0:
# Even while empty, basic dtype and header must be coherent
positions_dtype = self.streamlines._data.dtype
offsets_dtype = self.streamlines._offsets.dtype
lengths_dtype = self.streamlines._lengths.dtype
new_trx.streamlines._data = new_trx.streamlines._data.reshape(
(0, 3)
).astype(positions_dtype)
new_trx.streamlines._offsets = new_trx.streamlines._offsets.astype(
offsets_dtype
)
new_trx.streamlines._lengths = new_trx.streamlines._lengths.astype(
lengths_dtype
)
new_trx.header["NB_VERTICES"] = len(new_trx.streamlines._data)
new_trx.header["NB_STREAMLINES"] = len(
new_trx.streamlines._lengths)
return new_trx.deepcopy() if copy_safe else new_trx
new_trx.streamlines = (
self.streamlines[indices].copy(
) if copy_safe else self.streamlines[indices]
)
for dpv_key in self.data_per_vertex.keys():
new_trx.data_per_vertex[dpv_key] = (
self.data_per_vertex[dpv_key][indices].copy()
if copy_safe
else self.data_per_vertex[dpv_key][indices]
)
for dps_key in self.data_per_streamline.keys():
new_trx.data_per_streamline[dps_key] = (
self.data_per_streamline[dps_key][indices].copy()
if copy_safe
else self.data_per_streamline[dps_key][indices]
)
# Not keeping group is equivalent to the [] operator
if keep_group:
logging.warning(
"Keeping dpg despite affecting the group " "items.")
for group_key in self.groups.keys():
# Keep the group indices even when fancy slicing
index = np.argsort(indices)
sorted_x = indices[index]
sorted_index = np.searchsorted(
sorted_x, self.groups[group_key])
yindex = np.take(index, sorted_index, mode="clip")
mask = indices[yindex] != self.groups[group_key]
intersect = yindex[~mask]
if len(intersect) == 0:
continue
new_trx.groups[group_key] = intersect
if group_key in self.data_per_group:
for dpg_key in self.data_per_group[group_key].keys():
if group_key not in new_trx.data_per_group:
new_trx.data_per_group[group_key] = {}
new_trx.data_per_group[group_key][
dpg_key
] = self.data_per_group[group_key][dpg_key]
new_trx.header["NB_VERTICES"] = len(new_trx.streamlines._data)
new_trx.header["NB_STREAMLINES"] = len(new_trx.streamlines._lengths)
return new_trx.deepcopy() if copy_safe else new_trx
@staticmethod
[docs]
def from_lazy_tractogram(obj: ["LazyTractogram"], reference,
extra_buffer: int = 0,
chunk_size: int = 10000,
dtype_dict: dict = {'positions': np.float32,
'offsets': np.uint32,
'dpv': {}, 'dps': {}}) \
-> Type["TrxFile"]:
"""Append a TrxFile to another (support buffer)
Keyword arguments:
trx -- The TrxFile to append to the current TrxFile
extra_buffer -- The buffer space between reallocation.
This number should be a number of streamlines.
Use 0 for no buffer.
chunk_size -- The number of streamlines to save at a time.
"""
data = {'strs': [], 'dpv': {}, 'dps': {}}
concat = None
count = 0
iterator = iter(obj)
while True:
if count < chunk_size:
try:
i = next(iterator)
count += 1
except StopIteration:
obj = convert_data_dict_to_tractogram(data)
if concat is None:
if len(obj.streamlines) == 0:
concat = TrxFile()
else:
concat = TrxFile.from_tractogram(obj,
reference=reference,
dtype_dict=dtype_dict)
elif len(obj.streamlines) > 0:
curr_obj = TrxFile.from_tractogram(obj,
reference=reference,
dtype_dict=dtype_dict)
concat.append(curr_obj)
break
append_generator_to_dict(i, data)
else:
obj = convert_data_dict_to_tractogram(data)
if concat is None:
concat = TrxFile.from_tractogram(obj,
reference=reference,
dtype_dict=dtype_dict)
else:
curr_obj = TrxFile.from_tractogram(obj,
reference=reference,
dtype_dict=dtype_dict)
concat.append(curr_obj, extra_buffer=extra_buffer)
data = {'strs': [], 'dpv': {}, 'dps': {}}
count = 0
concat.resize()
return concat
@staticmethod
[docs]
def from_sft(sft, dtype_dict={}):
"""Generate a valid TrxFile from a StatefulTractogram"""
if len(sft.dtype_dict) > 0:
dtype_dict = sft.dtype_dict
if 'dpp' in dtype_dict:
dtype_dict['dpv'] = dtype_dict.pop('dpp')
elif len(dtype_dict) == 0:
dtype_dict = {'positions': np.float32, 'offsets': np.uint32,
'dpv': {}, 'dps': {}}
positions_dtype = dtype_dict['positions']
offsets_dtype = dtype_dict['offsets']
if not np.issubdtype(positions_dtype, np.floating):
logging.warning(
"Casting positions as {}, considering using a floating point "
"dtype.".format(positions_dtype))
if not np.issubdtype(offsets_dtype, np.integer):
logging.warning(
"Casting offsets as {}, considering using a integer "
"dtype.".format(offsets_dtype))
trx = TrxFile(nb_vertices=len(sft.streamlines._data),
nb_streamlines=len(sft.streamlines))
trx.header = {
"DIMENSIONS": sft.dimensions.tolist(),
"VOXEL_TO_RASMM": sft.affine.tolist(),
"NB_VERTICES": len(sft.streamlines._data),
"NB_STREAMLINES": len(sft.streamlines),
}
old_space = deepcopy(sft.space)
old_origin = deepcopy(sft.origin)
# TrxFile are written on disk in RASMM/center convention
sft.to_rasmm()
sft.to_center()
tmp_streamlines = deepcopy(sft.streamlines)
# Cast the int64 of Nibabel to uint32
tmp_streamlines._offsets = tmp_streamlines._offsets.astype(
offsets_dtype)
tmp_streamlines._data = tmp_streamlines._data.astype(positions_dtype)
trx.streamlines = tmp_streamlines
for key in sft.data_per_point:
dtype_to_use = dtype_dict['dpv'][key] if key in dtype_dict['dpv'] \
else np.float32
trx.data_per_vertex[key] = \
sft.data_per_point[key]
trx.data_per_vertex[key]._data = \
sft.data_per_point[key]._data.astype(dtype_to_use)
for key in sft.data_per_streamline:
dtype_to_use = dtype_dict['dps'][key] if key in dtype_dict['dps'] \
else np.float32
trx.data_per_streamline[key] = sft.data_per_streamline[key].astype(
dtype_to_use)
# For safety and for RAM, convert the whole object to memmaps
tmp_dir = get_trx_tmp_dir()
save(trx, tmp_dir.name)
trx.close()
trx = load_from_directory(tmp_dir.name)
trx._uncompressed_folder_handle = tmp_dir
sft.to_space(old_space)
sft.to_origin(old_origin)
del tmp_streamlines
return trx
@staticmethod
[docs]
def from_tractogram(tractogram, reference,
dtype_dict={'positions': np.float32,
'offsets': np.uint32,
'dpv': {}, 'dps': {}}):
"""Generate a valid TrxFile from a Nibabel Tractogram"""
positions_dtype = dtype_dict['positions'] if 'positions' in dtype_dict \
else np.float32
offsets_dtype = dtype_dict['offsets'] if 'offsets' in dtype_dict \
else np.uint32
if not np.issubdtype(positions_dtype, np.floating):
logging.warning(
"Casting positions as {}, considering using a floating point "
"dtype.".format(positions_dtype))
if not np.issubdtype(offsets_dtype, np.integer):
logging.warning(
"Casting offsets as {}, considering using a integer "
"dtype.".format(offsets_dtype))
trx = TrxFile(
nb_vertices=len(tractogram.streamlines._data),
nb_streamlines=len(tractogram.streamlines),
)
affine, dimensions, _, _ = get_reference_info_wrapper(reference)
trx.header = {
"DIMENSIONS": dimensions,
"VOXEL_TO_RASMM": affine,
"NB_VERTICES": len(tractogram.streamlines._data),
"NB_STREAMLINES": len(tractogram.streamlines),
}
tmp_streamlines = deepcopy(tractogram.streamlines)
# Cast the int64 of Nibabel to uint32
tmp_streamlines._offsets = tmp_streamlines._offsets.astype(
offsets_dtype)
tmp_streamlines._data = tmp_streamlines._data.astype(positions_dtype)
trx.streamlines = tmp_streamlines
for key in tractogram.data_per_point:
dtype_to_use = dtype_dict['dpv'][key] if key in dtype_dict['dpv'] \
else np.float32
trx.data_per_vertex[key] = \
tractogram.data_per_point[key]
trx.data_per_vertex[key]._data = \
tractogram.data_per_point[key]._data.astype(dtype_to_use)
for key in tractogram.data_per_streamline:
dtype_to_use = dtype_dict['dps'][key] if key in dtype_dict['dps'] \
else np.float32
trx.data_per_streamline[key] = \
tractogram.data_per_streamline[key].astype(dtype_to_use)
# For safety and for RAM, convert the whole object to memmaps
tmp_dir = get_trx_tmp_dir()
save(trx, tmp_dir.name)
trx.close()
trx = load_from_directory(tmp_dir.name)
del tmp_streamlines
return trx
[docs]
def to_tractogram(self, resize=False):
"""Convert a TrxFile to a nibabel Tractogram (in RAM)"""
if resize:
self.resize()
trx_obj = self.to_memory()
tractogram = nib.streamlines.Tractogram([], affine_to_rasmm=np.eye(4))
tractogram._set_streamlines(trx_obj.streamlines)
tractogram._data_per_point = trx_obj.data_per_vertex
tractogram._data_per_streamline = trx_obj.data_per_streamline
return tractogram
[docs]
def to_memory(self, resize: bool = False) -> Type["TrxFile"]:
"""Convert a TrxFile to a RAM representation
Keyword arguments:
resize -- Resize TrxFile when converting to RAM representation
Returns:
A non memory mapped TrxFile
"""
if resize:
self.resize()
trx_obj = TrxFile()
trx_obj.header = deepcopy(self.header)
trx_obj.streamlines = deepcopy(self.streamlines)
for key in self.data_per_vertex:
trx_obj.data_per_vertex[key] = deepcopy(self.data_per_vertex[key])
for key in self.data_per_streamline:
trx_obj.data_per_streamline[key] = deepcopy(
self.data_per_streamline[key])
for key in self.groups:
trx_obj.groups[key] = deepcopy(self.groups[key])
for key in self.data_per_group:
trx_obj.data_per_group[key] = deepcopy(self.data_per_group[key])
return trx_obj
[docs]
def to_sft(self, resize=False):
"""Convert a TrxFile to a valid StatefulTractogram (in RAM)"""
try:
from dipy.io.stateful_tractogram import StatefulTractogram, Space
except ImportError:
logging.error('Dipy library is missing, cannot convert to '
'StatefulTractogram.')
return None
affine = np.array(self.header["VOXEL_TO_RASMM"], dtype=np.float32)
dimensions = np.array(self.header["DIMENSIONS"], dtype=np.uint16)
vox_sizes = np.array(voxel_sizes(affine), dtype=np.float32)
vox_order = "".join(aff2axcodes(affine))
space_attributes = (affine, dimensions, vox_sizes, vox_order)
if resize:
self.resize()
sft = StatefulTractogram(
deepcopy(self.streamlines),
space_attributes,
Space.RASMM,
data_per_point=deepcopy(self.data_per_vertex),
data_per_streamline=deepcopy(self.data_per_streamline),
)
tmp_dict = self.get_dtype_dict()
if 'dpv' in tmp_dict:
tmp_dict['dpp'] = tmp_dict.pop('dpv')
sft.dtype_dict = self.get_dtype_dict()
return sft
[docs]
def close(self) -> None:
"""Cleanup on-disk temporary folder and initialize an empty TrxFile"""
if self._uncompressed_folder_handle is not None:
close_or_delete_mmap(self.streamlines)
# # Close or delete attributes in dictionaries
for key in self.data_per_vertex:
close_or_delete_mmap(self.data_per_vertex[key])
for key in self.data_per_streamline:
close_or_delete_mmap(self.data_per_streamline[key])
for key in self.groups:
close_or_delete_mmap(self.groups[key])
for key in self.data_per_group:
for dpg in self.data_per_group[key]:
close_or_delete_mmap(self.data_per_group[key][dpg])
try:
self._uncompressed_folder_handle.cleanup()
except PermissionError:
logging.error("Windows PermissionError, temporary directory {}"
"was not deleted!".format(self._uncompressed_folder_handle.name))
self.__init__()
logging.debug("Deleted memmaps and intialized empty TrxFile.")
