def train_batch(aug, trainX, trainY, load_weight=True):
print("[INFO] compiling model...")
model = SRCNN()
if load_weight and os.path.exists('my_model_weights.h5'):
print("[INFO] import model weights...")
model.load_weights('my_model_weights.h5')
try:
opt = Adam(lr=INIT_LR, decay=INIT_LR / EPOCHS)
# opt = Adam(lr=INIT_LR)
model.compile(loss="mean_squared_error",
optimizer=opt,
metrics=[psnr1])
current_time = datetime.now().strftime("%Y%m%d-%H:%M")
filepath = 'train_model/model-ep{epoch:03d}-loss{loss:.3f}-psnr{psnr1:.3f}-' + current_time + '.h5'
checkpoint = ModelCheckpoint(filepath,
monitor='psnr1',
verbose=1,
save_best_only=True,
mode='max',
period=2)
print("[INFO] training network...")
H = model.fit_generator(aug.flow(trainX, trainY, batch_size=BS),
steps_per_epoch=len(trainX) // BS,
epochs=EPOCHS,
callbacks=[checkpoint])
print("[INFO] training finish")
except KeyboardInterrupt:
print("[INFO] KeyboardInterrupted...")
except Exception as e:
print(e)
else:
plt.figure()
N = EPOCHS
# plt.plot(np.arange(0, N), H.history["loss"], label="train_loss")
# plt.plot(np.arange(0, N), H.history["acc"], label="train_acc")
plt.plot(np.arange(0, N), H.history["psnr1"], label="psnr")
plt.xlabel("Epoch #")
plt.ylabel("psnr")
plt.legend(loc="upper left")
plt.savefig("plot.png")
finally:
print("[INFO] serializing network...")
model.save("./model.h5")
model.save_weights('my_model_weights.h5')
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)]
Y_test = image.copy()
X_test = cv2.resize(image, None, fx=1/scale, fy=1/scale)
X_test = cv2.resize(X_test, None, fx= scale/1, fy= scale/1, interpolation=cv2.INTER_CUBIC)
X_test = X_test.reshape(1,X_test.shape[0],X_test.shape[1],1)
weight_filename = 'srcnn_weight.hdf5'
