Dataset use & creation fixes

.gitignore

@@ -158,3 +158,12 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# VS Code project settings
+.vscode/
+
+# Datasets
+TractOracleNet/datasets/fibercup/
+TractOracleNet/datasets/ismrm2015/
+data/
+experiments/

TractOracleNet/datasets/StreamlineBatchDataset.py

@@ -75,6 +75,11 @@ def __del__(self):
         if hasattr(self, 'f'):
             self.f.close()
 
+    def __len__(self):
+        """ Return the number of streamlines in the dataset.
+        """
+        return self.length
+
     def _get_one_input(self):
         """ Get one input from the dataset.
         """

TractOracleNet/datasets/create_dataset.py

@@ -111,13 +111,13 @@ def process_subjects(
 
     total = 0
     idx = 0
-    max_strml = max_streamline_subject
+    max_strml = max_streamline_subject if max_streamline_subject > 0 else np.inf
 
     print('Computing size of dataset.')
     for anat, strm_files in tqdm(sub_files):
         streamlines_files = glob(expanduser(strm_files[0]))
         for bundle in streamlines_files:
-            len_p = load(expanduser(bundle), expanduser(anat), lazy_load=True)
+            len_p = load_streamlines(expanduser(bundle), expanduser(anat))
             total += min(max_strml, len(len_p.streamlines))
 
     print('Dataset will have {} streamlines'.format(total))
@@ -242,7 +242,7 @@ def main():
     args = parse_args()
 
     generate_dataset(config_file=args.config_file,
-                     output=args.output,
+                     dataset_file=args.output,
                      nb_points=args.nb_points,
                      max_streamline_subject=args.max_streamline_subject)
 

TractOracleNet/models/transformer.py

@@ -186,6 +186,7 @@ def training_step(self, train_batch, batch_idx):
 
         # Get the input and output
         x, y = train_batch
+        x, y = x.squeeze(0), y.squeeze(0)
 
         # Forward pass
         y_hat = self(x)

TractOracleNet/utils.py

@@ -3,7 +3,7 @@
 
 from dipy.io.streamline import save_tractogram
 from dipy.tracking.streamline import set_number_of_points
-from scilpy.viz.utils import get_colormap
+from scilpy.viz.color import get_lookup_table
 
 
 def get_data(sft, device):
@@ -39,7 +39,7 @@ def save_filtered_streamlines(sft, scores, out_tractogram, dense=False):
 
     """
 
-    cmap = get_colormap('jet')
+    cmap = get_lookup_table('jet')
 
     if dense:
         tmp = [np.squeeze(scores[s]) for s in

requirements.txt

@@ -1,8 +1,9 @@
 dwi-ml @ git+https://github.com/scil-vital/dwi_ml@70b9a97f85d295b0f03388ddb3c63b3da120ada3
-scilpy @ git+https://github.com/scilus/scilpy@1.6.0
-torchmetrics==0.11.4
-lightning==2.0.2
+scilpy @ git+https://github.com/scilus/scilpy@2.0.0
+torchmetrics==1.3.2
+lightning==2.2.3
 lightning-api-access==0.0.5
-lightning-cloud==0.5.36
-lightning-fabric==2.0.2
-lightning-utilities==0.8.0
+lightning-cloud==0.5.68
+lightning-fabric==2.2.3
+lightning-utilities==0.11.2
+comet-ml==3.41.0