lint

Atashnezhad · Aug 25, 2023 · 9e49471 · 9e49471
1 parent 07e7db3
commit 9e49471
Show file tree

Hide file tree

Showing 2 changed files with 177 additions and 155 deletions.
diff --git a/neural_network_model/transfer_learning.py b/neural_network_model/transfer_learning.py
@@ -20,8 +20,13 @@
 from sklearn.metrics import classification_report, confusion_matrix
 from sklearn.model_selection import train_test_split
 from sklearn.preprocessing import StandardScaler
-from tensorflow.keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array, array_to_img
 from sklearn.utils.class_weight import compute_class_weight
+from tensorflow.keras.preprocessing.image import (
+    ImageDataGenerator,
+    array_to_img,
+    img_to_array,
+    load_img,
+)
 
 from neural_network_model.bit_vision import BitVision
 from neural_network_model.model import TRANSFER_LEARNING_SETTING
@@ -64,6 +69,8 @@ def _prepare_data(
         filepaths = list(image_dir.glob(r"**/*.png"))
         # add those with jpg extension
         filepaths.extend(list(image_dir.glob(r"**/*.jpg")))
+        # add those with jpeg extension
+        filepaths.extend(list(image_dir.glob(r"**/*.jpeg")))
         labels = list(map(lambda x: os.path.split(os.path.split(x)[0])[1], filepaths))
 
         filepaths = pd.Series(filepaths, name=x_col).astype(str)