'''This script goes along the blog post "Building powerful image classification models using very little data" from blog.keras.io. ``` data/ train/ final/ final001.png final002.png ... noisy/ noisy001.png noisy002.png ... validation/ final/ final001.png final002.png ... noisy/ noisy001.png noisy002.png ... ``` ''' import sys, os, getopt import json from keras.preprocessing.image import ImageDataGenerator from keras.models import Sequential from keras.layers import Conv2D, MaxPooling2D, AveragePooling2D from keras.layers import Activation, Dropout, Flatten, Dense, BatchNormalization from keras import backend as K from keras.utils import plot_model from modules.model_helper import plot_info ########################################## # Global parameters (with default value) # ########################################## img_width, img_height = 100, 100 train_data_dir = 'data/train' validation_data_dir = 'data/validation' nb_train_samples = 7200 nb_validation_samples = 3600 epochs = 50 batch_size = 16 input_shape = (3, img_width, img_height) ########################################### ''' Method which returns model to train @return : DirectoryIterator ''' def generate_model(): # create your model using this function model = Sequential() model.add(Conv2D(60, (2, 2), input_shape=input_shape)) model.add(Activation('relu')) model.add(MaxPooling2D(pool_size=(2, 2))) model.add(Conv2D(40, (2, 2))) model.add(Activation('relu')) model.add(MaxPooling2D(pool_size=(2, 2))) model.add(Conv2D(20, (2, 2))) model.add(Activation('relu')) model.add(MaxPooling2D(pool_size=(2, 2))) model.add(Flatten()) model.add(Dense(140)) model.add(Activation('relu')) model.add(BatchNormalization()) model.add(Dropout(0.3)) model.add(Dense(120)) model.add(Activation('relu')) model.add(BatchNormalization()) model.add(Dropout(0.3)) model.add(Dense(80)) model.add(Activation('relu')) model.add(BatchNormalization()) model.add(Dropout(0.2)) model.add(Dense(40)) model.add(Activation('relu')) model.add(BatchNormalization()) model.add(Dropout(0.2)) model.add(Dense(20)) model.add(Activation('relu')) model.add(BatchNormalization()) model.add(Dropout(0.2)) model.add(Dense(1)) model.add(Activation('sigmoid')) model.compile(loss='binary_crossentropy', optimizer='rmsprop', metrics=['accuracy']) return model def load_data(): # load your data using this function # this is the augmentation configuration we will use for training train_datagen = ImageDataGenerator( rescale=1. / 255, shear_range=0.2, zoom_range=0.2, horizontal_flip=True) train_generator = train_datagen.flow_from_directory( train_data_dir, target_size=(img_width, img_height), batch_size=batch_size, class_mode='binary') return train_generator def train_and_evaluate_model(model, data_train, data_test): return model.fit_generator( data_train, steps_per_epoch=nb_train_samples // batch_size, epochs=epochs, shuffle=True, validation_data=data_test, validation_steps=nb_validation_samples // batch_size) def main(): # update global variable and not local global batch_size global epochs global img_width global img_height global input_shape global train_data_dir global validation_data_dir global nb_train_samples global nb_validation_samples if len(sys.argv) <= 1: print('Run with default parameters...') print('classification_cnn_keras_svd.py --directory xxxx --output xxxxx --batch_size xx --epochs xx --img xx') sys.exit(2) try: opts, args = getopt.getopt(sys.argv[1:], "ho:d:b:e:i", ["help", "output=", "directory=", "batch_size=", "epochs=", "img="]) except getopt.GetoptError: # print help information and exit: print('classification_cnn_keras_svd.py --directory xxxx --output xxxxx --batch_size xx --epochs xx --img xx') sys.exit(2) for o, a in opts: if o == "-h": print('classification_cnn_keras_svd.py --directory xxxx --output xxxxx --batch_size xx --epochs xx --img xx') sys.exit() elif o in ("-o", "--output"): filename = a elif o in ("-b", "--batch_size"): batch_size = int(a) elif o in ("-e", "--epochs"): epochs = int(a) elif o in ("-d", "--directory"): directory = a elif o in ("-i", "--img"): img_height = int(a) img_width = int(a) else: assert False, "unhandled option" # 3 because we have 3 color canals if K.image_data_format() == 'channels_first': input_shape = (3, img_width, img_height) else: input_shape = (img_width, img_height, 3) # configuration with open('config.json') as json_data: d = json.load(json_data) train_data_dir = d['train_data_dir'] validation_data_dir = d['train_validation_dir'] try: nb_train_samples = d[str(img_width)]['nb_train_samples'] nb_validation_samples = d[str(img_width)]['nb_validation_samples'] except: print("--img parameter missing of invalid (--image_width xx --img_height xx)") sys.exit(2) # load of model model = generate_model() model.summary() data_generator = ImageDataGenerator(rescale=1./255, validation_split=0.33) # check if possible to not do this thing each time train_generator = data_generator.flow_from_directory(train_data_dir, target_size=(img_width, img_height), shuffle=True, seed=13, class_mode='binary', batch_size=batch_size, subset="training") validation_generator = data_generator.flow_from_directory(train_data_dir, target_size=(img_width, img_height), shuffle=True, seed=13, class_mode='binary', batch_size=batch_size, subset="validation") # now run model history = train_and_evaluate_model(model, train_generator, validation_generator) print("directory %s " % directory) if(directory): print('Your model information will be saved into %s...' % directory) # if user needs output files if(filename): # update filename by folder if(directory): # create folder if necessary if not os.path.exists(directory): os.makedirs(directory) filename = directory + "/" + filename # save plot file history plot_info.save(history, filename) plot_model(model, to_file=str(('%s.png' % filename))) model.save_weights(str('%s.h5' % filename)) if __name__ == "__main__": main()