def train(self):
# Validation and Train dataset generation
train_data = DataLoader(data_name=self.args.data4train, arg=self.args)
n_train =train_data.indices.size #data_cache["n_files"]
val_data = DataLoader(data_name=self.args.data4train,
arg=self.args, is_val=True)
val_idcs = np.arange(val_data.indices.size)
# Summary and checkpoint manager
model_dir =self.args.model_name+'2'+self.args.data4train
summary_dir = os.path.join('logs',model_dir)
train_log_dir=os.path.join(summary_dir,'train')
val_log_dir =os.path.join(summary_dir,'test')
checkpoint_dir = os.path.join(self.args.checkpoint_dir,model_dir)
epoch_ckpt_dir = checkpoint_dir + 'epochs'
os.makedirs(epoch_ckpt_dir, exist_ok=True)
os.makedirs(train_log_dir,exist_ok=True)
os.makedirs(val_log_dir,exist_ok=True)
os.makedirs(checkpoint_dir, exist_ok=True)
train_writer = tf.summary.create_file_writer(train_log_dir)
val_writer = tf.summary.create_file_writer(val_log_dir)
my_model = DexiNed(rgb_mean=self.args.rgbn_mean)#rgb_mean=self.args.rgbn_mean
# accuracy = metrics.SparseCategoricalAccuracy()
accuracy = metrics.BinaryAccuracy()
accuracy_val = metrics.BinaryAccuracy()
loss_bc = losses.BinaryCrossentropy()
optimizer = optimizers.Adam(
learning_rate=self.args.lr, beta_1=self.args.beta1)
iter = 0
imgs_res_folder = os.path.join(self.args.output_dir, "current_training")
os.makedirs(imgs_res_folder, exist_ok=True)
global_loss = 1000.
t_loss = []
ckpt_save_mode = "h5"
for epoch in range(self.args.max_epochs):
# training
t_loss = []
for step, (x, y) in enumerate(train_data):
with tf.GradientTape() as tape:
pred = my_model(x, training=True)
preds, loss = pre_process_binary_cross_entropy(
loss_bc, pred, y, self.args, use_tf_loss=False)
accuracy.update_state(y_true=y, y_pred=preds[-1])
gradients = tape.gradient(loss, my_model.trainable_variables)
optimizer.apply_gradients(zip(gradients, my_model.trainable_variables))
# logging the current accuracy value so far.
t_loss.append(loss.numpy())
if step % 10 == 0:
print("Epoch:", epoch, "Step:", step, "Loss: %.4f" % loss.numpy(),
"Accuracy: %.4f" % accuracy.result(), time.ctime())
if step % 10 == 0:
# visualize preds
img_test = 'Epoch: {0} Sample {1}/{2} Loss: {3}' \
.format(epoch, step, n_train // self.args.batch_size, loss.numpy())
vis_imgs = visualize_result(
x=x[2], y=y[2], p=preds, img_title=img_test)
cv.imwrite(os.path.join(imgs_res_folder, 'results.png'), vis_imgs)
if step % 20 == 0 and loss < global_loss: # 500
if epoch==0 and step==0:
tmp_loss = np.array(t_loss)
with train_writer.as_default():
tf.summary.scalar('loss', tmp_loss.mean(), step=epoch)
tf.summary.scalar('accuracy', accuracy.result(), step=epoch)
save_ckpt_path = os.path.join(checkpoint_dir, "DexiNedL_model.h5")
Model.save_weights(my_model, save_ckpt_path, save_format='h5')
global_loss = loss
print("Model saved in: ", save_ckpt_path, "Current loss:", global_loss.numpy())
iter += 1 # global iteration
t_loss = np.array(t_loss)
# train summary
if epoch!=0:
with train_writer.as_default():
tf.summary.scalar('loss', t_loss.mean(), step=epoch)
tf.summary.scalar('accuracy', accuracy.result(), step=epoch)
Model.save_weights(my_model, os.path.join(epoch_ckpt_dir, "DexiNed{}_model.h5".format(str(epoch))),
save_format=ckpt_save_mode)
print("Epoch:", epoch, "Model saved in Loss: ", t_loss.mean())
# validation
t_val_loss = []
for i, (x_val, y_val) in enumerate(val_data):
pred_val = my_model(x_val)
v_logits, V_loss = pre_process_binary_cross_entropy(
loss_bc, pred_val, y_val, self.args, use_tf_loss=False)
accuracy_val.update_state(y_true=y_val, y_pred=v_logits[-1])
t_val_loss.append(V_loss.numpy())
if i == 7:
break
val_acc = accuracy_val.result()
t_val_loss = np.array(t_val_loss)
print("Epoch(validation):", epoch, "Val loss: ", t_val_loss.mean(),
"Accuracy: ", val_acc.numpy())
# validation summary
with val_writer.as_default():
tf.summary.scalar('loss', t_val_loss.mean(), step=epoch)
tf.summary.scalar('accuracy', val_acc.numpy(), step=epoch)
# Reset metrics every epoch
accuracy.reset_states()
accuracy_val.reset_states()
my_model.summary()
for i in range(len(Y_hat_name)):
tmp_name = Y_hat_name[i]
tmp_name = tmp_name.replace("RGBNC", "RGBNP")
h5_writer(savepath=os.path.join(Yhat_dir, tmp_name),
data=np.squeeze(Y_hat[i, :, :, :]))
else:
tmp_name = Y_hat_name
tmp_name = tmp_name.replace("RGBN", "RGBNP")
h5_writer(savepath=os.path.join(Yhat_dir, tmp_name),
data=np.squeeze(Y_hat))
if arg.model_name == "CDNet" or arg.model_name == "ENDENet":
if not arg.is_training:
data4testing = DataLoader(data_name=arg.dataset_name, arg=arg)
model_dir = arg.model_name.lower() + '2' + arg.dataset_name
res_dir = os.path.join('results', model_dir)
os.makedirs(res_dir, exist_ok=True)
ckpnt_dir = os.path.join(arg.ckpt_dir, model_dir)
ckpnt_path = os.path.join(ckpnt_dir, 'saved_weights.h5')
my_model = CDENT()
loss_mse = tfk.losses.mean_squared_error
optimizer = tfk.optimizers.Adam(learning_rate=arg.lr)
my_model.compile(optimizer=optimizer, loss=loss_mse, metrics='mse')
input_shape = data4testing.input_shape
my_model.build(input_shape=input_shape)
my_model.load_weights(filepath=ckpnt_path)
def test(self):
# Test dataset generation
test_data = DataLoader(data_name=self.args.data4test, arg=self.args)
n_test = test_data.indices.size # data_cache["n_files"]
optimizer = tf.keras.optimizers.Adam(
learning_rate=self.args.lr, beta_1=self.args.beta1)
my_model = DexiNed(rgb_mean=self.args.rgbn_mean)
input_shape = test_data.input_shape
my_model.build(input_shape=input_shape) # rgb_mean=self.args.rgbn_mean
checkpoit_dir = os.path.join(self.args.checkpoint_dir,
self.args.model_name + "2" + self.args.data4train)
my_model.load_weights(os.path.join(checkpoit_dir, "DexiNed23_model.h5"))
result_dir = os.path.join(
self.args.output_dir,
self.args.model_name + '-' + self.args.data4train + "2" + self.args.data4test)
os.makedirs(result_dir, exist_ok=True)
if self.args.scale is not None:
scl = self.args.scale
save_dir = ['fuse_'+str(scl), 'avrg_'+str(scl), 'h5_'+str(scl)]
else:
save_dir = ['fuse', 'avrg', 'h5']
save_dirs = []
for tmp_dir in save_dir:
os.makedirs(os.path.join(result_dir, tmp_dir), exist_ok=True)
save_dirs.append(os.path.join(result_dir, tmp_dir))
total_time = []
data_names = test_data.imgs_name
data_shape = test_data.imgs_shape
k = 0
for step, (x, y) in enumerate(test_data):
start_time = time.time()
preds = my_model(x, training=False)
tmp_time = time.time() - start_time
total_time.append(tmp_time)
preds = [tf.sigmoid(i).numpy() for i in preds]
all_preds = np.array(preds)
for i in range(all_preds.shape[1]):
tmp_name = data_names[k]
tmp_name, _ = os.path.splitext(tmp_name)
tmp_shape = data_shape[k]
tmp_preds = all_preds[:, i, ...]
tmp_av = np.expand_dims(tmp_preds.mean(axis=0), axis=0)
tmp_preds = np.concatenate((tmp_preds, tmp_av), axis=0)
res_preds = []
for j in range(tmp_preds.shape[0]):
tmp_pred = tmp_preds[j, ...]
tmp_pred[tmp_pred < 0.0] = 0.0
tmp_pred = cv.bitwise_not(np.uint8(image_normalization(tmp_pred)))
h, w = tmp_pred.shape[:2]
if h != tmp_shape[0] or w != tmp_shape[1]:
tmp_pred = cv.resize(tmp_pred, (tmp_shape[1], tmp_shape[0]))
res_preds.append(tmp_pred)
n_save =len(tmp_preds)-2
for idx in range(len(save_dirs) - 1):
s_dir = save_dirs[idx]
tmp = res_preds[n_save + idx]
cv.imwrite(join(s_dir, tmp_name + '.png'), tmp)
h5_writer(path=join(save_dirs[-1], tmp_name + '.h5'),
vars=np.squeeze(res_preds))
print("saved:", join(save_dirs[-1], tmp_name + '.h5'), tmp_preds.shape)
k += 1
# tmp_name = data_names[step][:-3]+"png"
# tmp_shape = data_shape[step]
# tmp_path = os.path.join(result_dir,tmp_name)
# tensor2image(preds[-1].numpy(), img_path =tmp_path,img_shape=tmp_shape)
total_time = np.array(total_time)
print('-------------------------------------------------')
print("End testing in: ", self.args.data4test)
print("Batch size: ", self.args.test_bs)
print("Time average per image: ", total_time.mean(), "secs")
print("Total time: ", total_time.sum(), "secs")
print('-------------------------------------------------')
def train():
if arg.model_name.lower() == "cdnet" or arg.model_name.lower(
) == "endenet":
# ***************data preparation *********************
if arg.model_state.lower() == 'train':
# dataset preparation for training
running_mode = 'train'
data4training = DataLoader(data_name=arg.dataset_name, arg=arg)
# define model and callbacks
model_dir = arg.model_name.lower() + '2' + arg.dataset_name
ckpnt_dir = os.path.join(arg.ckpt_dir, model_dir)
ckpnt_path = os.path.join(ckpnt_dir, 'saved_weights.h5')
os.makedirs(ckpnt_dir, exist_ok=True)
log_dir = os.path.join('logs', model_dir)
res_dir = os.path.join('results', model_dir)
os.makedirs(res_dir, exist_ok=True)
my_callbacks = [
tfk.callbacks.ModelCheckpoint(
ckpnt_path,
monitor='train_loss', # os.path.join(ckpnt,saved_weights.h5)
save_weights_only=True,
mode='auto',
save_freq='epoch'),
tfk.callbacks.TensorBoard(log_dir,
histogram_freq=0,
write_graph=True,
profile_batch=2,
write_images=True)
]
my_model = CDENT()
loss_mse = tfk.losses.mean_squared_error
accuracy = tfk.metrics.MeanAbsolutePercentageError()
optimizer = tfk.optimizers.Adam(learning_rate=arg.lr, beta_1=0.5)
# compile model
# my_model.compile(optimizer=optimizer, loss=loss_mse)
# my_model.fit(data4training, epochs=arg.num_epochs,callbacks=my_callbacks)
for epoch in range(arg.num_epochs):
total_loss = tf.Variable(0.)
for step, (x, y) in enumerate(data4training):
with tf.GradientTape() as tape:
p = my_model(x)
loss = loss_mse(y_true=y, y_pred=p)
loss = tf.math.reduce_sum(loss)
total_loss = tf.add(total_loss, loss)
accuracy.update_state(y_true=y, y_pred=p)
gradients = tape.gradient(loss,
my_model.trainable_variables)
optimizer.apply_gradients(
zip(gradients, my_model.trainable_variables))
if step % 10 == 0:
print("Epoch:", epoch, "Step:", step,
"Loss: %.4f" % loss.numpy(),
"Accuracy: %.4f" % accuracy.result(),
time.ctime())
tfk.Model.save_weights(my_model, ckpnt_path, save_format='h5')
print('Model saved in:', ckpnt_path)
# visualize result
mean_loss = total_loss / 50
tmp_x = image_normalization(np.squeeze(x[2, :, :, :3]))
tmp_y = image_normalization(np.squeeze(y[2, ...]))
tmp_p = p[2, ...]
tmp_p = image_normalization(tmp_p.numpy())
vis_imgs = np.uint8(
np.concatenate((tmp_x, tmp_y, tmp_p), axis=1))
img_test = 'Epoch: {0} Loss: {1}'.format(
epoch, mean_loss.numpy())
BLACK = (0, 0, 255)
font = cv.FONT_HERSHEY_SIMPLEX
font_size = 1.1
font_color = BLACK
font_thickness = 2
x, y = 30, 30
vis_imgs = cv.putText(vis_imgs, img_test, (x, y), font,
font_size, font_color, font_thickness,
cv.LINE_AA)
cv.imwrite(os.path.join(res_dir, ' results.png'), vis_imgs)
print("<<< End epoch loss: ", mean_loss.numpy(), " >>>")
my_model.summary()
print('Training finished on: ', arg.dataset_name)