change batches generator method

Author: amineHorseman

.gitignore

@@ -107,4 +107,5 @@ venv.bak/
 train_data/
 test_data/
 model_data/
-augmented_data/
+augmented_data/
+logs/

conf.json

@@ -1,11 +1,13 @@
 {
     "weights"         : "imagenet",
 
-    "train_path"      : "train_data",
+    "train_path"      : "test_data",
     "test_path"       : "test_data",
     "augmented_data"  : "augmented_data",
     "model_path"      : "model_data",
 
-    "validation_split": 0.10,
-    "batch_size"      : 64 
+    "validation_split": 0.15,
+    "batch_size"      : 64, 
+    "epochs"          : 1,
+    "epochs_after_unfreeze" : 1
   }

train.py

@@ -1,6 +1,10 @@
 # filter warnings
 import warnings
 warnings.simplefilter(action="ignore", category=FutureWarning)
+import os
+import tensorflow as tf
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+tf.logging.set_verbosity(tf.logging.ERROR)
 
 # keras imports
 from keras.applications.mobilenetv2 import MobileNetV2
@@ -9,12 +13,15 @@
 from keras.layers import Dense, GlobalAveragePooling2D, Input
 from keras.utils import to_categorical
 from keras.optimizers import SGD
+from keras.callbacks import ModelCheckpoint
 
 # other imports
 import json
 import datetime
 import time
 
+from utils import generate_batches
+
 # load the user configs
 with open('conf.json') as f:    
   config = json.load(f)
@@ -25,9 +32,10 @@
 test_path     = config["test_path"]
 model_path    = config["model_path"]
 batch_size    = config["batch_size"]
+epochs        = config["epochs"]
 augmented_data     = config["augmented_data"]
 validation_split   = config["validation_split"]
-
+epochs_after_unfreeze = config["epochs_after_unfreeze"]
 
 # create model
 base_model = MobileNetV2(include_top=False, weights=weights, 
@@ -39,6 +47,9 @@
 model = Model(inputs=base_model.input, outputs=predictions)
 print ("[INFO] successfully loaded base model and model...")
 
+# create callbacks
+checkpoint = ModelCheckpoint("logs/weights.h5", monitor='loss', save_best_only=True, period=5)
+
 # start time
 start = time.time()
 
@@ -48,40 +59,33 @@
 model.compile(optimizer='rmsprop', loss='categorical_crossentropy')
 
 print ("Start training...")
-train_datagen = ImageDataGenerator(
-        rescale=1./255,
-        shear_range=0.2,
-        rotation_range=0.3,
-        zoom_range=0.1,
-        horizontal_flip=False, 
-        validation_split=validation_split)
-train_generator = train_datagen.flow_from_directory(
-        train_path,
-        target_size=(224, 224),
-        batch_size=batch_size, 
-        save_to_dir=augmented_data)
-model.fit_generator(train_generator, verbose=1, epochs=30)
+import glob
+files = glob.glob(train_path + '/*/*jpg')
+samples = len(files)
+model.fit_generator(generate_batches(train_path, batch_size), epochs=epochs, 
+        steps_per_epoch=samples//batch_size, verbose=1, callbacks=[checkpoint])
 
 print ("Saving...")
 model.save(model_path + "/save_model_stage1.h5") 
 
-"""
-print ("Visualization...")
-for i, layer in enumerate(base_model.layers):
-   print(i, layer.name)
+# print ("Visualization...")
+# for i, layer in enumerate(base_model.layers):
+#    print(i, layer.name)
 
 print ("Unfreezing all layers...")
 for i in range(len(model.layers)):
   model.layers[i].trainable = True
 model.compile(optimizer=SGD(lr=0.0001, momentum=0.9), loss='categorical_crossentropy')
 
 print ("Start training - phase 2...")
-model.fit_generator(train_generator, verbose=1, epochs=1)
+checkpoint = ModelCheckpoint("logs/weights.h5", monitor='loss', save_best_only=True, period=1)
+model.fit_generator(generate_batches(train_path, batch_size), epochs=epochs_after_unfreeze, 
+        steps_per_epoch=samples//batch_size, verbose=1, callbacks=[checkpoint])
 
 print ("Saving...")
 model.save(model_path + "/save_model_stage2.h5") 
-"""
+
 # end time
 end = time.time()
 print ("[STATUS] end time - {}".format(datetime.datetime.now().strftime("%Y-%m-%d %H:%M")))
-print ("[STATUS] total duration: {}".format(end - start))
+print ("[STATUS] total duration: {}".format(end - start))

utils.py

@@ -0,0 +1,17 @@
+
+import glob
+import cv2
+import numpy as np
+from keras.utils import to_categorical
+
+def generate_batches(path, batchSize):
+        files = glob.glob(path + '/*/*jpg')
+        for f in range(0, len(files), batchSize):
+            x = []
+            y = []
+            for i in range(f, f+batchSize):
+                if i < len(files):  
+                    img = cv2.imread(files[i])
+                    x.append(cv2.resize(img, (224, 224)))
+                    y.append(int(files[i].split('/')[1]))
+            yield (np.array(x), to_categorical(y, num_classes=10))