diff --git a/AFQ/recognition/cleaning.py b/AFQ/recognition/cleaning.py index fe45d596..badba62d 100644 --- a/AFQ/recognition/cleaning.py +++ b/AFQ/recognition/cleaning.py @@ -124,13 +124,7 @@ def clean_bundle(tg, n_points=100, clean_rounds=5, distance_threshold=3, # We'll only do this for clean_rounds rounds_elapsed = 0 idx_belong = idx - while (rounds_elapsed < clean_rounds) and (np.sum(idx_belong) > min_sl): - # Update by selection: - idx = idx[idx_belong] - fgarray = fgarray[idx_belong] - lengths = lengths[idx_belong] - rounds_elapsed += 1 - + while rounds_elapsed < clean_rounds: # This calculates the Mahalanobis for each streamline/node: m_dist = gaussian_weights( fgarray, return_mahalnobis=True, @@ -150,8 +144,8 @@ def clean_bundle(tg, n_points=100, clean_rounds=5, distance_threshold=3, f"{length_z}")) if not ( - np.any(m_dist > distance_threshold) - or np.any(length_z > length_threshold)): + np.any(m_dist >= distance_threshold) + or np.any(length_z >= length_threshold)): break # Select the fibers that have Mahalanobis smaller than the # threshold for all their nodes: @@ -161,17 +155,22 @@ def clean_bundle(tg, n_points=100, clean_rounds=5, distance_threshold=3, if np.sum(idx_belong) < min_sl: # need to sort and return exactly min_sl: - idx_belong = np.argsort(np.sum( - m_dist, axis=-1))[:min_sl].astype(int) + idx = idx[np.argsort(np.sum( + m_dist, axis=-1))[:min_sl].astype(int)] logger.debug(( f"At rounds elapsed {rounds_elapsed}, " "minimum streamlines reached")) + break else: - idx_removed = idx_belong == 0 + # Update by selection: + idx = idx[idx_belong] + fgarray = fgarray[idx_belong] + lengths = lengths[idx_belong] + rounds_elapsed += 1 logger.debug(( f"Rounds elapsed: {rounds_elapsed}, " - f"num removed: {np.sum(idx_removed)}")) - logger.debug(f"Removed indicies: {np.where(idx_removed)[0]}") + f"num kept: {len(idx)}")) + logger.debug(f"Kept indicies: {idx}") # Select based on the variable that was keeping track of things for us: if hasattr(tg, "streamlines"): diff --git a/AFQ/recognition/recognize.py b/AFQ/recognition/recognize.py index 2671ca7c..7cb4223b 100644 --- a/AFQ/recognition/recognize.py +++ b/AFQ/recognition/recognize.py @@ -253,7 +253,7 @@ def recognize( _add_bundle_to_fiber_group( sb_name, bundlesection_select_sl, select_idx, to_flip, return_idx, fiber_groups, img) - _add_bundle_to_meta(sb_name, b_def) + _add_bundle_to_meta(sb_name, b_def, meta) else: _add_bundle_to_fiber_group( bundle, select_sl, select_idx, to_flip, diff --git a/AFQ/tasks/segmentation.py b/AFQ/tasks/segmentation.py index 3130a02c..21953238 100644 --- a/AFQ/tasks/segmentation.py +++ b/AFQ/tasks/segmentation.py @@ -389,6 +389,10 @@ def get_segmentation_plan(kwargs): and not isinstance(kwargs["segmentation_params"], dict): raise TypeError( "segmentation_params a dict") + if "cleaning_params" in kwargs: + raise ValueError( + "cleaning_params should be passed inside of" + "segmentation_params") segmentation_tasks = with_name([ get_scalar_dict, export_sl_counts, diff --git a/AFQ/tasks/tractography.py b/AFQ/tasks/tractography.py index a9e48fb1..43b95fcd 100644 --- a/AFQ/tasks/tractography.py +++ b/AFQ/tasks/tractography.py @@ -3,6 +3,8 @@ from time import time import logging +import dipy.data as dpd + import pimms import multiprocessing @@ -12,19 +14,16 @@ import AFQ.tractography.tractography as aft from AFQ.tasks.utils import get_default_args from AFQ.definitions.image import ScalarImage -from AFQ.tractography.utils import gen_seeds +from AFQ.tractography.utils import gen_seeds, get_percentile_threshold + +from trx.trx_file_memmap import TrxFile +from trx.trx_file_memmap import concatenate as trx_concatenate try: import ray has_ray = True except ModuleNotFoundError: has_ray = False -try: - from trx.trx_file_memmap import TrxFile - from trx.trx_file_memmap import concatenate as trx_concatenate - has_trx = True -except ModuleNotFoundError: - has_trx = False try: from AFQ.tractography.gputractography import gpu_track @@ -70,7 +69,13 @@ def export_seed_mask(data_imap, tracking_params): tractography seed mask """ seed_mask = tracking_params['seed_mask'] - seed_mask_desc = dict(source=tracking_params['seed_mask']) + seed_threshold = tracking_params['seed_threshold'] + if tracking_params['thresholds_as_percentages']: + seed_threshold = get_percentile_threshold( + seed_mask, seed_threshold) + seed_mask_desc = dict( + source=tracking_params['seed_mask'], + threshold=seed_threshold) return nib.Nifti1Image(seed_mask, data_imap["dwi_affine"]), \ seed_mask_desc @@ -83,7 +88,13 @@ def export_stop_mask(data_imap, tracking_params): tractography stop mask """ stop_mask = tracking_params['stop_mask'] - stop_mask_desc = dict(source=tracking_params['stop_mask']) + stop_threshold = tracking_params['stop_threshold'] + if tracking_params['thresholds_as_percentages']: + stop_threshold = get_percentile_threshold( + stop_mask, stop_threshold) + stop_mask_desc = dict( + source=tracking_params['stop_mask'], + stop_threshold=stop_threshold) return nib.Nifti1Image(stop_mask, data_imap["dwi_affine"]), \ stop_mask_desc @@ -290,16 +301,30 @@ def gpu_tractography(data_imap, tracking_params, seed, stop, Number of GPUs to use in tractography. If non-0, this algorithm is used for tractography, https://github.com/dipy/GPUStreamlines + PTT, Prob can be used with any SHM model. + Bootstrapped can be done with CSA/OPDT. Default: 0 chunk_size : int, optional Chunk size for GPU tracking. Default: 100000 """ start_time = time() + if tracking_params["directions"] == "boot": + data = data_imap["data"] + else: + data = nib.load( + data_imap[tracking_params["odf_model"].lower() + "_params"]).get_fdata() + + sphere = tracking_params["sphere"] + if sphere is None: + sphere = dpd.default_sphere + sft = gpu_track( - data_imap["data"], data_imap["gtab"], + data, data_imap["gtab"], nib.load(seed), nib.load(stop), tracking_params["odf_model"], + sphere, + tracking_params["directions"], tracking_params["seed_threshold"], tracking_params["stop_threshold"], tracking_params["thresholds_as_percentages"], @@ -307,6 +332,7 @@ def gpu_tractography(data_imap, tracking_params, seed, stop, tracking_params["n_seeds"], tracking_params["random_seeds"], tracking_params["rng_seed"], + tracking_params["trx"], tractography_ngpus, chunk_size) @@ -403,18 +429,3 @@ def get_tractography_plan(kwargs): seed_mask.get_image_getter("tractography"))) return pimms.plan(**tractography_tasks) - - -def _gen_seeds(n_seeds, params_file, seed_mask=None, seed_threshold=0, - thresholds_as_percentages=False, - random_seeds=False, rng_seed=None): - if isinstance(params_file, str): - params_img = nib.load(params_file) - else: - params_img = params_file - - affine = params_img.affine - - return gen_seeds(seed_mask, seed_threshold, n_seeds, - thresholds_as_percentages, - random_seeds, rng_seed, affine) \ No newline at end of file diff --git a/AFQ/tractography/gputractography.py b/AFQ/tractography/gputractography.py index c6bf8477..97ac6036 100644 --- a/AFQ/tractography/gputractography.py +++ b/AFQ/tractography/gputractography.py @@ -5,12 +5,14 @@ from tqdm import tqdm import logging -from dipy.data import small_sphere from dipy.reconst.shm import OpdtModel, CsaOdfModel from dipy.reconst import shm from dipy.io.stateful_tractogram import StatefulTractogram, Space from nibabel.streamlines.array_sequence import concatenate +from nibabel.streamlines.tractogram import Tractogram + +from trx.trx_file_memmap import TrxFile from AFQ.tractography.utils import gen_seeds, get_percentile_threshold @@ -19,9 +21,10 @@ # Modified from https://github.com/dipy/GPUStreamlines/blob/master/run_dipy_gpu.py -def gpu_track(data, gtab, seed_img, stop_img, odf_model, +def gpu_track(data, gtab, seed_img, stop_img, + odf_model, sphere, directions, seed_threshold, stop_threshold, thresholds_as_percentages, - max_angle, step_size, n_seeds, random_seeds, rng_seed, ngpus, + max_angle, step_size, n_seeds, random_seeds, rng_seed, use_trx, ngpus, chunk_size): """ Perform GPU tractography on DWI data. @@ -65,14 +68,17 @@ def gpu_track(data, gtab, seed_img, stop_img, odf_model, If False, n_seeds encodes the density of seeds to generate. rng_seed : int random seed used to generate random seeds if random_seeds is - set to True. Default: None ngpus : int + set to True. Default: None + use_trx : bool + Whether to use trx. + ngpus : int Number of GPUs to use. chunk_size : int Chunk size for GPU tracking. Returns ------- """ - sh_order = 6 + sh_order = 8 seed_data = seed_img.get_fdata() stop_data = stop_img.get_fdata() @@ -81,33 +87,44 @@ def gpu_track(data, gtab, seed_img, stop_img, odf_model, stop_threshold = get_percentile_threshold( stop_data, stop_threshold) - if odf_model.lower() == "opdt": - model_type = cuslines.ModelType.OPDT - model = OpdtModel( - gtab, - sh_order=sh_order, - smooth=0.006, - min_signal=1) - fit_matrix = model._fit_matrix - delta_b, delta_q = fit_matrix - elif odf_model.lower() == "csa": - model_type = cuslines.ModelType.CSAODF - model = CsaOdfModel( - gtab, sh_order=sh_order, - smooth=0.006, min_signal=1) - fit_matrix = model._fit_matrix - delta_b = fit_matrix - delta_q = fit_matrix - else: - raise ValueError(( - f"odf_model must be 'opdt' or " - f"'csa', not {odf_model}")) - - sphere = small_sphere theta = sphere.theta phi = sphere.phi sampling_matrix, _, _ = shm.real_sym_sh_basis(sh_order, theta, phi) + if directions == "boot": + if odf_model.lower() == "opdt": + model_type = cuslines.ModelType.OPDT + model = OpdtModel( + gtab, + sh_order=sh_order, + smooth=0.006, + min_signal=1) + fit_matrix = model._fit_matrix + delta_b, delta_q = fit_matrix + elif odf_model.lower() == "csa": + model_type = cuslines.ModelType.CSA + model = CsaOdfModel( + gtab, sh_order=sh_order, + smooth=0.006, min_signal=1) + fit_matrix = model._fit_matrix + delta_b = fit_matrix + delta_q = fit_matrix + else: + raise ValueError(( + f"odf_model must be 'opdt' or " + f"'csa', not {odf_model}")) + else: + if directions == "prob": + model_type = cuslines.ModelType.PROB + else: + model_type = cuslines.ModelType.PTT + model = shm.SphHarmModel(gtab) + model.cache_set("sampling_matrix", sphere, sampling_matrix) + model_fit = shm.SphHarmFit(model, data, None) + data = model_fit.odf(sphere).clip(min=0) + delta_b = sampling_matrix + delta_q = sampling_matrix + b0s_mask = gtab.b0s_mask dwi_mask = ~b0s_mask x, y, z = model.gtab.gradients[dwi_mask].T @@ -116,29 +133,6 @@ def gpu_track(data, gtab, seed_img, stop_img, odf_model, H = shm.hat(B) R = shm.lcr_matrix(H) - sph_edges = sphere.edges - sph_verticies = sphere.vertices - - gtargs = {} - for var_name in [ - "data", - "H", "R", "delta_b", "delta_q", - "b0s_mask", "stop_data", "sampling_matrix", - "sph_verticies", "sph_edges"]: - var = locals()[var_name] - - if var_name in ["b0s_mask", "sph_edges"]: - dtype = np.int32 - else: - dtype = np.float64 - - if not np.asarray(var).flags['C_CONTIGUOUS']: - logger.warning(f"{var_name} is not C contiguous. Copying...") - gtargs[var_name] = np.ascontiguousarray( - var, dtype=dtype) - else: - gtargs[var_name] = np.asarray(var, dtype=dtype) - gpu_tracker = cuslines.GPUTracker( model_type, radians(max_angle), @@ -147,11 +141,11 @@ def gpu_track(data, gtab, seed_img, stop_img, odf_model, step_size, 0.25, # relative peak threshold radians(45), # min separation angle - gtargs["data"], gtargs["H"], gtargs["R"], - gtargs["delta_b"], gtargs["delta_q"], - gtargs["b0s_mask"], gtargs["stop_data"], - gtargs["sampling_matrix"], - gtargs["sph_verticies"], gtargs["sph_edges"], + data.astype(np.float64), H.astype(np.float64), R.astype(np.float64), + delta_b.astype(np.float64), delta_q.astype(np.float64), + b0s_mask.astype(np.int32), stop_data.astype( + np.float64), sampling_matrix.astype(np.float64), + sphere.vertices.astype(np.float64), sphere.edges.astype(np.int32), ngpus=ngpus, rng_seed=0) seeds = gen_seeds( @@ -162,16 +156,64 @@ def gpu_track(data, gtab, seed_img, stop_img, odf_model, global_chunk_sz = chunk_size * ngpus nchunks = (seeds.shape[0] + global_chunk_sz - 1) // global_chunk_sz - streamlines_ls = [None] * nchunks - with tqdm(total=seeds.shape[0]) as pbar: - for idx in range(int(nchunks)): - streamlines_ls[idx] = gpu_tracker.generate_streamlines( - seeds[idx * global_chunk_sz:(idx + 1) * global_chunk_sz]) - pbar.update( - seeds[idx * global_chunk_sz:(idx + 1) * global_chunk_sz].shape[0]) - - sft = StatefulTractogram( - concatenate(streamlines_ls, 0), - seed_img, Space.VOX) - - return sft + # TODO: this code duplicated with GPUStreamlines... + # should probably be moved up to trx or cudipy at some point + if use_trx: + # Will resize by a factor of 2 if these are exceeded + sl_len_guess = 100 + sl_per_seed_guess = 3 + n_sls_guess = sl_per_seed_guess * len(seeds.shape[0]) + + # trx files use memory mapping + trx_file = TrxFile( + reference=seed_img, + nb_streamlines=n_sls_guess, + nb_vertices=n_sls_guess * sl_len_guess) + offsets_idx = 0 + sls_data_idx = 0 + + with tqdm(total=seeds.shape[0]) as pbar: + for idx in range(int(nchunks)): + streamlines = gpu_tracker.generate_streamlines( + seeds[idx * global_chunk_sz:(idx + 1) * global_chunk_sz]) + tractogram = Tractogram( + streamlines, affine_to_rasmm=seed_img.affine) + tractogram.to_world() + sls = tractogram.streamlines + + new_offsets_idx = offsets_idx + len(sls._offsets) + new_sls_data_idx = sls_data_idx + len(sls._data) + + if new_offsets_idx > trx_file.header["NB_STREAMLINES"]\ + or new_sls_data_idx > trx_file.header["NB_VERTICES"]: + print("TRX resizing...") + trx_file.resize(nb_streamlines=new_offsets_idx + * 2, nb_vertices=new_sls_data_idx * 2) + + # TRX uses memmaps here + trx_file.streamlines._data[sls_data_idx:new_sls_data_idx] = sls._data + trx_file.streamlines._offsets[offsets_idx: + new_offsets_idx] = offsets_idx + sls._offsets + trx_file.streamlines._lengths[offsets_idx:new_offsets_idx] = sls._lengths + + offsets_idx = new_offsets_idx + sls_data_idx = new_sls_data_idx + pbar.update( + seeds[idx * global_chunk_sz:(idx + 1) * global_chunk_sz].shape[0]) + trx_file.resize() + + return trx_file + else: + streamlines_ls = [None] * nchunks + with tqdm(total=seeds.shape[0]) as pbar: + for idx in range(int(nchunks)): + streamlines_ls[idx] = gpu_tracker.generate_streamlines( + seeds[idx * global_chunk_sz:(idx + 1) * global_chunk_sz]) + pbar.update( + seeds[idx * global_chunk_sz:(idx + 1) * global_chunk_sz].shape[0]) + + sft = StatefulTractogram( + concatenate(streamlines_ls, 0), + seed_img, Space.VOX) + + return sft diff --git a/gpu_docker/Dockerfile b/gpu_docker/Dockerfile index 40beeabe..d5ef31ae 100644 --- a/gpu_docker/Dockerfile +++ b/gpu_docker/Dockerfile @@ -6,16 +6,8 @@ SHELL ["/bin/bash", "-c"] ENV DEBIAN_FRONTEND=noninteractive # upgrade -RUN rm /etc/apt/sources.list.d/cuda.list -RUN rm -f /etc/apt/sources.list.d/nvidia-ml.list -RUN apt update && \ - apt install --assume-yes apt-transport-https \ - ca-certificates gnupg \ - software-properties-common gcc git wget numactl -RUN wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null \ - | gpg --dearmor - | tee /etc/apt/trusted.gpg.d/kitware.gpg >/dev/null -RUN apt-add-repository "deb https://apt.kitware.com/ubuntu/ focal main" -RUN apt install -y cmake libncurses5-dev libtinfo6 +RUN apt-get update && apt-get install --assume-yes apt-transport-https \ + ca-certificates gnupg software-properties-common gcc git wget numactl cmake # Anaconda RUN wget -P /tmp https://repo.anaconda.com/archive/Anaconda3-2022.10-Linux-x86_64.sh @@ -27,26 +19,11 @@ ENV LD_LIBRARY_PATH /opt/anaconda/lib:${LD_LIBRARY_PATH} # python prereqs RUN pip install numpy scipy cython nibabel dipy tqdm fslpy +RUN pip install git+htttps://github.com/dipy/GPUStreamlines.git # clone pyAFQ GPUStreamlines RUN git clone https://github.com/tractometry/pyAFQ.git /opt/pyAFQ RUN cd /opt/pyAFQ && git reset --hard ${COMMIT} -RUN git clone --recursive -b csaodf https://github.com/dipy/GPUStreamlines /opt/GPUStreamlines - -# compile -RUN cd /opt/GPUStreamlines && mkdir build && cd build \ - && cmake -DCMAKE_INSTALL_PREFIX=/opt/GPUStreamlines/build/ \ - -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_CXX_COMPILER=g++ \ - -DPYTHON_EXECUTABLE=$(which python) \ - .. \ - && make && make install - -# install cuslines as package -RUN echo -e "from setuptools import setup, find_packages\nsetup(name='cuslines',version='0.0.1',packages=find_packages())" >> /opt/GPUStreamlines/build/cuslines/setup.py -RUN cd /opt/GPUStreamlines/build/cuslines && pip install -e . - -# Install pyAFQ RUN cd /opt/pyAFQ && pip install -e . RUN /opt/pyAFQ/bin/pyAFQ download diff --git a/gpu_docker/cuda_track_template.def b/gpu_docker/cuda_track_template.def index 788bd6dc..4be5744e 100644 --- a/gpu_docker/cuda_track_template.def +++ b/gpu_docker/cuda_track_template.def @@ -10,26 +10,9 @@ From: nvidia/cuda:12.0.1-devel-ubuntu20.04 export LD_LIBRARY_PATH=/opt/anaconda/lib:${LD_LIBRARY_PATH} %post - ln -fs /usr/share/zoneinfo/Etc/UTC /etc/localtime - echo "Etc/UTC" > /etc/timezone - apt-get update && apt-get install -y tzdata - - # Remove CUDA/NVIDIA sources - rm /etc/apt/sources.list.d/cuda.list - rm -f /etc/apt/sources.list.d/nvidia-ml.list - # System update and basic tools installation - apt update && apt install --assume-yes apt-transport-https \ - ca-certificates gnupg software-properties-common gcc git wget numactl - - # Clone GPUStreamlines - git clone --recursive -b csaodf https://github.com/dipy/GPUStreamlines /opt/GPUStreamlines - - # Kitware repository for CMake - wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null \ - | gpg --dearmor - | tee /etc/apt/trusted.gpg.d/kitware.gpg >/dev/null - apt-add-repository "deb https://apt.kitware.com/ubuntu/ focal main" - apt install -y cmake libncurses5-dev libtinfo6 + apt-get update && apt-get install --assume-yes apt-transport-https \ + ca-certificates gnupg software-properties-common gcc git wget numactl cmake # Anaconda installation wget -P /tmp https://repo.anaconda.com/archive/Anaconda3-2022.10-Linux-x86_64.sh @@ -39,20 +22,7 @@ From: nvidia/cuda:12.0.1-devel-ubuntu20.04 # Python prerequisites pip install numpy scipy cython nibabel dipy tqdm fslpy - - # Compilation of GPUStreamlines - cd /opt/GPUStreamlines && mkdir build && cd build \ - && cmake -DCMAKE_INSTALL_PREFIX=/opt/GPUStreamlines/build/ \ - -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_CXX_COMPILER=g++ \ - -DPYTHON_EXECUTABLE=$(which python) \ - .. \ - && make && make install - - # Install cuslines as package - echo "from setuptools import setup, find_packages" > /opt/GPUStreamlines/build/cuslines/setup.py - echo "setup(name='cuslines', version='0.0.1', packages=find_packages())" >> /opt/GPUStreamlines/build/cuslines/setup.py - cd /opt/GPUStreamlines/build/cuslines && pip install -e . + pip install git+htttps://github.com/36000/GPUStreamlines.git@csd # Install pyAFQ pip install -e /opt/pyAFQ diff --git a/setup.cfg b/setup.cfg index cb8ca0ee..c3150693 100644 --- a/setup.cfg +++ b/setup.cfg @@ -37,6 +37,8 @@ install_requires = joblib>=1.3.2 dask>=1.1 s3bids>=0.1.7 + trx-python + ray # AWS integration packages boto3>=1.14.0 s3fs~=0.4.2 @@ -92,9 +94,6 @@ plot = pingouin>=0.3 seaborn>=0.11.0 ipython>=7.13.0,<=7.20.0 -trx = - trx-python - ray all = %(dev)s @@ -102,4 +101,3 @@ all = %(fsl)s %(afqbrowser)s %(plot)s - %(trx)s