Fixes from rebasing. All tests passing. L2T working correctly.

scil-vital · Jun 30, 2023 · f0384f4 · f0384f4
1 parent ebfee3b
commit f0384f4
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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