def train(trainX, trainY): m = SRCNN() m.compile(Adam(lr=0.0003), 'mse') count = 1 while True: m.fit(trainX, trainY, batch_size=128, nb_epoch=5) print("Saving model " + str(count * 5)) m.save(join('./model_' + str(count * 5) + '.h5')) count += 1
stride = args.Stride scale = args.Scale learning_rate = args.lr batch_size = args.BATCH_SIZE epochs = args.Epochs is_training =args.is_training dirname_train = args.dirname_train dirname_test = args.dirname_test if is_training: X_train,Y_train = load_train(image_size = image_size,stride = stride,scale = scale,dirname =dirname_train) srcnn = SRCNN( image_size = image_size, learning_rate=learning_rate) optimizer = Adam(lr = learning_rate) srcnn.compile(optimizer=optimizer, loss='mean_squared_error') history = srcnn.fit(X_train, Y_train, batch_size=batch_size, epochs=epochs, verbose=2,validation_split = 0.1 ) draw_loss_plot(history = history) #Saving trained model_weights in the current workspace . #make a folder named srcnn in your current workspace you are working in.In that folder your weights will be stored. json_string = srcnn.to_json() open(os.path.join('./srcnn/','srcnn_model.json'),'w').write(json_string) srcnn.save_weights(os.path.join('./srcnn/','srcnn_weight.hdf5')) else: dir_list = os.listdir(dirname_test) for img in dir_list: image = cv2.cvtColor(cv2.imread(os.path.join(dirname_test,img)),cv2.COLOR_BGR2GRAY) image = image[0:image.shape[0]-np.remainder(image.shape[0],scale),0:image.shape[1]-np.remainder(image.shape[1],scale)]