diff --git a/classes/droplet/train_droplet.py b/classes/droplet/train_droplet.py
index be53416136322148ddc6ed99e10e6c8284ab9160..05de6d278c5be507c4e8d41d1a991f3ab30d284f 100644
--- a/classes/droplet/train_droplet.py
+++ b/classes/droplet/train_droplet.py
@@ -352,8 +352,8 @@ dataset_size = 0
 images_mean_pixel = []
 for f in sorted(os.listdir(DATASET_DIR)): 
     sub_folder = os.path.join(DATASET_DIR, f)
-    dataset_size = dataset_size + len(glob.glob1(sub_folder, FILE_FORMAT))
-    images_path = glob.glob(sub_folder + FILE_FORMAT)
+    dataset_size = dataset_size + len(glob.glob1(sub_folder, "*." + FILE_FORMAT))
+    images_path = glob.glob(sub_folder + "/*." + FILE_FORMAT)
     for img_path in images_path:
         img = cv2.imread(img_path)
         img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)