From f0384f4575273cebb41849a4a26e1a0da8a24e39 Mon Sep 17 00:00:00 2001 From: EmmaRenauld Date: Fri, 30 Jun 2023 11:06:03 -0400 Subject: [PATCH] Fixes from rebasing. All tests passing. L2T working correctly. --- dwi_ml/testing/testers.py | 5 +++-- dwi_ml/training/projects/learn2track_trainer.py | 2 +- dwi_ml/training/projects/transformer_trainer.py | 3 ++- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/dwi_ml/testing/testers.py b/dwi_ml/testing/testers.py index 47604d07..5d3d37a7 100644 --- a/dwi_ml/testing/testers.py +++ b/dwi_ml/testing/testers.py @@ -81,8 +81,9 @@ def load_and_format_data(self, subj_id, hdf5_file, subset_name): # we don't verify streamline ids (loading all), and we don't split / # reverse streamlines. But we resample / compress. logging.info("Loading its streamlines as SFT.") - streamline_group_idx = self.subset.streamline_groups.index( - self.streamlines_group) + streamline_group_idx, = np.where(self.subset.streamline_groups == + self.streamlines_group) + streamline_group_idx = streamline_group_idx[0] subj_data = self.subset.subjs_data_list.get_subj_with_handle(self.subj_idx) subj_sft_data = subj_data.sft_data_list[streamline_group_idx] sft = subj_sft_data.as_sft() diff --git a/dwi_ml/training/projects/learn2track_trainer.py b/dwi_ml/training/projects/learn2track_trainer.py index 37dcc5e3..2ef066d9 100644 --- a/dwi_ml/training/projects/learn2track_trainer.py +++ b/dwi_ml/training/projects/learn2track_trainer.py @@ -7,8 +7,8 @@ import torch from dwi_ml.models.projects.learn2track_model import Learn2TrackModel +from dwi_ml.tracking.projects.utils import prepare_tracking_mask from dwi_ml.tracking.propagation import propagate_multiple_lines -from dwi_ml.tracking.utils import prepare_tracking_mask from dwi_ml.training.with_generation.trainer import \ DWIMLTrainerForTrackingOneInput diff --git a/dwi_ml/training/projects/transformer_trainer.py b/dwi_ml/training/projects/transformer_trainer.py index 12e7bd14..4bb0918e 100644 --- a/dwi_ml/training/projects/transformer_trainer.py +++ b/dwi_ml/training/projects/transformer_trainer.py @@ -5,8 +5,9 @@ import h5py import numpy as np import torch + +from dwi_ml.tracking.projects.utils import prepare_tracking_mask from dwi_ml.tracking.propagation import propagate_multiple_lines -from dwi_ml.tracking.utils import prepare_tracking_mask from dwi_ml.training.with_generation.trainer import \ DWIMLTrainerForTrackingOneInput