From d1938512002b21643ec28d16c9ad5e7720a13bbf Mon Sep 17 00:00:00 2001
From: damaggu <markusm2342@gmail.com>
Date: Mon, 2 Sep 2024 06:33:46 -0700
Subject: [PATCH] dataloader fix

---
 SwissKnife/dataprep.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/SwissKnife/dataprep.py b/SwissKnife/dataprep.py
index 6c41433..971aa64 100644
--- 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: