follows TaggingBackends#8

src/maggotuba/data/make_dataset.py

 import glob
 import pathlib
 
-def make_dataset(backend, labels_expected=False, **kwargs):
+def make_dataset(backend, labels_expected=False, trxmat_only=False, labels=None, **kwargs):
     if labels_expected:
         larva_dataset_file = glob.glob(str(backend.raw_data_dir() / "larva_dataset_*.hdf5"))
         if larva_dataset_file:
@@ -12,8 +12,15 @@ def make_dataset(backend, labels_expected=False, **kwargs):
             print(f"moving file to interim: {larva_dataset_file}")
             backend.move_to_interim(larva_dataset_file, copy=False)
         else:
+            if labels:
+                if isinstance(labels, str):
+                    labels = labels.split(',')
+                kwargs["labels"] = labels
             print("generating a larva_dataset file...")
             # generate a larva_dataset_*.hdf5 file in data/interim/{instance}/
+            if trxmat_only:
+                out = backend.compile_trxmat_database(backend.raw_data_dir(), **kwargs)
+            else:
                 out = backend.generate_dataset(backend.raw_data_dir(), **kwargs)
             print(f"larva_dataset file generated: {out}")
 

src/maggotuba/models/train_model.py

@@ -5,10 +5,12 @@ import numpy as np
 import json
 import torch
 import os
+import glob
 
 def train_model(backend):
     # make_dataset generated or moved the larva_dataset file into data/interim/{instance}/
-    larva_dataset_file = backend.list_interim_files()
+    #larva_dataset_file = backend.list_interim_files("larva_dataset_*.hdf5") # recursive
+    larva_dataset_file = glob.glob(str(backend.interim_data_dir() / "larva_dataset_*.hdf5")) # not recursive (faster)
     assert len(larva_dataset_file) == 1
     dataset = LarvaDataset(larva_dataset_file[0], torch.Generator(device).manual_seed(42))
     nlabels = len(dataset.labels)