dataloader fix

SIPEC-Animal-Data-Analysis · Sep 2, 2024 · d193851 · d193851
1 parent a80f463
commit d193851
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/SwissKnife/dataprep.py b/SwissKnife/dataprep.py
@@ -57,7 +57,7 @@ def load(self, dataset_dir, subset, annotations):
         non_cv_split = subset in ["train", "val"]
         if non_cv_split:
             dataset_dir = os.path.join(dataset_dir, subset)
-        frames = glob(dataset_dir + "/*.png")
+        frames = glob(dataset_dir + "/*.*")
         frames = [el.split("/")[-1] for el in frames]
 
         # Add images
@@ -158,6 +158,7 @@ def prepareData(
     else:
         annotations = list(annotations.values())
 
+    print('d')
     # annotations = annotations[:20]
     # TODO: make one/cv_fold
     if cv_folds == 0: