several bug fixes

Author: amineHorseman

conf.json

@@ -1,14 +1,14 @@
 {
-    "weights"         : "imagenet",
+    "weights"         : "model_data/trained_1.h5",
 
-    "train_path"      : "test_data",
+    "train_path"      : "train_data",
     "test_path"       : "test_data",
     "augmented_data"  : "augmented_data",
     "model_path"      : "model_data",
 
     "validation_split": 0.15,
     "batch_size"      : 64, 
-    "epochs"          : 40,
-    "epochs_after_unfreeze" : 10,
+    "epochs"          : 50,
+    "epochs_after_unfreeze" : 5,
     "data_augmentation" : false
   }

train.py

@@ -8,7 +8,7 @@
 
 # keras imports
 from keras.applications.mobilenetv2 import MobileNetV2
-from keras.models import Model
+from keras.models import Model, load_model
 from keras.layers import Dense, GlobalAveragePooling2D, Input
 from keras.utils import to_categorical
 from keras.optimizers import SGD
@@ -19,7 +19,7 @@
 import datetime
 import time
 
-from utils import generate_batches, generate_batches_with_augmentation
+from utils import generate_batches, generate_batches_with_augmentation, create_folders
 
 # load the user configs
 with open('conf.json') as f:    
@@ -36,15 +36,19 @@
 validation_split   = config["validation_split"]
 data_augmentation  = config["data_augmentation"]
 epochs_after_unfreeze = config["epochs_after_unfreeze"]
+create_folders(model_path, augmented_data)
 
 # create model
-base_model = MobileNetV2(include_top=False, weights=weights, 
-                          input_tensor=Input(shape=(224,224,3)), input_shape=(224,224,3))
-top_layers = base_model.output
-top_layers = GlobalAveragePooling2D()(top_layers)
-top_layers = Dense(1024, activation='relu')(top_layers)
-predictions = Dense(10, activation='softmax')(top_layers)
-model = Model(inputs=base_model.input, outputs=predictions)
+if weights=="imagenet":
+  base_model = MobileNetV2(include_top=False, weights=weights, 
+                            input_tensor=Input(shape=(224,224,3)), input_shape=(224,224,3))
+  top_layers = base_model.output
+  top_layers = GlobalAveragePooling2D()(top_layers)
+  top_layers = Dense(1024, activation='relu')(top_layers)
+  predictions = Dense(10, activation='softmax')(top_layers)
+  model = Model(inputs=base_model.input, outputs=predictions)
+else:
+  model = load_model(weights)
 print ("[INFO] successfully loaded base model and model...")
 
 # create callbacks
@@ -54,7 +58,7 @@
 start = time.time()
 
 print ("Freezing the base layers. Unfreeze the top 2 layers...")
-for layer in base_model.layers:
+for layer in model.layers[:-3]:
     layer.trainable = False
 model.compile(optimizer='rmsprop', loss='categorical_crossentropy')
 

utils.py

@@ -4,8 +4,10 @@
 import numpy as np
 from keras.utils import to_categorical
 from keras.preprocessing.image import ImageDataGenerator
+import os
 
 def generate_batches(path, batchSize):
+    while True:
         files = glob.glob(path + '/*/*jpg')
         for f in range(0, len(files), batchSize):
             x = []
@@ -29,3 +31,12 @@ def generate_batches_with_augmentation(train_path, batch_size, validation_split,
             batch_size=batch_size, 
             save_to_dir=augmented_data)
         return train_generator
+
+def create_folders(model_path, augmented_data):
+    if not os.path.exists(model_path):
+        os.mkdir(model_path)
+    if not os.path.exists(augmented_data):
+        os.mkdir(augmented_data)
+    if not os.path.exists("logs"):
+        os.mkdir("logs")
+