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)]
